feat: add idempotent batch task submission and recovery examples

internal/storage/mysql/task_batch_integration_test.go

@@ -0,0 +1,230 @@
+package mysql
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"os"
+	"sync"
+	"testing"
+	"time"
+
+	core "git.yuntuo.hk/k2safe/orbit/model"
+	"git.yuntuo.hk/k2safe/orbit/platform/auth"
+	"git.yuntuo.hk/k2safe/orbit/pulse/service"
+)
+
+func taskBatch(f *executionFixture, n int) service.BatchSubmitInput {
+	in := service.BatchSubmitInput{}
+	for i := 0; i < n; i++ {
+		in.Items = append(in.Items, service.BatchSubmitItem{IdempotencyKey: fmt.Sprintf("batch-%d", i), Task: service.SubmitInput{QueueID: f.resources.QueueID, Handler: "object.prepare", HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{"n":9007199254740993}`)}})
+	}
+	return in
+}
+func requireBatchStatuses(t *testing.T, out service.BatchSubmitResult, err error, wants ...service.BatchItemStatus) {
+	t.Helper()
+	if err != nil || len(out.Items) != len(wants) {
+		t.Fatal(out, err)
+	}
+	for i, want := range wants {
+		if out.Items[i].Index != i || out.Items[i].Status != want {
+			t.Fatalf("item %d=%+v want=%s", i, out.Items[i], want)
+		}
+	}
+}
+
+type batchAfterCommitStore struct {
+	service.Store
+	after func()
+}
+
+func (s *batchAfterCommitStore) Submit(ctx context.Context, token auth.Token, key string, in service.SubmitInput) (service.SubmitResult, error) {
+	out, err := s.Store.Submit(ctx, token, key, in)
+	if err == nil && s.after != nil {
+		fn := s.after
+		s.after = nil
+		fn()
+	}
+	return out, err
+}
+
+func TestIntegrationTaskBatch(t *testing.T) {
+	dsn := os.Getenv("ORBIT_TEST_MYSQL_DSN")
+	if dsn == "" {
+		t.Skip("requires disposable MySQL")
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
+	defer cancel()
+	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 { return newExecutionFixture(t, ctx, db) }
+	t.Run("partial_rejections_replay_and_quota", func(t *testing.T) {
+		f := fixture(t)
+		configureResources(t, f, 10, 4, 2, true)
+		in := taskBatch(f, 5)
+		in.Items[1].Task.Handler = ""
+		in.Items[2].Task.QueueID = executionID(t)
+		out, err := f.tasks.SubmitBatch(ctx, f.token, in)
+		requireBatchStatuses(t, out, err, service.BatchAccepted, service.BatchRejected, service.BatchRejected, service.BatchAccepted, service.BatchRejected)
+		if out.Items[1].Error.Code != "INVALID_ARGUMENT" || out.Items[2].Error.Code != "NOT_FOUND" || out.Items[4].Error.Code != "QUEUE_FULL" || out.Items[4].Error.RetryAfterMS != 1000 {
+			t.Fatal(out)
+		}
+		assertQueueCount(t, f, 2)
+		assertApplicationCounts(t, f, 2, 0, 0)
+		for table, filter := range map[string]string{"orbit_idempotency_keys": "operation='task.submit'", "orbit_audit_events": "action='task.submit'", "orbit_events": "type='task.accepted'"} {
+			if f.count(table, filter) != 2 {
+				t.Fatal("partial batch leaked writes", table)
+			}
+		}
+		if string(f.get(out.Items[0].Receipt.TaskID).Payload) != `{"n":9007199254740993}` {
+			t.Fatal("payload precision lost")
+		}
+		one, err := f.tasks.Submit(ctx, f.token, in.Items[0].IdempotencyKey, in.Items[0].Task)
+		if err != nil || !one.Deduplicated || one.Task.ID != out.Items[0].Receipt.TaskID {
+			t.Fatal("single batch identity diverged", err)
+		}
+		first := out.Items[0].Receipt
+		if _, err := f.tasks.CancelTask(ctx, f.token, first.TaskID, "cancel", service.CancelInput{ExpectedRevision: first.Revision, Reason: "free batch slot"}); err != nil {
+			t.Fatal(err)
+		}
+		in.Items = []service.BatchSubmitItem{in.Items[0], in.Items[4]}
+		again, err := f.tasks.SubmitBatch(ctx, f.token, in)
+		requireBatchStatuses(t, again, err, service.BatchAccepted, service.BatchAccepted)
+		if !again.Items[0].Receipt.Deduplicated || again.Items[0].Receipt.Revision != 1 || again.Items[1].Receipt.Deduplicated {
+			t.Fatal("receipt/rejected identity not reusable")
+		}
+		in.Items[0].Task.Payload = json.RawMessage(`{"changed":true}`)
+		conflict, err := f.tasks.SubmitBatch(ctx, f.token, in)
+		requireBatchStatuses(t, conflict, err, service.BatchRejected, service.BatchAccepted)
+		if conflict.Items[0].Error.Code != "IDEMPOTENCY_CONFLICT" {
+			t.Fatal(conflict)
+		}
+		assertQueueCount(t, f, 2)
+		assertApplicationCounts(t, f, 2, 0, 0)
+	})
+	t.Run("cross_namespace_and_queue_scope", func(t *testing.T) {
+		f, g := fixture(t), fixture(t)
+		in := taskBatch(f, 2)
+		in.Items[0].Task.QueueID = g.resources.QueueID
+		out, err := f.tasks.SubmitBatch(ctx, f.token, in)
+		requireBatchStatuses(t, out, err, service.BatchRejected, service.BatchAccepted)
+		if out.Items[0].Error.Code != "NOT_FOUND" {
+			t.Fatal(out)
+		}
+		f.exec(`UPDATE orbit_credentials SET permissions=? WHERE namespace_id=? AND id=?`, `{"actions":["tasks:submit"],"queues":["`+string(g.resources.QueueID)+`"]}`, f.resources.NamespaceID, f.resources.CredentialID)
+		out, err = f.tasks.SubmitBatch(ctx, f.token, service.BatchSubmitInput{Items: in.Items[1:]})
+		requireBatchStatuses(t, out, err, service.BatchRejected)
+		if out.Items[0].Error.Code != "PERMISSION_DENIED" || out.Items[0].Receipt != nil {
+			t.Fatal("cached receipt bypassed current queue scope")
+		}
+		if g.count("pulse_tasks", "1=1") != 0 {
+			t.Fatal("cross namespace write")
+		}
+	})
+	t.Run("revoke_between_items_is_rechecked", func(t *testing.T) {
+		f := fixture(t)
+		s := &batchAfterCommitStore{Store: NewTaskStore(db), after: func() {
+			f.exec(`UPDATE orbit_credentials SET revoked_at=UTC_TIMESTAMP(6) WHERE namespace_id=? AND id=?`, f.resources.NamespaceID, f.resources.CredentialID)
+		}}
+		out, err := service.New(s).SubmitBatch(ctx, f.token, taskBatch(f, 3))
+		requireBatchStatuses(t, out, err, service.BatchAccepted, service.BatchRejected, service.BatchNotAttempted)
+		if out.Items[1].Error.Code != "UNAUTHENTICATED" || f.count("pulse_tasks", "1=1") != 1 {
+			t.Fatal(out)
+		}
+	})
+	t.Run("cancel_after_first_commit_leaves_recoverable_prefix", func(t *testing.T) {
+		f := fixture(t)
+		run, stop := context.WithCancel(ctx)
+		defer stop()
+		s := &batchAfterCommitStore{Store: NewTaskStore(db), after: stop}
+		in := taskBatch(f, 3)
+		out, err := service.New(s).SubmitBatch(run, f.token, in)
+		requireBatchStatuses(t, out, err, service.BatchAccepted, service.BatchNotAttempted, service.BatchNotAttempted)
+		if f.count("pulse_tasks", "1=1") != 1 {
+			t.Fatal("cancelled suffix wrote tasks")
+		}
+		again, err := f.tasks.SubmitBatch(ctx, f.token, in)
+		requireBatchStatuses(t, again, err, service.BatchAccepted, service.BatchAccepted, service.BatchAccepted)
+		if !again.Items[0].Receipt.Deduplicated || again.Items[0].Receipt.TaskID != out.Items[0].Receipt.TaskID {
+			t.Fatal("lost committed prefix")
+		}
+		assertApplicationCounts(t, f, 3, 0, 0)
+	})
+	t.Run("audit_failure_rolls_back_item_and_preserves_previous_commit", func(t *testing.T) {
+		f := fixture(t)
+		in := taskBatch(f, 3)
+		if _, err := f.tasks.Submit(ctx, f.token, in.Items[0].IdempotencyKey, in.Items[0].Task); err != nil {
+			t.Fatal(err)
+		}
+		trigger := "batch_audit_" + string(executionID(t))
+		f.exec("CREATE TRIGGER " + trigger + " BEFORE INSERT ON orbit_audit_events FOR EACH ROW SIGNAL SQLSTATE '45000' SET MESSAGE_TEXT='private injected audit fault'")
+		out, err := f.tasks.SubmitBatch(ctx, f.token, in)
+		f.exec("DROP TRIGGER " + trigger)
+		requireBatchStatuses(t, out, err, service.BatchAccepted, service.BatchUnknown, service.BatchNotAttempted)
+		if out.Items[1].Error.Code != "UNAVAILABLE" {
+			t.Fatal(out)
+		}
+		assertQueueCount(t, f, 1)
+		assertApplicationCounts(t, f, 1, 0, 0)
+		for table, filter := range map[string]string{"orbit_idempotency_keys": "operation='task.submit'", "orbit_audit_events": "action='task.submit'", "orbit_events": "type='task.accepted'"} {
+			if f.count(table, filter) != 1 {
+				t.Fatal("atomic item rollback failed", table)
+			}
+		}
+		again, err := f.tasks.SubmitBatch(ctx, f.token, in)
+		requireBatchStatuses(t, again, err, service.BatchAccepted, service.BatchAccepted, service.BatchAccepted)
+		assertQueueCount(t, f, 3)
+	})
+	t.Run("concurrent_replay_and_shared_application_backpressure", func(t *testing.T) {
+		f := fixture(t)
+		grantApplicationResources(f)
+		configureApplication(t, f, 100, 10, 3)
+		in := taskBatch(f, 4)
+		type result struct {
+			out service.BatchSubmitResult
+			err error
+		}
+		ch := make(chan result, 8)
+		var wg sync.WaitGroup
+		for i := 0; i < 8; i++ {
+			wg.Add(1)
+			go func() { defer wg.Done(); out, err := f.tasks.SubmitBatch(ctx, f.token, in); ch <- result{out, err} }()
+		}
+		wg.Wait()
+		close(ch)
+		ids := map[int]core.ID{}
+		fresh := 0
+		for r := range ch {
+			requireBatchStatuses(t, r.out, r.err, service.BatchAccepted, service.BatchAccepted, service.BatchAccepted, service.BatchRejected)
+			for i := 0; i < 3; i++ {
+				receipt := r.out.Items[i].Receipt
+				if previous := ids[i]; previous != "" && previous != receipt.TaskID {
+					t.Fatal("duplicate identity")
+				}
+				ids[i] = receipt.TaskID
+				if !receipt.Deduplicated {
+					fresh++
+				}
+			}
+			if r.out.Items[3].Error.Code != "APPLICATION_FULL" {
+				t.Fatal(r.out)
+			}
+		}
+		if fresh != 3 || f.count("pulse_tasks", "1=1") != 3 {
+			t.Fatal("duplicate creation", fresh)
+		}
+		assertApplicationCounts(t, f, 3, 0, 0)
+		g := newExecutionFixture(t, ctx, db, f.resources.ApplicationID)
+		out, err := g.tasks.SubmitBatch(ctx, g.token, taskBatch(g, 1))
+		requireBatchStatuses(t, out, err, service.BatchRejected)
+		if out.Items[0].Error.Code != "APPLICATION_FULL" {
+			t.Fatal(out)
+		}
+	})
+}

internal/storage/mysql/task_store.go

@@ -28,6 +28,11 @@ func NewTaskStore(db *sql.DB) *TaskStore { return &TaskStore{db: db} }
 
 var _ service.Store = (*TaskStore)(nil)
 
+// AuthorizeSubmit 用于批量请求的只读预认证；每项事务仍独立校验凭据、动作与队列范围。
+func (s *TaskStore) AuthorizeSubmit(ctx context.Context, token auth.Token) error {
+	return s.transaction(ctx, token, "tasks:submit", func(*sql.Tx, auth.Permissions) error { return nil })
+}
+
 func (s *TaskStore) transaction(ctx context.Context, token auth.Token, action string, fn func(*sql.Tx, auth.Permissions) error) error {
 	return WithTx(ctx, s.db, func(tx *sql.Tx) error {
 		p, err := authenticate(ctx, tx, token)
@@ -159,6 +164,12 @@ func (s *TaskStore) Submit(ctx context.Context, token auth.Token, key string, in
 		if !p.AllowsQueue(in.QueueID) {
 			return fault.Forbidden
 		}
+		// 先按既有应用额度锁串行进入受理，再占用幂等唯一键。
+		// 满额回滚时，同键并发INSERT的重复检查锁不能与应用额度锁互相等待。
+		// 此处只取得锁，不检查/扣减额度；已有回执在满额时仍可返回。
+		if _, err := lockApplicationResources(ctx, tx, token.NamespaceID, false); err != nil {
+			return err
+		}
 		taskID, err := core.NewID()
 		if err != nil {
 			return err
@@ -174,7 +185,7 @@ func (s *TaskStore) Submit(ctx context.Context, token auth.Token, key string, in
 			out.Deduplicated = true
 			return nil
 		}
-		if _, _, err := lockQueue(ctx, tx, token.NamespaceID, in.QueueID); err != nil {
+		if _, _, err := lockQueueAfterApplication(ctx, tx, token.NamespaceID, in.QueueID); err != nil {
 			return err
 		}
 		now, err := dbNow(ctx, tx)

internal/storage/mysql/worker_execution.go

@@ -49,6 +49,11 @@ func lockQueue(ctx context.Context, tx *sql.Tx, ns, id core.ID) (paused bool, ma
 	if _, err = lockApplicationResources(ctx, tx, ns, false); err != nil {
 		return
 	}
+	return lockQueueAfterApplication(ctx, tx, ns, id)
+}
+
+// lockQueueAfterApplication 仅在同一事务已持有应用额度排他锁时调用。
+func lockQueueAfterApplication(ctx context.Context, tx *sql.Tx, ns, id core.ID) (paused bool, max uint32, err error) {
 	err = tx.QueryRowContext(ctx, "SELECT paused,max_concurrency FROM pulse_queues WHERE namespace_id=? AND id=? FOR UPDATE", ns, id).Scan(&paused, &max)
 	if errors.Is(err, sql.ErrNoRows) {
 		err = fault.NotFound