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;