From 0a12ed65336b348239a372945a469e25de1a922d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20Gra=C3=B1a?= Date: Thu, 17 Sep 2026 16:08:49 -0300 Subject: [PATCH] A cancel request racing an attempt's registration still cuts it requestCancel preempts the Run's registered attempt, but a worker that read its record before the request landed may be between reserving an attempt and registering it: the preempt found nothing and the attempt ran on a live context to its own end. TestCascadeSurvivesRestart hung this way when the cascade reached the child in that window. preemptAttempt now leaves the cause behind when no attempt is registered, and attemptContext applies it to the next forward registration, so the attempt starts already cut. The entry is consumed there or dropped at terminality. --- engine/cancel_test.go | 62 +++++++++++++++++++++++++++++++++++++++++++ engine/engine.go | 47 ++++++++++++++++++++++++-------- spec/01-model.md | 6 +++-- 3 files changed, 102 insertions(+), 13 deletions(-) diff --git a/engine/cancel_test.go b/engine/cancel_test.go index 6b68090..e37117f 100644 --- a/engine/cancel_test.go +++ b/engine/cancel_test.go @@ -537,3 +537,65 @@ func TestCascadeSurvivesRestart(t *testing.T) { t.Fatalf("child Wait = %+v, %v; want canceled with the parent's cause", res, err) } } + +// holdingStore is a driver.Store that, after applying the transition +// hold selects, blocks until released: it widens the window between an +// attempt's reservation and its registration for preemption. +type holdingStore struct { + driver.Store + hold func(driver.Transition) bool + held chan struct{} + release chan struct{} + once sync.Once +} + +func (s *holdingStore) ApplyTransition(ctx context.Context, id durable.RunID, t driver.Transition) error { + err := s.Store.ApplyTransition(ctx, id, t) + if err == nil && s.hold(t) { + s.once.Do(func() { + close(s.held) + <-s.release + }) + } + return err +} + +// A cancel request that lands after the worker read the record and +// reserved the attempt, but before the attempt registered for +// preemption, still cuts the attempt: it starts with its context +// already canceled instead of running to its own end. +func TestCancelBetweenReservationAndRegistrationCutsTheAttempt(t *testing.T) { + store := &holdingStore{ + Store: mem.New(), + hold: func(tr driver.Transition) bool { return tr.Cursor.StepID == "late/v1" }, + held: make(chan struct{}), + release: make(chan struct{}), + } + def := pipelinedef.New(pipelinedef.Config{ + ID: "late-cancel", + Steps: []pipelinedef.Step{stateless("late/v1", func(ctx context.Context, inv durable.Invocation) error { + <-ctx.Done() + return ctx.Err() + })}, + }) + _, pipes := startEngine(t, store, def) + run, _, err := pipes[0].Schedule(context.Background(), "r", nil) + if err != nil { + t.Fatalf("Schedule: %v", err) + } + select { + case <-store.held: + case <-time.After(5 * time.Second): + t.Fatal("attempt reservation never reached the store") + } + if err := run.Cancel(context.Background(), "late"); err != nil { + t.Fatalf("Cancel: %v", err) + } + close(store.release) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + res, err := run.Wait(ctx) + if err != nil || !res.Canceled() || res.Failure.Message != "late" { + t.Fatalf("Wait = %+v, %v; want canceled with cause late", res, err) + } +} diff --git a/engine/engine.go b/engine/engine.go index 61e9a00..c93e419 100644 --- a/engine/engine.go +++ b/engine/engine.go @@ -275,6 +275,14 @@ type Engine struct { started bool invalid map[durable.RunID]*InvalidRunError attemptCancel map[durable.RunID]attemptHandle + // pendingPreempt holds the cause of a cancel request that found no + // attempt registered for its Run. The Run's worker may have read + // its record before the request landed and be reserving a forward + // attempt on that stale read; attemptContext applies the cause when + // the attempt registers, so the request cuts it as if it had been + // in flight. An entry is consumed by that registration or dropped + // at terminality. + pendingPreempt map[durable.RunID]string // dirty marks the Runs whose store record changed under their // worker's carried copy — today only by a cancel request, which goes // through the engine. The worker takes the mark each iteration and @@ -320,15 +328,16 @@ type Engine struct { // scheduling is accepted. func New(store driver.Store, opts ...Option) *Engine { e := &Engine{ - store: store, - clock: wallClock{}, - logger: slog.Default(), - retry: defaultRetryPolicy, - textLimit: DefaultTextLimit, - invalid: make(map[durable.RunID]*InvalidRunError), - attemptCancel: make(map[durable.RunID]attemptHandle), - classCapacity: make(map[string]int), - awaitTimers: make(map[durable.RunID]chan struct{}), + store: store, + clock: wallClock{}, + logger: slog.Default(), + retry: defaultRetryPolicy, + textLimit: DefaultTextLimit, + invalid: make(map[durable.RunID]*InvalidRunError), + attemptCancel: make(map[durable.RunID]attemptHandle), + pendingPreempt: make(map[durable.RunID]string), + classCapacity: make(map[string]int), + awaitTimers: make(map[durable.RunID]chan struct{}), runClassCapacity: make(map[string]int), runClassQueue: make(map[string]int), @@ -701,12 +710,20 @@ func inAttempt(ctx context.Context) bool { } // attemptContext derives the per-attempt handler context and registers its -// cancel so a cancellation request can preempt the in-flight attempt. +// cancel so a cancellation request can preempt the in-flight attempt. A +// request that arrived before the registration, while the worker was +// reserving the attempt on a record read before the request landed, is +// applied here: a forward attempt starts already cut. func (e *Engine) attemptContext(id durable.RunID, phase durable.Phase) (context.Context, func()) { ctx, cancel := context.WithCancelCause(context.WithValue(e.baseCtx, attemptKey{}, id)) e.mu.Lock() e.attemptCancel[id] = attemptHandle{cancel: cancel, phase: phase} + cause, pending := e.pendingPreempt[id] + delete(e.pendingPreempt, id) e.mu.Unlock() + if pending && phase == durable.PhaseForward { + cancel(&durable.PreemptedError{Cause: cause}) + } return ctx, func() { e.mu.Lock() delete(e.attemptCancel, id) @@ -719,10 +736,15 @@ func (e *Engine) attemptContext(id durable.RunID, phase durable.Phase) (context. // if any, with a *PreemptedError cause carrying the request's cause. // The attempt resolves as canceled unless it returns success (see // runForward). An unwind attempt is never preempted: during a -// cancellation the unwind is the work. +// cancellation the unwind is the work. With no attempt registered the +// cause is left for the next registration (see attemptContext): the +// worker may be between reserving an attempt and registering it. func (e *Engine) preemptAttempt(id durable.RunID, cause string) { e.mu.Lock() h, ok := e.attemptCancel[id] + if !ok { + e.pendingPreempt[id] = cause + } e.mu.Unlock() if ok && h.phase == durable.PhaseForward { h.cancel(&durable.PreemptedError{Cause: cause}) @@ -1010,6 +1032,9 @@ func (e *Engine) completeRun(rec *driver.RunRecord) { "elapsed", rec.UpdatedAt.Sub(rec.CreatedAt)) e.emitRunTerminal(rec) e.releaseRunClass(rec.RunID) + e.mu.Lock() + delete(e.pendingPreempt, rec.RunID) + e.mu.Unlock() e.waiters.Notify(rec.RunID) e.awaitTargetDone(rec.RunID) } diff --git a/spec/01-model.md b/spec/01-model.md index 9771d4b..06d1b53 100644 --- a/spec/01-model.md +++ b/spec/01-model.md @@ -412,8 +412,10 @@ for a Step never attempted), and the Run's Failure names the Step. An attempt is executing: its context is canceled, with `context.Cause` a `*PreemptedError` carrying the request's cause — informational, for -middleware that labels spans; a handler needs only `ctx.Done()`. What -the attempt returns decides the Step, never the Run: +middleware that labels spans; a handler needs only `ctx.Done()`. An +attempt the worker was reserving on a record read before the request +landed starts with its context already canceled. What the attempt +returns decides the Step, never the Run: - success commits the Step's State as usual, and the Step unwinds with the others;