orbit/internal/storage/mysql/keyed_concurrency_integration_test.go

package mysql

import (
	"context"
	"encoding/json"
	"os"
	"sync"
	"testing"
	"time"

	core "git.yuntuo.hk/k2safe/orbit/model"
	pulse "git.yuntuo.hk/k2safe/orbit/pulse/model"
	"git.yuntuo.hk/k2safe/orbit/pulse/service"
	"git.yuntuo.hk/k2safe/orbit/pulse/worker"
)

func keyedWorker(t *testing.T, f *executionFixture, queue core.ID, node *worker.Service) *executionFixture {
	t.Helper()
	g := *f
	g.resources.QueueID = queue
	g.a = node
	g.hello.InstanceID = executionID(t)
	g.hello.QueueID = queue
	g.hello.ConcurrencyKeys = true
	var err error
	g.binding, _, err = node.Connect(g.ctx, g.token, g.hello)
	if err != nil {
		t.Fatal(err)
	}
	return &g
}
func keyedSubmit(t *testing.T, f *executionFixture, key string, max uint32) service.TaskView {
	t.Helper()
	out, err := f.tasks.Submit(f.ctx, f.token, string(executionID(t)), service.SubmitInput{QueueID: f.resources.QueueID, Handler: "object.prepare", HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{}`), ConcurrencyKey: key, Retry: pulse.RetryPolicy{MaxAttempts: max, InitialDelayMS: 1000, MaxDelayMS: 30000, AttemptTimeoutMS: 30000}})
	if err != nil {
		t.Fatal(err)
	}
	return out.Task
}
func keyedUsed(t *testing.T, f *executionFixture, key string) uint32 {
	t.Helper()
	var used uint32
	if err := f.db.QueryRowContext(f.ctx, `SELECT used_permits FROM pulse_concurrency_limits WHERE namespace_id=? AND scope_key='key' AND key_digest=?`, f.resources.NamespaceID, concurrencyDigest(key)).Scan(&used); err != nil {
		t.Fatal(err)
	}
	return used
}
func noKeyClaim(t *testing.T, f *executionFixture) {
	t.Helper()
	if a, err := f.a.Claim(f.ctx, f.token, f.binding); err != nil || a != nil {
		t.Fatal("expected no eligible claim", a, err)
	}
}
func TestIntegrationKeyedConcurrency(t *testing.T) {
	dsn := os.Getenv("ORBIT_TEST_MYSQL_DSN")
	if dsn == "" {
		t.Skip("requires disposable MySQL")
	}
	ctx, stop := context.WithTimeout(context.Background(), 70*time.Second)
	defer stop()
	db, err := Open(ctx, dsn, 32, 8)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	if err := Migrate(ctx, db); err != nil {
		t.Fatal(err)
	}
	fixture := func(t *testing.T) *executionFixture {
		f := newExecutionFixture(t, ctx, db)
		return keyedWorker(t, f, f.resources.QueueID, f.a)
	}
	t.Run("legacy_workers_and_tracking_keys_keep_their_contract", func(t *testing.T) {
		old := newExecutionFixture(t, ctx, db)
		keyed := keyedSubmit(t, old, "resource/one", 1)
		noKeyClaim(t, old)
		if old.get(keyed.ID).AttemptCount != 0 {
			t.Fatal("legacy worker consumed keyed attempt")
		}
		for i := 0; i < 2; i++ {
			if _, err := old.tasks.Submit(ctx, old.token, string(executionID(t)), service.SubmitInput{QueueID: old.resources.QueueID, Handler: "object.prepare", HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{}`), BusinessKey: "same-tracking-key"}); err != nil {
				t.Fatal(err)
			}
		}
		a, _ := old.claim()
		b, _ := old.claim()
		if a.ConcurrencyToken != 0 || b.ConcurrencyToken != 0 {
			t.Fatal("tracking key changed execution semantics")
		}
		f := keyedWorker(t, old, old.resources.QueueID, old.b)
		a, _ = f.claim()
		if a.Task.ID != keyed.ID || a.ConcurrencyToken == 0 {
			t.Fatal("opted-in worker did not receive key token")
		}
		hello := f.hello
		hello.ConcurrencyKeys = false
		_, _, err := f.a.Connect(ctx, f.token, hello)
		expectWorkerCode(t, err, "STATE_CONFLICT")
		// 旧 flag=false 摘要必须继续允许旧进程重连。
		if _, _, err := old.a.Connect(ctx, old.token, old.hello); err != nil {
			t.Fatal("legacy registration digest changed", err)
		}
	})
	t.Run("cross_queue_race_single_permit_and_monotonic_tokens", func(t *testing.T) {
		f := fixture(t)
		q, err := f.tasks.CreateQueue(ctx, f.token, "second-queue", service.QueueInput{Name: "second", MaxConcurrency: 4, RatePerSecond: 100})
		if err != nil {
			t.Fatal(err)
		}
		g := keyedWorker(t, f, q.Queue.ID, f.b)
		keyedSubmit(t, f, "object/42", 1)
		keyedSubmit(t, g, "object/42", 1)
		type claim struct {
			f   *executionFixture
			a   *worker.Assignment
			err error
		}
		ch := make(chan claim, 2)
		var wg sync.WaitGroup
		for _, target := range []*executionFixture{f, g} {
			wg.Add(1)
			go func() {
				defer wg.Done()
				a, err := target.a.Claim(ctx, target.token, target.binding)
				ch <- claim{target, a, err}
			}()
		}
		wg.Wait()
		close(ch)
		var won claim
		for v := range ch {
			if v.err != nil {
				t.Fatal(v.err)
			}
			if v.a != nil {
				if won.a != nil {
					t.Fatal("same key dispatched twice")
				}
				won = v
			}
		}
		if won.a == nil || won.a.ConcurrencyToken == 0 || keyedUsed(t, f, "object/42") != 1 {
			t.Fatal("missing key permit")
		}
		assertApplicationCounts(t, f, 2, 1, 1)
		if f.count("pulse_permits", "attempt_id=? AND released_at IS NULL", won.a.AttemptID) != 2 {
			t.Fatal("expected queue and key permits")
		}
		task := won.a.Task
		if _, err := won.f.tasks.CancelTask(ctx, won.f.token, task.ID, "cancel", service.CancelInput{ExpectedRevision: task.Revision, Reason: "release unstarted"}); err != nil {
			t.Fatal(err)
		}
		other := f
		if won.f == f {
			other = g
		}
		next, ref := other.claim()
		if next.ConcurrencyToken <= won.a.ConcurrencyToken {
			t.Fatal("key sequence reset across tasks")
		}
		message := executionID(t)
		if _, err := other.a.Start(ctx, other.token, other.binding, message, ref); err != nil {
			t.Fatal(err)
		}
		if _, err := other.a.Start(ctx, other.token, other.binding, message, ref); err != nil {
			t.Fatal(err)
		}
		if _, err := other.a.Renew(ctx, other.token, other.binding, executionID(t), ref); err != nil {
			t.Fatal(err)
		}
		history, err := other.tasks.GetTaskAttempt(ctx, other.token, ref.TaskID, ref.AttemptID)
		if err != nil || history.ConcurrencyToken != next.ConcurrencyToken {
			t.Fatal("token not preserved in history", err)
		}
		retryFinish(t, other, ref, "SUCCEEDED", false)
		assertApplicationCounts(t, f, 0, 0, 0)
		if keyedUsed(t, f, "object/42") != 0 {
			t.Fatal("key permit leaked")
		}
	})
	t.Run("busy_keys_do_not_hide_later_keys_or_other_namespaces", func(t *testing.T) {
		f := fixture(t)
		keyedSubmit(t, f, "busy", 1)
		f.claim()
		for i := 0; i < 20; i++ {
			keyedSubmit(t, f, "busy", 1)
		}
		free := keyedSubmit(t, f, "free", 1)
		a, _ := f.claim()
		if a.Task.ID != free.ID {
			t.Fatal("blocked prefix hid free key")
		}
		if f.count("pulse_tasks", "concurrency_key='busy' AND state='WAITING' AND attempt_count<>0") != 0 {
			t.Fatal("blocked keys consumed attempts")
		}
		g := newExecutionFixture(t, ctx, db, f.resources.ApplicationID)
		g = keyedWorker(t, g, g.resources.QueueID, g.a)
		keyedSubmit(t, g, "busy", 1)
		a, _ = g.claim()
		if a.ConcurrencyToken == 0 {
			t.Fatal("namespace isolation lost")
		}
	})
	t.Run("unknown_late_result_and_retry_release_exactly_once", func(t *testing.T) {
		f := fixture(t)
		task := keyedSubmit(t, f, "resource/retry", 2)
		first, ref := f.claim()
		if _, err := f.a.Start(ctx, f.token, f.binding, executionID(t), ref); err != nil {
			t.Fatal(err)
		}
		retryFinish(t, f, ref, "FAILED", true)
		if keyedUsed(t, f, "resource/retry") != 0 {
			t.Fatal("retry holds key while waiting")
		}
		makeRetryDue(f, task.ID)
		second, ref2 := f.claim()
		if second.ConcurrencyToken <= first.ConcurrencyToken {
			t.Fatal("retry reused token")
		}
		if _, err := f.a.Start(ctx, f.token, f.binding, executionID(t), ref2); err != nil {
			t.Fatal(err)
		}
		retryFinish(t, f, ref2, "UNKNOWN", false)
		assertApplicationCounts(t, f, 1, 0, 0)
		if f.get(task.ID).State != pulse.TaskNeedsReview {
			t.Fatal("unknown result lost")
		}
		keyedSubmit(t, f, "resource/retry", 1)
		third, ref3 := f.claim()
		if third.ConcurrencyToken <= second.ConcurrencyToken {
			t.Fatal("key epoch did not advance")
		}
		late, err := f.a.Finish(ctx, f.token, f.binding, executionID(t), worker.FinishInput{Ref: ref2, Outcome: "SUCCEEDED", Result: json.RawMessage(`{}`)})
		if err != nil || late.Receipt.Disposition != "EVIDENCE_ONLY" {
			t.Fatal("late result applied", err)
		}
		if keyedUsed(t, f, "resource/retry") != 1 {
			t.Fatal("late result released another attempt's key")
		}
		f.expire(ref3)
		if _, err := f.storeB.Recover(ctx, 100); err != nil {
			t.Fatal(err)
		}
		if keyedUsed(t, f, "resource/retry") != 0 {
			t.Fatal("expiry leaked key")
		}
		assertApplicationCounts(t, f, 1, 0, 0)
	})
	t.Run("claim_rollback_and_exact_utf8_key_identity", func(t *testing.T) {
		f := fixture(t)
		keyedSubmit(t, f, "asset/A", 1)
		trigger := "test_key_event_" + string(executionID(t))
		f.exec("CREATE TRIGGER " + trigger + " BEFORE INSERT ON orbit_events FOR EACH ROW SIGNAL SQLSTATE '45000' SET MESSAGE_TEXT='injected claim failure'")
		out, err := f.a.Claim(ctx, f.token, f.binding)
		f.exec("DROP TRIGGER " + trigger)
		if err == nil || out != nil {
			t.Fatal("failed claim exposed permit")
		}
		assertApplicationCounts(t, f, 1, 0, 0)
		if f.count("pulse_attempts", "1=1") != 0 || f.count("pulse_permits", "1=1") != 0 || f.count("pulse_concurrency_limits", "scope_key='key'") != 0 {
			t.Fatal("claim did not fully roll back")
		}
		f.claim()
		for _, key := range []string{"asset/a", "asset/é", "asset/e\u0301"} {
			keyedSubmit(t, f, key, 1)
			f.claim()
		}
		if f.count("pulse_concurrency_limits", "scope_key='key' AND used_permits=1") != 4 {
			t.Fatal("resource keys normalized unexpectedly")
		}
	})
	t.Run("sequence_exhaustion_is_not_allowed_to_wrap", func(t *testing.T) {
		f := fixture(t)
		keyedSubmit(t, f, "exhausted", 1)
		ref := retryStart(t, f)
		retryFinish(t, f, ref, "SUCCEEDED", false)
		f.exec(`UPDATE pulse_concurrency_limits SET revision=18446744073709551613 WHERE namespace_id=? AND scope_key='key'`, f.resources.NamespaceID)
		keyedSubmit(t, f, "exhausted", 1)
		a, ref := f.claim()
		if a.ConcurrencyToken != ^uint64(0)-1 {
			t.Fatal("high token truncated")
		}
		if _, err := f.a.Start(ctx, f.token, f.binding, executionID(t), ref); err != nil {
			t.Fatal(err)
		}
		retryFinish(t, f, ref, "SUCCEEDED", false)
		waiting := keyedSubmit(t, f, "exhausted", 1)
		_, err := f.a.Claim(ctx, f.token, f.binding)
		expectWorkerCode(t, err, "STATE_CONFLICT")
		if f.get(waiting.ID).AttemptCount != 0 {
			t.Fatal("exhausted sequence consumed task attempt")
		}
	})
	t.Run("nonempty_keys_do_not_use_text_collation_for_null", func(t *testing.T) {
		f := fixture(t)
		for _, key := range []string{"\u0301", "\u200b"} {
			task := keyedSubmit(t, f, key, 1)
			if task.ConcurrencyKey != key {
				t.Fatal("nonempty key collapsed by database collation")
			}
			a, _ := f.claim()
			if a.Task.ID != task.ID || a.Task.ConcurrencyKey != key || a.ConcurrencyToken == 0 {
				t.Fatal("exact key did not reach assignment")
			}
		}
	})
	t.Run("schedule_template_and_key_validation", func(t *testing.T) {
		f := fixture(t)
		sf := newScheduleFixture(t, f)
		def := sf.definition()
		def.Task.ConcurrencyKey = "schedule/resource"
		def.Task.Priority = 7
		v := sf.create(def)
		sf.due(v, 2*time.Second)
		sf.tick()
		a, ref := f.claim()
		if a.Task.ConcurrencyKey != def.Task.ConcurrencyKey || a.ConcurrencyToken == 0 || a.Task.Priority != 7 {
			t.Fatal("schedule omitted frozen key")
		}
		noKeyClaim(t, f)
		if _, err := f.a.Start(ctx, f.token, f.binding, executionID(t), ref); err != nil {
			t.Fatal(err)
		}
		retryFinish(t, f, ref, "SUCCEEDED", false)
		b, _ := f.claim()
		if b.ConcurrencyToken <= a.ConcurrencyToken {
			t.Fatal("schedule reused key token")
		}
		for _, key := range []string{" leading", "trailing ", "line\nfeed", "\u0000"} {
			_, err := f.tasks.Submit(ctx, f.token, string(executionID(t)), service.SubmitInput{QueueID: f.resources.QueueID, Handler: "object.prepare", HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{}`), ConcurrencyKey: key})
			expectWorkerCode(t, err, "INVALID_ARGUMENT")
		}
	})
}