feat: add idempotent batch task submission and recovery examples

This commit is contained in:
2026-09-17 21:21:18 +08:00
parent a8a8dcfe1e
commit 571107f96d
27 changed files with 1375 additions and 24 deletions

View File

@@ -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)
}
})
}

View File

@@ -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)

View File

@@ -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

View File

@@ -0,0 +1,79 @@
package httpserver
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"git.yuntuo.hk/k2safe/orbit/internal/fault"
core "git.yuntuo.hk/k2safe/orbit/model"
"git.yuntuo.hk/k2safe/orbit/platform/auth"
"git.yuntuo.hk/k2safe/orbit/pulse/service"
)
type batchAuthStore struct {
service.Store
authCalls int
}
func batchHTTPToken(t *testing.T) string {
t.Helper()
token, _, err := auth.Issue(core.ID(strings.Repeat("1", 32)), core.ID(strings.Repeat("2", 32)))
if err != nil {
t.Fatal(err)
}
return "Bearer " + token
}
func (s *batchAuthStore) AuthorizeSubmit(context.Context, auth.Token) error {
s.authCalls++
return fault.Unauthorized
}
func TestBatchHTTPRejectsWholeMalformedEnvelopeBeforeWrites(t *testing.T) {
for _, test := range []struct{ body, key, code string }{
{`{"items":[{"idempotency_key":"a","task":{}},{"idempotency_key":"b","task":{"unexpected":1}}]}`, "", "INVALID_ARGUMENT"},
{`{"items":[{"idempotency_key":"a","task":{"retry":{"MaxAttempts":1}}}]}`, "", "INVALID_ARGUMENT"},
{`{"items":[{"idempotency_key":"a","task":{"payload":{"x":1,"x":2}}}]}`, "", "INVALID_ARGUMENT"},
{`{"items":[null]}`, "", "INVALID_ARGUMENT"},
{`{"items":[{"idempotency_key":"a","task":null}]}`, "", "INVALID_ARGUMENT"},
{`{"items":[{"idempotency_key":"a","task":{}},{"idempotency_key":"a","task":{}}]}`, "", "INVALID_ARGUMENT"},
{`{"items":[]}`, "", "INVALID_ARGUMENT"},
{`{"items":[{"idempotency_key":"a","task":{}}]}`, "whole-batch", "INVALID_ARGUMENT"},
{`{"items":[{"idempotency_key":"a","task":{"payload":"` + strings.Repeat("x", 512<<10) + `"}}]}`, "", "BODY_TOO_LARGE"},
} {
s := &batchAuthStore{}
h := New("", nil, service.New(s)).Handler
r := httptest.NewRequest("POST", "/v1/tasks:batch", strings.NewReader(test.body))
r.Header.Set("Content-Type", "application/json")
r.Header.Set("Authorization", batchHTTPToken(t))
if test.key != "" {
r.Header.Set("Idempotency-Key", test.key)
}
w := httptest.NewRecorder()
h.ServeHTTP(w, r)
var got struct {
Error struct {
Code string `json:"code"`
} `json:"error"`
}
if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil || got.Error.Code != test.code || s.authCalls != 0 {
t.Fatalf("got=%s auth=%d want=%s", w.Body.String(), s.authCalls, test.code)
}
}
}
func TestBatchHTTPAuthenticatesAllInvalidItems(t *testing.T) {
s := &batchAuthStore{}
r := httptest.NewRequest("POST", "/v1/tasks:batch", strings.NewReader(`{"items":[{"idempotency_key":"a","task":{}}]}`))
r.Header.Set("Content-Type", "application/json")
r.Header.Set("Authorization", batchHTTPToken(t))
w := httptest.NewRecorder()
New("", nil, service.New(s)).Handler.ServeHTTP(w, r)
if w.Code != http.StatusUnauthorized || s.authCalls != 1 {
t.Fatal(w.Body.String(), s.authCalls)
}
}

View File

@@ -27,6 +27,17 @@ import (
type endpoint func(*http.Request, auth.Token) (any, int, error)
func registerTasks(mux *http.ServeMux, s *service.Service) {
mux.HandleFunc("POST /v1/tasks:batch", wrap(func(r *http.Request, t auth.Token) (any, int, error) {
if len(r.Header.Values("Idempotency-Key")) != 0 {
return nil, 0, fault.Invalid("batch requires per-item idempotency_key, not an Idempotency-Key header")
}
var in service.BatchSubmitInput
if err := decodeBody(r, &in); err != nil {
return nil, 0, err
}
out, err := s.SubmitBatch(r.Context(), t, in)
return out, http.StatusOK, err
}))
mux.HandleFunc("POST /v1/tasks", wrap(func(r *http.Request, t auth.Token) (any, int, error) {
var in service.SubmitInput
if err := decodeBody(r, &in); err != nil {