From c395a5a9ad9c364e69853b8aff18edbcc11fe48e Mon Sep 17 00:00:00 2001 From: ecordell Date: Sun, 8 Feb 2026 11:07:40 -0500 Subject: [PATCH] middleware: new middleware package built on top of state --- .golangci.yaml | 5 + go.mod | 2 +- go.sum | 4 +- queue/handlers.go | 55 +- queue/handlers_test.go | 489 ++++------------ state/FORMAL.md | 165 ++++-- state/doc.go | 3 +- state/formal_test.go | 252 +++++++++ state/middleware/doc.go | 23 + state/middleware/middleware.go | 75 +++ state/middleware/middleware_test.go | 251 ++++++++ state/state.go | 311 ++++++++-- state/state_test.go | 848 +++++++++++++++++++++++----- 13 files changed, 1791 insertions(+), 692 deletions(-) create mode 100644 state/middleware/doc.go create mode 100644 state/middleware/middleware.go create mode 100644 state/middleware/middleware_test.go diff --git a/.golangci.yaml b/.golangci.yaml index ae5af32..5747e3c 100644 --- a/.golangci.yaml +++ b/.golangci.yaml @@ -32,6 +32,11 @@ linters: - "whitespace" - "unused" settings: + govet: + disable: + # inline reports spurious "cannot inline: type parameter inference is + # not yet supported" errors on calls to generic stdlib functions. + - "inline" staticcheck: checks: - "all" diff --git a/go.mod b/go.mod index 4fe78f6..9b04973 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,7 @@ module github.com/authzed/controller-idioms go 1.25.0 require ( - github.com/authzed/ctxkey v0.0.0-20260210154927-ca132876f62c + github.com/authzed/ctxkey v0.1.0 github.com/cespare/xxhash/v2 v2.3.0 github.com/davecgh/go-spew v1.1.1 github.com/fsnotify/fsnotify v1.9.0 diff --git a/go.sum b/go.sum index 3df6ce6..2565f13 100644 --- a/go.sum +++ b/go.sum @@ -10,8 +10,8 @@ github.com/NYTimes/gziphandler v1.1.1 h1:ZUDjpQae29j0ryrS0u/B8HZfJBtBQHjqw2rQ2cq github.com/NYTimes/gziphandler v1.1.1/go.mod h1:n/CVRwUEOgIxrgPvAQhUUr9oeUtvrhMomdKFjzJNB0c= github.com/antlr4-go/antlr/v4 v4.13.1 h1:SqQKkuVZ+zWkMMNkjy5FZe5mr5WURWnlpmOuzYWrPrQ= github.com/antlr4-go/antlr/v4 v4.13.1/go.mod h1:GKmUxMtwp6ZgGwZSva4eWPC5mS6vUAmOABFgjdkM7Nw= -github.com/authzed/ctxkey v0.0.0-20260210154927-ca132876f62c h1:kcjpI9wcj5yFYWJ7bEeiY6hqH2I6Es3lVz6PB7RKblQ= -github.com/authzed/ctxkey v0.0.0-20260210154927-ca132876f62c/go.mod h1:ve6DXRv9l2HcM2lxSUmZKWXl7Iq/00xYixBuHOtST+4= +github.com/authzed/ctxkey v0.1.0 h1:wtQnpKgjzxF//qotuiR/Toh691lbY2ZFynsQFQ57mrA= +github.com/authzed/ctxkey v0.1.0/go.mod h1:ve6DXRv9l2HcM2lxSUmZKWXl7Iq/00xYixBuHOtST+4= github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= github.com/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM= diff --git a/queue/handlers.go b/queue/handlers.go index 77398be..2a543de 100644 --- a/queue/handlers.go +++ b/queue/handlers.go @@ -57,7 +57,7 @@ import ( "github.com/authzed/controller-idioms/state" ) -// Done creates a handler that marks the current queue key as finished and terminates the pipeline. +// Done is a handler that marks the current queue key as finished and terminates the pipeline. // This is equivalent to calling queue.NewQueueOperationsCtx().Done(ctx) and returning. // // Usage: @@ -65,15 +65,13 @@ import ( // pipeline := state.Sequence( // validateInput, // processResource, -// queue.Done(), // Stop here - processing complete +// queue.Done, // Stop here - processing complete // ) -func Done() state.NewStep { - return state.NewTerminalStepFunc(func(ctx context.Context) { - NewQueueOperationsCtx().Done(ctx) - }) -} +var Done = state.NewTerminalStepFunc(func(ctx context.Context) { + NewQueueOperationsCtx().Done(ctx) +}) -// Requeue creates a handler that requeues the current key immediately and terminates the pipeline. +// Requeue is a handler that requeues the current key immediately and terminates the pipeline. // This is equivalent to calling queue.NewQueueOperationsCtx().Requeue(ctx) and returning. // // Usage: @@ -81,13 +79,11 @@ func Done() state.NewStep { // pipeline := state.Decision( // resourceReady, // continueProcessing, -// queue.Requeue(), // Not ready - try again immediately +// queue.Requeue, // Not ready - try again immediately // ) -func Requeue() state.NewStep { - return state.NewTerminalStepFunc(func(ctx context.Context) { - NewQueueOperationsCtx().Requeue(ctx) - }) -} +var Requeue = state.NewTerminalStepFunc(func(ctx context.Context) { + NewQueueOperationsCtx().Requeue(ctx) +}) // RequeueAfter creates a handler that requeues the current key after the specified duration // and terminates the pipeline. @@ -160,34 +156,3 @@ func RequeueAPIErr(err error) state.NewStep { NewQueueOperationsCtx().RequeueAPIErr(ctx, err) }) } - -// OnError creates a handler that executes different queue operations based on whether -// an error occurred in the context. -// -// Usage: -// -// pipeline := state.Sequence( -// riskyOperation, -// queue.OnError( -// queue.RequeueErr(fmt.Errorf("operation failed")), // If error -// queue.Done(), // If success -// ), -// ) -func OnError(errorHandler, successHandler state.NewStep) state.NewStep { - return func(next state.Step) state.Step { - errStep := errorHandler(next) - okStep := successHandler(next) - return state.StepFunc(func(ctx context.Context) state.Step { - if ctx.Err() != nil { - if errStep != nil { - return errStep.Run(ctx) - } - return nil - } - if okStep != nil { - return okStep.Run(ctx) - } - return nil - }) - } -} diff --git a/queue/handlers_test.go b/queue/handlers_test.go index d98dcfa..a1ea069 100644 --- a/queue/handlers_test.go +++ b/queue/handlers_test.go @@ -7,6 +7,8 @@ import ( "testing" "time" + "github.com/stretchr/testify/require" + "github.com/authzed/ctxkey" "github.com/authzed/controller-idioms/queue" @@ -15,151 +17,95 @@ import ( ) func TestDone(t *testing.T) { - ctx := t.Context() fakeQueue := &fake.FakeInterface{} + ctxWithQueue := queue.NewQueueOperationsCtx().WithValue(t.Context(), fakeQueue) - // Set up context with fake queue - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) - - // Create and run pipeline + var afterExecuted bool pipeline := state.Sequence( + state.Do(func(ctx context.Context) context.Context { return ctx }), + queue.Done, state.Do(func(ctx context.Context) context.Context { - // This should execute - return ctx - }), - queue.Done(), - state.Do(func(ctx context.Context) context.Context { - t.Error("This should not execute - pipeline should have terminated") + afterExecuted = true return ctx }), ) state.Run(ctxWithQueue, pipeline) - // Verify Done was called - if fakeQueue.DoneCallCount() != 1 { - t.Errorf("Expected Done to be called once, got %d calls", fakeQueue.DoneCallCount()) - } + require.Equal(t, 1, fakeQueue.DoneCallCount()) + require.False(t, afterExecuted, "pipeline should terminate after Done") } func TestRequeue(t *testing.T) { - ctx := t.Context() fakeQueue := &fake.FakeInterface{} + ctxWithQueue := queue.NewQueueOperationsCtx().WithValue(t.Context(), fakeQueue) - // Set up context with fake queue - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) - - executed := false + var afterExecuted bool pipeline := state.Sequence( + state.Do(func(ctx context.Context) context.Context { return ctx }), + queue.Requeue, state.Do(func(ctx context.Context) context.Context { - executed = true - return ctx - }), - queue.Requeue(), - state.Do(func(ctx context.Context) context.Context { - t.Error("This should not execute - pipeline should have terminated") + afterExecuted = true return ctx }), ) state.Run(ctxWithQueue, pipeline) - // Verify execution and requeue - if !executed { - t.Error("Expected first action to execute") - } - if fakeQueue.RequeueCallCount() != 1 { - t.Errorf("Expected Requeue to be called once, got %d calls", fakeQueue.RequeueCallCount()) - } + require.Equal(t, 1, fakeQueue.RequeueCallCount()) + require.False(t, afterExecuted, "pipeline should terminate after Requeue") } func TestRequeueAfter(t *testing.T) { - ctx := t.Context() fakeQueue := &fake.FakeInterface{} + ctxWithQueue := queue.NewQueueOperationsCtx().WithValue(t.Context(), fakeQueue) duration := 30 * time.Second - // Set up context with fake queue - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) - pipeline := state.Sequence( - state.Do(func(ctx context.Context) context.Context { - // Setup action - return ctx - }), + state.Do(func(ctx context.Context) context.Context { return ctx }), queue.RequeueAfter(duration), state.Do(func(ctx context.Context) context.Context { - t.Error("This should not execute - pipeline should have terminated") + t.Error("pipeline should terminate after RequeueAfter") return ctx }), ) state.Run(ctxWithQueue, pipeline) - // Verify RequeueAfter was called with correct duration - if fakeQueue.RequeueAfterCallCount() != 1 { - t.Errorf("Expected RequeueAfter to be called once, got %d calls", fakeQueue.RequeueAfterCallCount()) - } - if fakeQueue.RequeueAfterArgsForCall(0) != duration { - t.Errorf("Expected RequeueAfter duration %v, got %v", duration, fakeQueue.RequeueAfterArgsForCall(0)) - } + require.Equal(t, 1, fakeQueue.RequeueAfterCallCount()) + require.Equal(t, duration, fakeQueue.RequeueAfterArgsForCall(0)) } func TestRequeueErr(t *testing.T) { - ctx := t.Context() fakeQueue := &fake.FakeInterface{} + ctxWithQueue := queue.NewQueueOperationsCtx().WithValue(t.Context(), fakeQueue) testErr := errors.New("test error") - // Set up context with fake queue - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) - pipeline := state.Sequence( - state.Do(func(ctx context.Context) context.Context { - // Setup action - return ctx - }), + state.Do(func(ctx context.Context) context.Context { return ctx }), queue.RequeueErr(testErr), state.Do(func(ctx context.Context) context.Context { - t.Error("This should not execute - pipeline should have terminated") + t.Error("pipeline should terminate after RequeueErr") return ctx }), ) state.Run(ctxWithQueue, pipeline) - // Verify RequeueErr was called with correct error - if fakeQueue.RequeueErrCallCount() != 1 { - t.Errorf("Expected RequeueErr to be called once, got %d calls", fakeQueue.RequeueErrCallCount()) - } - if !errors.Is(fakeQueue.RequeueErrArgsForCall(0), testErr) { - t.Errorf("Expected RequeueErr error %v, got %v", testErr, fakeQueue.RequeueErrArgsForCall(0)) - } + require.Equal(t, 1, fakeQueue.RequeueErrCallCount()) + require.ErrorIs(t, fakeQueue.RequeueErrArgsForCall(0), testErr) } func TestRequeueErrFromWithinStep(t *testing.T) { - ctx := t.Context() fakeQueue := &fake.FakeInterface{} + ctxWithQueue := queue.NewQueueOperationsCtx().WithValue(t.Context(), fakeQueue) - // Set up context with fake queue - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) - - var executedBefore bool - var executedAfter bool + var executedBefore, executedAfter bool - // Simulate a step that encounters an error and returns the error handler directly - riskyOperation := state.NewStepFunc(func(ctx context.Context, next state.Step) state.Step { + riskyOperation := state.NewStepFunc(func(ctx context.Context, _ state.Step) state.Step { executedBefore = true - // Simulate some operation that fails err := errors.New("operation failed") - if err != nil { - // Return the error handler directly, short-circuiting the pipeline - return queue.RequeueErr(err).Step().Run(ctx) - } - return state.Continue(ctx, next) + return queue.RequeueErr(err).Step().Run(ctx) }) pipeline := state.Sequence( @@ -172,62 +118,32 @@ func TestRequeueErrFromWithinStep(t *testing.T) { state.Run(ctxWithQueue, pipeline) - // Verify the operation executed but subsequent steps did not - if !executedBefore { - t.Error("Expected risky operation to execute") - } - if executedAfter { - t.Error("Expected pipeline to terminate after error, but subsequent step executed") - } - - // Verify RequeueErr was called - if fakeQueue.RequeueErrCallCount() != 1 { - t.Errorf("Expected RequeueErr to be called once, got %d calls", fakeQueue.RequeueErrCallCount()) - } + require.True(t, executedBefore) + require.False(t, executedAfter, "pipeline should terminate after inline error") + require.Equal(t, 1, fakeQueue.RequeueErrCallCount()) } func TestRequeueAPIErr(t *testing.T) { - ctx := t.Context() fakeQueue := &fake.FakeInterface{} + ctxWithQueue := queue.NewQueueOperationsCtx().WithValue(t.Context(), fakeQueue) testErr := errors.New("API error") - // Set up context with fake queue - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) - - pipeline := queue.RequeueAPIErr(testErr) - state.Run(ctxWithQueue, pipeline) + state.Run(ctxWithQueue, queue.RequeueAPIErr(testErr)) - // Verify RequeueAPIErr was called with correct error - if fakeQueue.RequeueAPIErrCallCount() != 1 { - t.Errorf("Expected RequeueAPIErr to be called once, got %d calls", fakeQueue.RequeueAPIErrCallCount()) - } - if !errors.Is(fakeQueue.RequeueAPIErrArgsForCall(0), testErr) { - t.Errorf("Expected RequeueAPIErr error %v, got %v", testErr, fakeQueue.RequeueAPIErrArgsForCall(0)) - } + require.Equal(t, 1, fakeQueue.RequeueAPIErrCallCount()) + require.ErrorIs(t, fakeQueue.RequeueAPIErrArgsForCall(0), testErr) } func TestRequeueAPIErrFromWithinStep(t *testing.T) { - ctx := t.Context() fakeQueue := &fake.FakeInterface{} + ctxWithQueue := queue.NewQueueOperationsCtx().WithValue(t.Context(), fakeQueue) - // Set up context with fake queue - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) - - var executedBefore bool - var executedAfter bool + var executedBefore, executedAfter bool - // Simulate a Kubernetes API call that fails - callKubernetesAPI := state.NewStepFunc(func(ctx context.Context, next state.Step) state.Step { + callKubernetesAPI := state.NewStepFunc(func(ctx context.Context, _ state.Step) state.Step { executedBefore = true - // Simulate an API error apiErr := errors.New("deployments.apps \"myapp\" not found") - if apiErr != nil { - // Return the API error handler directly - return queue.RequeueAPIErr(apiErr).Step().Run(ctx) - } - return state.Continue(ctx, next) + return queue.RequeueAPIErr(apiErr).Step().Run(ctx) }) pipeline := state.Sequence( @@ -240,18 +156,9 @@ func TestRequeueAPIErrFromWithinStep(t *testing.T) { state.Run(ctxWithQueue, pipeline) - // Verify the API call executed but subsequent steps did not - if !executedBefore { - t.Error("Expected API call to execute") - } - if executedAfter { - t.Error("Expected pipeline to terminate after API error, but subsequent step executed") - } - - // Verify RequeueAPIErr was called - if fakeQueue.RequeueAPIErrCallCount() != 1 { - t.Errorf("Expected RequeueAPIErr to be called once, got %d calls", fakeQueue.RequeueAPIErrCallCount()) - } + require.True(t, executedBefore) + require.False(t, executedAfter, "pipeline should terminate after inline API error") + require.Equal(t, 1, fakeQueue.RequeueAPIErrCallCount()) } func TestConditionalRequeue(t *testing.T) { @@ -261,34 +168,20 @@ func TestConditionalRequeue(t *testing.T) { expectRequeue bool expectContinuation bool }{ - { - name: "requeue when condition is true", - condition: true, - expectRequeue: true, - expectContinuation: false, - }, - { - name: "continue when condition is false", - condition: false, - expectRequeue: false, - expectContinuation: true, - }, + {"requeue when condition is true", true, true, false}, + {"continue when condition is false", false, false, true}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - ctx := context.Background() fakeQueue := &fake.FakeInterface{} + ctxWithQueue := queue.NewQueueOperationsCtx().WithValue(t.Context(), fakeQueue) - // Set up context with fake queue - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) - - continued := false + var continued bool pipeline := state.Sequence( state.When( func(_ context.Context) bool { return tt.condition }, - queue.Requeue(), + queue.Requeue, ), state.Do(func(ctx context.Context) context.Context { continued = true @@ -298,22 +191,12 @@ func TestConditionalRequeue(t *testing.T) { state.Run(ctxWithQueue, pipeline) - // Verify requeue behavior - requeueCount := fakeQueue.RequeueCallCount() - if tt.expectRequeue && requeueCount != 1 { - t.Errorf("Expected requeue to be called once, got %d calls", requeueCount) - } - if !tt.expectRequeue && requeueCount != 0 { - t.Errorf("Expected requeue not to be called, got %d calls", requeueCount) - } - - // Verify continuation behavior - if tt.expectContinuation && !continued { - t.Error("Expected pipeline to continue") - } - if !tt.expectContinuation && continued { - t.Error("Expected pipeline to terminate") + if tt.expectRequeue { + require.Equal(t, 1, fakeQueue.RequeueCallCount()) + } else { + require.Equal(t, 0, fakeQueue.RequeueCallCount()) } + require.Equal(t, tt.expectContinuation, continued) }) } } @@ -327,30 +210,16 @@ func TestConditionalRequeueAfter(t *testing.T) { expectRequeue bool expectContinuation bool }{ - { - name: "requeue after when condition is true", - condition: true, - expectRequeue: true, - expectContinuation: false, - }, - { - name: "continue when condition is false", - condition: false, - expectRequeue: false, - expectContinuation: true, - }, + {"requeue after when condition is true", true, true, false}, + {"continue when condition is false", false, false, true}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - ctx := context.Background() fakeQueue := &fake.FakeInterface{} + ctxWithQueue := queue.NewQueueOperationsCtx().WithValue(t.Context(), fakeQueue) - // Set up context with fake queue - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) - - continued := false + var continued bool pipeline := state.Sequence( state.When( func(_ context.Context) bool { return tt.condition }, @@ -364,27 +233,13 @@ func TestConditionalRequeueAfter(t *testing.T) { state.Run(ctxWithQueue, pipeline) - // Verify requeue behavior - requeueCount := fakeQueue.RequeueAfterCallCount() - if tt.expectRequeue && requeueCount != 1 { - t.Errorf("Expected requeueAfter to be called once, got %d calls", requeueCount) - } - if !tt.expectRequeue && requeueCount != 0 { - t.Errorf("Expected requeueAfter not to be called, got %d calls", requeueCount) - } - - // Verify duration if requeued - if tt.expectRequeue && fakeQueue.RequeueAfterArgsForCall(0) != duration { - t.Errorf("Expected duration %v, got %v", duration, fakeQueue.RequeueAfterArgsForCall(0)) - } - - // Verify continuation behavior - if tt.expectContinuation && !continued { - t.Error("Expected pipeline to continue") - } - if !tt.expectContinuation && continued { - t.Error("Expected pipeline to terminate") + if tt.expectRequeue { + require.Equal(t, 1, fakeQueue.RequeueAfterCallCount()) + require.Equal(t, duration, fakeQueue.RequeueAfterArgsForCall(0)) + } else { + require.Equal(t, 0, fakeQueue.RequeueAfterCallCount()) } + require.Equal(t, tt.expectContinuation, continued) }) } } @@ -396,34 +251,20 @@ func TestConditionalDone(t *testing.T) { expectDone bool expectContinuation bool }{ - { - name: "done when condition is true", - condition: true, - expectDone: true, - expectContinuation: false, - }, - { - name: "continue when condition is false", - condition: false, - expectDone: false, - expectContinuation: true, - }, + {"done when condition is true", true, true, false}, + {"continue when condition is false", false, false, true}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - ctx := context.Background() fakeQueue := &fake.FakeInterface{} + ctxWithQueue := queue.NewQueueOperationsCtx().WithValue(t.Context(), fakeQueue) - // Set up context with fake queue - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) - - continued := false + var continued bool pipeline := state.Sequence( state.When( func(_ context.Context) bool { return tt.condition }, - queue.Done(), + queue.Done, ), state.Do(func(ctx context.Context) context.Context { continued = true @@ -433,99 +274,21 @@ func TestConditionalDone(t *testing.T) { state.Run(ctxWithQueue, pipeline) - // Verify done behavior - doneCount := fakeQueue.DoneCallCount() - if tt.expectDone && doneCount != 1 { - t.Errorf("Expected done to be called once, got %d calls", doneCount) - } - if !tt.expectDone && doneCount != 0 { - t.Errorf("Expected done not to be called, got %d calls", doneCount) - } - - // Verify continuation behavior - if tt.expectContinuation && !continued { - t.Error("Expected pipeline to continue") - } - if !tt.expectContinuation && continued { - t.Error("Expected pipeline to terminate") + if tt.expectDone { + require.Equal(t, 1, fakeQueue.DoneCallCount()) + } else { + require.Equal(t, 0, fakeQueue.DoneCallCount()) } + require.Equal(t, tt.expectContinuation, continued) }) } } -func TestOnError(t *testing.T) { - tests := []struct { - name string - hasError bool - expectDone bool - expectRequeue bool - }{ - { - name: "error handler when error present", - hasError: true, - expectDone: false, - expectRequeue: true, - }, - { - name: "success handler when no error", - hasError: false, - expectDone: true, - expectRequeue: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - ctx := context.Background() - fakeQueue := &fake.FakeInterface{} - - // Set up context with fake queue - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) - - // Add error to context if needed - if tt.hasError { - cancelCtx, cancel := context.WithCancel(ctxWithQueue) - cancel() // This sets ctx.Err() - ctxWithQueue = cancelCtx - } - - pipeline := queue.OnError( - queue.Requeue(), // Error handler - queue.Done(), // Success handler - ) - - state.Run(ctxWithQueue, pipeline) - - // Verify behavior based on error state - doneCount := fakeQueue.DoneCallCount() - requeueCount := fakeQueue.RequeueCallCount() - - if tt.expectDone && doneCount != 1 { - t.Errorf("Expected done to be called once, got %d calls", doneCount) - } - if !tt.expectDone && doneCount != 0 { - t.Errorf("Expected done not to be called, got %d calls", doneCount) - } - - if tt.expectRequeue && requeueCount != 1 { - t.Errorf("Expected requeue to be called once, got %d calls", requeueCount) - } - if !tt.expectRequeue && requeueCount != 0 { - t.Errorf("Expected requeue not to be called, got %d calls", requeueCount) - } - }) - } -} - -// Test realistic controller scenarios +// TestControllerScenarios tests realistic controller pipelines. func TestControllerScenarios(t *testing.T) { t.Run("successful processing flow", func(t *testing.T) { - ctx := context.Background() fakeQueue := &fake.FakeInterface{} - - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) + ctxWithQueue := queue.NewQueueOperationsCtx().WithValue(t.Context(), fakeQueue) var executionOrder []string @@ -542,33 +305,18 @@ func TestControllerScenarios(t *testing.T) { executionOrder = append(executionOrder, "finalize") return ctx }), - queue.Done(), + queue.Done, ) state.Run(ctxWithQueue, pipeline) - expectedOrder := []string{"validate", "process", "finalize"} - if len(executionOrder) != len(expectedOrder) { - t.Fatalf("Expected %d executions, got %d: %v", len(expectedOrder), len(executionOrder), executionOrder) - } - - for i, expected := range expectedOrder { - if executionOrder[i] != expected { - t.Errorf("Expected execution[%d] = %q, got %q", i, expected, executionOrder[i]) - } - } - - if fakeQueue.DoneCallCount() != 1 { - t.Errorf("Expected done to be called once, got %d calls", fakeQueue.DoneCallCount()) - } + require.Equal(t, []string{"validate", "process", "finalize"}, executionOrder) + require.Equal(t, 1, fakeQueue.DoneCallCount()) }) t.Run("resource not ready - requeue after delay", func(t *testing.T) { - ctx := context.Background() fakeQueue := &fake.FakeInterface{} - - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) + ctxWithQueue := queue.NewQueueOperationsCtx().WithValue(t.Context(), fakeQueue) var executionOrder []string @@ -578,46 +326,30 @@ func TestControllerScenarios(t *testing.T) { return ctx }), state.Decision( - func(_ context.Context) bool { - return false // Dependencies not ready - }, - // Dependencies ready + func(_ context.Context) bool { return false }, state.Sequence( state.Do(func(ctx context.Context) context.Context { executionOrder = append(executionOrder, "process") return ctx }), - queue.Done(), + queue.Done, ), - // Dependencies not ready queue.RequeueAfter(5*time.Minute), ), ) state.Run(ctxWithQueue, pipeline) - expectedOrder := []string{"check-dependencies"} - if len(executionOrder) != len(expectedOrder) { - t.Fatalf("Expected %d executions, got %d: %v", len(expectedOrder), len(executionOrder), executionOrder) - } - - if fakeQueue.RequeueAfterCallCount() != 1 { - t.Errorf("Expected requeueAfter to be called once, got %d calls", fakeQueue.RequeueAfterCallCount()) - } - - if fakeQueue.RequeueAfterArgsForCall(0) != 5*time.Minute { - t.Errorf("Expected 5 minute delay, got %v", fakeQueue.RequeueAfterArgsForCall(0)) - } + require.Equal(t, []string{"check-dependencies"}, executionOrder) + require.Equal(t, 1, fakeQueue.RequeueAfterCallCount()) + require.Equal(t, 5*time.Minute, fakeQueue.RequeueAfterArgsForCall(0)) }) t.Run("validation error - requeue with error", func(t *testing.T) { - ctx := context.Background() fakeQueue := &fake.FakeInterface{} - - queueCtx := queue.NewQueueOperationsCtx() - ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) - + ctxWithQueue := queue.NewQueueOperationsCtx().WithValue(t.Context(), fakeQueue) testErr := errors.New("validation failed") + var executionOrder []string pipeline := state.Sequence( @@ -626,40 +358,27 @@ func TestControllerScenarios(t *testing.T) { return ctx }), state.Decision( - func(_ context.Context) bool { - return false // Validation failed - }, - // Validation passed + func(_ context.Context) bool { return false }, state.Sequence( state.Do(func(ctx context.Context) context.Context { executionOrder = append(executionOrder, "process") return ctx }), - queue.Done(), + queue.Done, ), - // Validation failed queue.RequeueErr(testErr), ), ) state.Run(ctxWithQueue, pipeline) - expectedOrder := []string{"validate"} - if len(executionOrder) != len(expectedOrder) { - t.Fatalf("Expected %d executions, got %d: %v", len(expectedOrder), len(executionOrder), executionOrder) - } - - if fakeQueue.RequeueErrCallCount() != 1 { - t.Errorf("Expected requeueErr to be called once, got %d calls", fakeQueue.RequeueErrCallCount()) - } - - if !errors.Is(fakeQueue.RequeueErrArgsForCall(0), testErr) { - t.Errorf("Expected error %v, got %v", testErr, fakeQueue.RequeueErrArgsForCall(0)) - } + require.Equal(t, []string{"validate"}, executionOrder) + require.Equal(t, 1, fakeQueue.RequeueErrCallCount()) + require.ErrorIs(t, fakeQueue.RequeueErrArgsForCall(0), testErr) }) } -// Example demonstrating the clean controller pattern this enables +// Example demonstrates the clean controller pattern this enables. func Example() { ctx := context.Background() fakeQueue := &fake.FakeInterface{} @@ -667,43 +386,29 @@ func Example() { queueCtx := queue.NewQueueOperationsCtx() ctxWithQueue := queueCtx.WithValue(ctx, fakeQueue) - // Simulate resource state resourceReadyKey := ctxkey.New[bool]() ctxWithResource := resourceReadyKey.Set(ctxWithQueue, true) - // Clean controller pipeline using queue handlers controllerPipeline := state.Sequence( - // Set finalizer state.Do(func(ctx context.Context) context.Context { fmt.Println("Setting finalizer") return ctx }), - - // Check if resource is ready to process state.When( func(ctx context.Context) bool { return !resourceReadyKey.MustValue(ctx) }, - queue.RequeueAfter(30*time.Second), // Wait 30 seconds if not ready + queue.RequeueAfter(30*time.Second), ), - - // Process the resource state.Do(func(ctx context.Context) context.Context { fmt.Println("Processing resource") return ctx }), - - // Validate processing completed successfully state.When( - func(_ context.Context) bool { - // In real code, this would check if processing completed - return false // Assume success - }, - queue.Requeue(), + func(_ context.Context) bool { return false }, + queue.Requeue, ), - - // Mark as done - queue.Done(), + queue.Done, ) state.Run(ctxWithResource, controllerPipeline) diff --git a/state/FORMAL.md b/state/FORMAL.md index 7a453a0..963a51d 100644 --- a/state/FORMAL.md +++ b/state/FORMAL.md @@ -33,11 +33,11 @@ Sequence( Without laws: context values can vanish, causing hard-to-debug failures. -**4. Middleware Correctness** - Kleisli category structure tells us middleware must preserve composition: +**4. Middleware Correctness** - the categorical structure tells us what middleware can and cannot do: -- Build middleware from safe combinators (correct by construction) -- Or write property tests verifying endofunctor laws -- No ad-hoc wrappers that break when composed +- Middleware built from `Sequence`/`Do` stays inside the algebra and inherits its laws — but its after-hooks are provably absorbed by terminal steps (see "Bracket Middleware" below) +- `MakeMiddleware` is a **bracket**: a control operator with its own laws (B1–B4) that guarantees paired before/after hooks on every outcome +- Either way: no ad-hoc wrappers with unstated properties ## Category Theory Foundation @@ -134,7 +134,7 @@ Every monad automatically generates a Kleisli category - so just having one isn' The Kleisli category perspective becomes valuable when we have **transformations on `NewStep` that preserve composition**: -**Middleware as Endofunctors**: Functions of type `func(NewStep) NewStep` are endofunctors on the Kleisli category. When middleware preserves the categorical structure (identity and composition), we get: +**Middleware as transformations on Kleisli arrows**: Functions of type `func(NewStep) NewStep` map arrows of the Kleisli category to arrows of the Kleisli category: ```go // Middleware transforms NewSteps while preserving composition @@ -153,46 +153,34 @@ func LoggingMiddleware(name string) Middleware { } } } - -// Middleware composition preserves the categorical laws: -// middleware(Sequence(a, b)) ≈ Sequence(middleware(a), middleware(b)) ``` -The Kleisli category perspective tells us that **correct middleware must be an endofunctor** - it must preserve identity and composition. While Go's type system doesn't enforce this, the categorical perspective guides us to: - -**1. Build safe middleware combinators** that structurally preserve composition: +**A note on the endofunctor laws.** An idealized middleware would preserve +identity and composition: -```go -// This combinator guarantees correct behavior by construction -func MakeMiddleware( - before func(context.Context) context.Context, - after func(context.Context) context.Context, -) Middleware { - return func(step NewStep) NewStep { - return Sequence( - Do(before), // ← Built from composition primitives - step, // ← that already satisfy the laws - Do(after), - ) - } -} +```text +mw(Sequence(a, b)) ≈ Sequence(mw(a), mw(b)) +mw(Noop) ≈ Noop ``` -Since `Sequence` and `Do` already satisfy the categorical laws, middleware built from them **inherits** those properties. No manual proof needed - it's correct by construction. - -**2. Test endofunctor properties** for custom middleware: +Observe that no middleware with effectful hooks satisfies these laws strictly: +the left side of the first law fires one hook pair around the whole sequence, +while the right side fires a pair around *each* step — observably different +for any hook that logs, times, or counts. The laws describe a choice of +**granularity**, not a property a hook-bearing middleware can have. -```go -// Property-based test that middleware is an endofunctor -func TestMiddlewarePreservesComposition(t *testing.T, mw Middleware) { - // Law 1: mw(Sequence(a, b)) ≈ Sequence(mw(a), mw(b)) - // Law 2: mw(Noop) ≈ Noop -} -``` +The library therefore does not rely on middleware being endofunctors. Instead, +`AmbientDispatch(a, b, ...)` makes the granularity explicit: it applies the +ambient middleware to each listed step — literally constructing +`Sequence(mw(a), mw(b), ...)`, the per-step side of the law. Passing a +composite (e.g. a `Sequence`) as one argument selects the unit side. The +caller chooses; no law is silently assumed. -**Without the Kleisli perspective**: We'd write ad-hoc wrappers and hope they work correctly when composed. +**Without this perspective**: We'd write ad-hoc wrappers and hope they compose. -**With the Kleisli perspective**: We know middleware must preserve categorical structure, so we either build it from compositional primitives (correct by construction) or write property tests to verify the endofunctor laws. +**With it**: Middleware granularity is an explicit decision at the dispatch +site, and `MakeMiddleware` provides a constructor with proven laws of its own +(next section). #### Kleisli Laws @@ -207,6 +195,107 @@ func TestMiddlewarePreservesComposition(t *testing.T, mw Middleware) { These laws are verified in our test suite (see `formal_test.go`). +### Bracket Middleware: MakeMiddleware and the Absorption Theorem + +`MakeMiddleware(before, after)` is **not** a Kleisli composition — and cannot +be. This section states why, and the laws it satisfies instead. + +#### The Absorption Theorem + +A **terminal step** ignores its continuation and returns nil +(`NewTerminalStepFunc`, `Terminal`, `queue.Done`, `queue.Requeue`, ...). +Terminal steps are **left zeros** of `Sequence`: + +```text +Sequence(t, x) = t for terminal t and any x +``` + +This is immediate from the definition: `Sequence(t, x)(next) = t(x(next))`, +and `t` discards its argument. It follows that any middleware built purely +from Kleisli composition loses its after-hook to a terminal step: + +```text +Sequence(Do(before), t, Do(after)) = Sequence(Do(before), t) +``` + +The after-hook sits downstream of `t` — the absorbed position. The same holds +on the cancellation path: `Continue` stops the pipeline before reaching it. + +**Consequence**: "stays inside the `Sequence`/`Do` algebra" and "after-hook +always fires" are mutually exclusive. A finalizer is not a Kleisli arrow. +(This is the same fact as in effect systems generally: `finally`/`bracket` is +not definable from `bind` alone; it requires a control operator.) Since +`MakeMiddleware` exists for finalizer-shaped concerns — timers, spans, +metrics — and every controller pipeline ends in a terminal step, it is +implemented as a control operator: it intercepts the wrapped step's +continuation with a sentinel, observes whether the step continued, runs the +after-hook unconditionally, and resumes the pipeline outside the step's +dynamic extent (a small delimited continuation). + +#### Bracket Laws + +In exchange for leaving the pure algebra, `MakeMiddleware` satisfies the +bracket laws: + +- **B1 (pairing)**: the after-hook runs exactly once whenever the before-hook + ran, on every outcome **the program survives** — continue, terminal, + cancelled, or a panic that recovery middleware above will catch. On an + unrecovered panic the process terminates: no hooks fire during fatal + unwinding and the panic value propagates untouched. + + Panic recovery is a dynamic-scope property, and context is this library's + dynamic scope: `middleware.Recover` marks its extent via + `WithPanicRecovery`, and brackets release during unwinding only inside a + marked scope (without calling `recover` — the panic continues to the + recovery middleware). `Parallel` clears the marker for branch goroutines, + since recovery cannot cross goroutine boundaries; a branch with its own + recovery middleware re-establishes it. Frameworks that recover panics + outside the pipeline can opt in by calling `WithPanicRecovery` before + `Run`. +- **B2 (nesting)**: brackets are properly nested (LIFO), like `defer` or + `try/finally`: `mwA(mwB(s))` yields `beforeA, beforeB, s, afterB, afterA`. +- **B3 (transparency)**: the bracket is invisible to the surrounding pipeline. + Terminality is preserved (`mw(t)` is terminal when `t` is), and when the + step continues with context `ctx′`, downstream observes `after(ctx′)`. +- **B4 (identity, qualified)**: `MakeMiddleware(nil, nil)(s) ≈ s` for steps + that invoke their continuation at most once and do no work after invoking + it. (A step that does inline work *after* calling its continuation would + see that work reordered before downstream execution — the price of + delimiting. Steps built from `Do`, `Sequence`, `Decision`, etc. never do + this.) + +These laws are verified in `formal_test.go` (`TestBracketPairingLaw`, +`TestBracketNestingLaw`, `TestBracketTransparencyLaw`, +`TestBracketIdentityLaw`, and the panic sides of B1 in +`TestBracketReleaseOnRecoveredPanic` / +`TestBracketNoReleaseOnUnrecoveredPanic`). + +**On the name "bracket"**: the term is borrowed from effect systems — +Haskell's `Control.Exception.bracket` and the bracket laws of cats-effect's +`Bracket` typeclass — where it names exactly this contract: paired, +guaranteed finalization with LIFO nesting, independent of how the +computation exits. The parallel runs deep: in those systems `bracket` is +famously not definable from `bind` alone, just as `MakeMiddleware` is not +definable from `Sequence`/`Do` (the Absorption Theorem above). It is +unrelated to the Lie or Poisson bracket of algebra. + +A useful side effect of B3: because the sentinel delimits the step, an +after-hook brackets *only its step*, not the entire downstream tail that CPS +would otherwise execute inline within the step's call frame. Bracket-based +timing middleware therefore attributes duration to the step it wraps. + +#### Choosing a Middleware Form + +- **Observability (timers, spans, metrics, logging)**: use `MakeMiddleware`. + You want the finalizer guarantee (B1) and per-step attribution. +- **Pipeline logic that should respect termination** (e.g. a follow-up step + that must *not* run if the pipeline stopped): stay in the algebra — + `func(s NewStep) NewStep { return Sequence(Do(before), s, Do(after)) }`. + Its after-hook is absorbed by terminal steps and skipped on cancellation, + which is exactly the semantics you asked for by writing it as a `Sequence`. +- **Defer semantics around panics**: neither — see `middleware.Recover` in + `state/middleware`. + ## Algebraic Operations These operations give us the algebraic structure to build complex pipelines. They correspond to fundamental categorical constructions. @@ -374,7 +463,7 @@ The `Box` is thread-safe, so parallel branches can safely write to it without ra **Failure Handling in Parallel**: -- **Panics**: Goroutine panics propagate naturally and crash the program (programming errors should be loud). Wrap a step with `Recover(step)` to opt into panic suppression for that specific step. +- **Panics**: Goroutine panics propagate naturally and crash the program (programming errors should be loud). Wrap a step with `middleware.Recover` (in `state/middleware`) to opt into panic suppression for that specific step. - **Context Cancellation**: Checked after all branches complete via `Continue`; cancellation stops the pipeline from advancing to the next step but does not prevent branches from running. - **Errors**: Should be communicated via typedctx.Box or context values, then handled after parallel completes - **Assumption**: All branches run to completion. Context cancellation stops the continuation but not the in-flight goroutines. diff --git a/state/doc.go b/state/doc.go index be5ff44..bf121d6 100644 --- a/state/doc.go +++ b/state/doc.go @@ -135,7 +135,8 @@ // - Decision(predicate, ifTrue, ifFalse) - Binary branching (if/else) // - Enum(selector, cases, default) - Multi-way branching on comparable types // - Switch(selector, cases, default) - Multi-way branching on strings -// - Recover(step) - Wrap a step to suppress panics (opt-in; panics propagate by default) +// +// For panic recovery, use the Recover middleware in the state/middleware sub-package. // // # Choosing a Pattern // diff --git a/state/formal_test.go b/state/formal_test.go index ffebac6..68f57ff 100644 --- a/state/formal_test.go +++ b/state/formal_test.go @@ -482,3 +482,255 @@ func TestCompleteCategoryIntegration(t *testing.T) { // Verify branch-false was not executed require.NotContains(t, trace, "branch-false", "False branch should not execute") } + +// ============================================================================= +// BRACKET MIDDLEWARE LAWS +// +// These tests verify the laws of MakeMiddleware, which is a *bracket* (a +// control operator), not a Kleisli composition. A pure Kleisli middleware +// Sequence(Do(before), step, Do(after)) provably cannot run its after-hook +// around a terminal step: terminal steps are left zeros of Sequence +// (Sequence(t, x) = t), so the after-hook is absorbed. MakeMiddleware +// therefore intercepts the step's continuation with a sentinel and re-emits +// it, trading residency in the Sequence/Do algebra for a finalizer guarantee. +// +// The laws it satisfies instead (see FORMAL.md "Bracket Middleware"): +// B1 (pairing): after runs exactly once whenever before ran, on every +// outcome the program survives (continue, terminal, +// cancel, recovered panic); on an unrecovered panic no +// hooks fire and the panic propagates untouched +// B2 (nesting): mwA(mwB(s)) brackets properly: +// beforeA, beforeB, s, afterB, afterA +// B3 (transparency): terminality is preserved, and downstream observes +// after(ctx') where ctx' is the context the step threaded +// B4 (identity): MakeMiddleware(nil, nil)(s) ≈ s for steps that invoke +// their continuation at most once +// ============================================================================= + +// TestBracketPairingLaw verifies B1: the after-hook fires exactly once +// whenever the before-hook fired, regardless of the wrapped step's outcome. +// +// What this prevents: +// - Spans that never close / timers that never stop on terminal steps, +// which are how every controller pipeline ends (queue.Done, queue.Requeue) +// - Metrics silently dropped on the cancellation (failure) path +func TestBracketPairingLaw(t *testing.T) { + outcomes := map[string]NewStep{ + "continues": Do(func(ctx context.Context) context.Context { return ctx }), + "terminal": Terminal, + "cancels": Do(func(ctx context.Context) context.Context { + c, cancel := context.WithCancel(ctx) + cancel() + return c + }), + } + + for name, step := range outcomes { + t.Run(name, func(t *testing.T) { + var before, after int + mw := MakeMiddleware( + func(ctx context.Context) context.Context { before++; return ctx }, + func(ctx context.Context) context.Context { after++; return ctx }, + ) + Run(t.Context(), mw(step)) + require.Equal(t, 1, before, "before-hook should fire exactly once") + require.Equal(t, before, after, "after-hook must pair with before-hook") + }) + } +} + +// TestBracketNestingLaw verifies B2: nested brackets open and close in +// properly nested (LIFO) order, like defer or try/finally. +// +// What this prevents: +// - Interleaved spans (A opens, B opens, A closes, B closes) that break +// tracing tools expecting well-nested extents +func TestBracketNestingLaw(t *testing.T) { + var trace []string + hook := func(s string) ContextFunc { + return func(ctx context.Context) context.Context { + trace = append(trace, s) + return ctx + } + } + mwA := MakeMiddleware(hook("beforeA"), hook("afterA")) + mwB := MakeMiddleware(hook("beforeB"), hook("afterB")) + + step := Do(func(ctx context.Context) context.Context { + trace = append(trace, "step") + return ctx + }) + + Run(t.Context(), mwA(mwB(step))) + require.Equal(t, []string{"beforeA", "beforeB", "step", "afterB", "afterA"}, trace) +} + +// TestBracketTransparencyLaw verifies B3: the bracket is transparent to the +// pipeline around it. Terminality is preserved (a bracketed terminal step +// still stops the pipeline), and when the step continues, downstream steps +// observe after(ctx') where ctx' is the context the step threaded forward. +// +// What this prevents: +// - A bracket resurrecting a terminated pipeline (running steps after +// queue.Done) +// - Context values written by the step or the after-hook vanishing before +// downstream steps +func TestBracketTransparencyLaw(t *testing.T) { + t.Run("terminality preserved", func(t *testing.T) { + var downstream bool + mw := MakeMiddleware( + func(ctx context.Context) context.Context { return ctx }, + func(ctx context.Context) context.Context { return ctx }, + ) + Run(t.Context(), Sequence( + mw(Terminal), + Do(func(ctx context.Context) context.Context { + downstream = true + return ctx + }), + )) + require.False(t, downstream, "bracketed terminal step must still terminate the pipeline") + }) + + t.Run("context threads through step and after-hook", func(t *testing.T) { + var sawUser string + var sawProcessed bool + mw := MakeMiddleware( + nil, + func(ctx context.Context) context.Context { + // after-hook sees the context the step threaded forward + return processedCtxKey.Set(ctx, true) + }, + ) + Run(t.Context(), Sequence( + mw(Do(func(ctx context.Context) context.Context { + return userCtxKey.Set(ctx, "alice") + })), + Do(func(ctx context.Context) context.Context { + sawUser, _ = userCtxKey.Value(ctx) + sawProcessed, _ = processedCtxKey.Value(ctx) + return ctx + }), + )) + require.Equal(t, "alice", sawUser, "downstream must see context written by the bracketed step") + require.True(t, sawProcessed, "downstream must see context written by the after-hook") + }) +} + +// TestBracketIdentityLaw verifies B4: the empty bracket is observationally +// identity for well-behaved steps (those that invoke their continuation at +// most once and do no work after invoking it). +// +// Known caveat (why this law is qualified): a step that does work inline +// *after* invoking its continuation would observe that work reordered before +// downstream execution, because the bracket delimits the step's dynamic +// extent. Kleisli-composed steps (Do, Sequence, Decision, ...) never do this. +func TestBracketIdentityLaw(t *testing.T) { + identity := MakeMiddleware(nil, nil) + + var trace []string + step := Do(func(ctx context.Context) context.Context { + trace = append(trace, "step") + return userCtxKey.Set(ctx, "alice") + }) + downstream := Do(func(ctx context.Context) context.Context { + user, _ := userCtxKey.Value(ctx) + trace = append(trace, "downstream:"+user) + return ctx + }) + + Run(t.Context(), Sequence(identity(step), downstream)) + require.Equal(t, []string{"step", "downstream:alice"}, trace, + "MakeMiddleware(nil, nil) should be observationally identity") + + // Terminality is also preserved by the empty bracket. + trace = nil + Run(t.Context(), Sequence(identity(Terminal), downstream)) + require.Empty(t, trace, "empty bracket must preserve terminality") +} + +// TestBracketReleaseOnRecoveredPanic verifies the panic side of B1: when a +// panic-recovery scope is active (marked via WithPanicRecovery, as +// middleware.Recover does), the bracket closes during unwinding — the panic +// itself continues to propagate to the recovery middleware untouched. +func TestBracketReleaseOnRecoveredPanic(t *testing.T) { + var trace []string + hook := func(s string) ContextFunc { + return func(ctx context.Context) context.Context { + trace = append(trace, s) + return ctx + } + } + mw := MakeMiddleware(hook("before"), hook("after")) + + // Simulate a recovery middleware: mark the context, recover the panic. + recovery := func(step NewStep) NewStep { + return func(next Step) Step { + return StepFunc(func(ctx context.Context) (result Step) { + defer func() { + if r := recover(); r != nil { + trace = append(trace, "recovered") + } + }() + return step(next).Run(WithPanicRecovery(ctx)) + }) + } + } + + panicStep := NewStepFunc(func(_ context.Context, _ Step) Step { + panic("boom") + }) + + require.NotPanics(t, func() { + Run(t.Context(), recovery(mw(panicStep))) + }) + require.Equal(t, []string{"before", "after", "recovered"}, trace, + "bracket must close during unwinding, before recovery observes the panic") +} + +// TestBracketNoReleaseOnUnrecoveredPanic verifies the other panic side of B1: +// with no recovery scope active, the bracket does not intercept unwinding at +// all — no hooks fire during a fatal panic, and the panic value is untouched. +func TestBracketNoReleaseOnUnrecoveredPanic(t *testing.T) { + var afterRan bool + mw := MakeMiddleware(nil, func(ctx context.Context) context.Context { + afterRan = true + return ctx + }) + panicStep := NewStepFunc(func(_ context.Context, _ Step) Step { + panic("boom") + }) + + require.PanicsWithValue(t, "boom", func() { + Run(t.Context(), mw(panicStep)) + }, "panic must propagate unmodified") + require.False(t, afterRan, + "no recovery scope: hooks must not fire during fatal unwinding") +} + +// TestPanicRecoveryMarkerClearedInParallelBranches verifies that Parallel +// clears the recovery marker for its branches: panic recovery cannot cross +// goroutine boundaries, so a branch must not observe a recovery scope +// established outside the Parallel. The continuation after Parallel (which +// runs inline in the protected frame) keeps the marker. +func TestPanicRecoveryMarkerClearedInParallelBranches(t *testing.T) { + ctx := WithPanicRecovery(t.Context()) + require.True(t, PanicRecoveryActive(ctx)) + + var inBranch, inContinuation bool + Run(ctx, Sequence( + Parallel( + Do(func(c context.Context) context.Context { + inBranch = PanicRecoveryActive(c) + return c + }), + ), + Do(func(c context.Context) context.Context { + inContinuation = PanicRecoveryActive(c) + return c + }), + )) + + require.False(t, inBranch, "recovery cannot cross goroutines; branches must not be marked") + require.True(t, inContinuation, "continuation after Parallel stays in the recovery scope") +} diff --git a/state/middleware/doc.go b/state/middleware/doc.go new file mode 100644 index 0000000..fcf9726 --- /dev/null +++ b/state/middleware/doc.go @@ -0,0 +1,23 @@ +// Package middleware provides ready-made state.Middleware implementations +// for common cross-cutting concerns. +// +// # Available Middleware +// +// - Log: logs each step's name and elapsed time +// - Recover: catches panics and routes them through the error handler +// +// # Usage with AmbientDispatch +// +// The typical pattern is to register middleware via WithAmbientMiddleware and +// run the pipeline with AmbientDispatch: +// +// ctx = state.WithAmbientMiddleware(ctx, middleware.Log(slog.Default())) +// state.Run(ctx, state.AmbientDispatch(step1, step2, step3)) +// +// # Recover +// +// Recover is applied per-step, not ambient. Use it explicitly on steps that +// may panic: +// +// state.Parallel(state.Map(middleware.Recover, step1, step2, step3)...) +package middleware diff --git a/state/middleware/middleware.go b/state/middleware/middleware.go new file mode 100644 index 0000000..76947d8 --- /dev/null +++ b/state/middleware/middleware.go @@ -0,0 +1,75 @@ +// Package middleware provides ready-made state.Middleware implementations +// for common cross-cutting concerns. +package middleware + +import ( + "context" + "log/slog" + "time" + + "github.com/authzed/controller-idioms/state" +) + +// Log returns a Middleware that logs each step's name and elapsed time after +// each step executes. It uses state.Named for the step name (empty string if +// the step was not annotated). Elapsed time is logged as duration_ms +// (float64, milliseconds) at Info level. Works for both terminal and +// non-terminal named steps. +// +// Note: because this wraps a NewStep, the measured duration includes the step +// body and any continuation steps that execute inline via CPS. +// +// Typical controller usage: +// +// ctx = state.WithAmbientMiddleware(ctx, middleware.Log(slog.Default())) +func Log(logger *slog.Logger) state.Middleware { + return func(step state.NewStep) state.NewStep { + return func(next state.Step) state.Step { + return state.StepFunc(func(ctx context.Context) state.Step { + ctxWithBox := state.WithStepNameCapture(ctx) + + start := time.Now() + result := step(next).Run(ctxWithBox) + elapsed := time.Since(start) + + logger.InfoContext(ctx, "step executed", + "step", state.CapturedStepName(ctxWithBox), + "duration_ms", float64(elapsed.Microseconds())/1000.0, + ) + return result + }) + } + } +} + +// Recover is a Middleware that wraps a step with panic recovery. On panic, +// it cancels a child context with a *state.PanicError cause and routes through +// state.Continue, so any state.WithErrorHandler registered on the context will +// be called with the *state.PanicError. The pipeline does not continue past the +// recovered panic. +// +// In most cases, panics should propagate naturally — only use Recover when +// you explicitly need to handle a panic from a specific step. +// +// Recover marks its dynamic extent as a panic-recovery scope (see +// state.WithPanicRecovery), so state.MakeMiddleware brackets inside it close +// their after hooks during unwinding before the panic is recovered here. +// +// Note: Recover cannot catch panics from goroutines spawned by the wrapped +// step (e.g. a Parallel step). Go's runtime does not allow cross-goroutine +// panic recovery; those panics will still crash the program. +func Recover(step state.NewStep) state.NewStep { + return func(next state.Step) state.Step { + return state.StepFunc(func(ctx context.Context) (result state.Step) { + childCtx, cancel := context.WithCancelCause(ctx) + defer cancel(nil) + defer func() { + if r := recover(); r != nil { + cancel(&state.PanicError{Value: r}) + result = state.Continue(childCtx, next) + } + }() + return step(next).Run(state.WithPanicRecovery(childCtx)) + }) + } +} diff --git a/state/middleware/middleware_test.go b/state/middleware/middleware_test.go new file mode 100644 index 0000000..29cb8b3 --- /dev/null +++ b/state/middleware/middleware_test.go @@ -0,0 +1,251 @@ +package middleware_test + +import ( + "bytes" + "context" + "encoding/json" + "log/slog" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/authzed/controller-idioms/state" + "github.com/authzed/controller-idioms/state/middleware" +) + +func TestRecoverStopsPipelineOnPanic(t *testing.T) { + ctx := t.Context() + var afterPanicExecuted bool + + pipeline := state.Sequence( + middleware.Recover(state.Do(func(_ context.Context) context.Context { + panic("intentional panic") + })), + state.Do(func(ctx context.Context) context.Context { + afterPanicExecuted = true + return ctx + }), + ) + + require.NotPanics(t, func() { + state.Run(ctx, pipeline) + }, "Recover should suppress the panic") + require.False(t, afterPanicExecuted, "pipeline should stop after recovered panic") +} + +func TestRecoverContinuesNormally(t *testing.T) { + ctx := t.Context() + var executed []string + + pipeline := state.Sequence( + middleware.Recover(state.Do(func(ctx context.Context) context.Context { + executed = append(executed, "step1") + return ctx + })), + state.Do(func(ctx context.Context) context.Context { + executed = append(executed, "step2") + return ctx + }), + ) + + state.Run(ctx, pipeline) + + require.Equal(t, []string{"step1", "step2"}, executed) +} + +func TestRecoverRoutesToErrorHandler(t *testing.T) { + var handlerErr error + ctx := state.WithErrorHandler(t.Context(), func(err error) state.Step { + handlerErr = err + return nil + }) + + pipeline := middleware.Recover(state.Do(func(_ context.Context) context.Context { + panic("intentional panic") + })) + + require.NotPanics(t, func() { + state.Run(ctx, pipeline) + }) + require.Error(t, handlerErr, "error handler should be called after recovered panic") +} + +func TestRecoverPanicValueAvailableViaCause(t *testing.T) { + type sentinelType struct{ msg string } + panicValue := sentinelType{"something went wrong"} + + var causeErr error + ctx := state.WithErrorHandler(t.Context(), func(err error) state.Step { + causeErr = err + return nil + }) + + pipeline := middleware.Recover(state.Do(func(_ context.Context) context.Context { + panic(panicValue) + })) + + state.Run(ctx, pipeline) + + var p *state.PanicError + require.ErrorAs(t, causeErr, &p, "error should be a *state.PanicError") + require.Equal(t, panicValue, p.Value, "PanicError.Value should be the original panic value") +} + +func TestRecoverWrappingSequence(t *testing.T) { + var innerAfterPanicExecuted bool + var outerAfterPanicExecuted bool + var handlerErr error + + ctx := state.WithErrorHandler(t.Context(), func(err error) state.Step { + handlerErr = err + return nil + }) + + pipeline := state.Sequence( + middleware.Recover(state.Sequence( + state.Do(func(ctx context.Context) context.Context { return ctx }), + state.Do(func(_ context.Context) context.Context { panic("mid-sequence panic") }), + state.Do(func(ctx context.Context) context.Context { + innerAfterPanicExecuted = true + return ctx + }), + )), + state.Do(func(ctx context.Context) context.Context { + outerAfterPanicExecuted = true + return ctx + }), + ) + + require.NotPanics(t, func() { state.Run(ctx, pipeline) }) + require.False(t, innerAfterPanicExecuted, "step inside recovered sequence after panic should not execute") + require.False(t, outerAfterPanicExecuted, "step outside recovered sequence after panic should not execute") + var p *state.PanicError + require.ErrorAs(t, handlerErr, &p) + require.Equal(t, "mid-sequence panic", p.Value) +} + +func TestMapWithRecoverRunsAllBranches(t *testing.T) { + // Verify that wrapping with Recover does not suppress branch execution. + ctx := t.Context() + var counter int32 + + pipeline := state.Parallel(state.Map(middleware.Recover, + state.Do(func(ctx context.Context) context.Context { + atomic.AddInt32(&counter, 1) + return ctx + }), + state.Do(func(ctx context.Context) context.Context { + atomic.AddInt32(&counter, 1) + return ctx + }), + )...) + + state.Run(ctx, pipeline) + require.Equal(t, int32(2), atomic.LoadInt32(&counter)) +} + +func TestMapWithRecoverRoutesPanicToErrorHandler(t *testing.T) { + var handlerErr error + ctx := state.WithErrorHandler(t.Context(), func(err error) state.Step { + handlerErr = err + return nil + }) + + pipeline := state.Parallel(state.Map(middleware.Recover, + state.Do(func(ctx context.Context) context.Context { return ctx }), + state.Do(func(_ context.Context) context.Context { panic("branch panic") }), + state.Do(func(ctx context.Context) context.Context { return ctx }), + )...) + + require.NotPanics(t, func() { state.Run(ctx, pipeline) }) + + var p *state.PanicError + require.ErrorAs(t, handlerErr, &p, "error handler should receive a *state.PanicError from the panicking branch") + require.Equal(t, "branch panic", p.Value) +} + +// Compile-time check: Recover must be assignable to state.Middleware. +var _ state.Middleware = middleware.Recover + +func TestLogRecordsStepNameAndDuration(t *testing.T) { + var buf bytes.Buffer + logger := slog.New(slog.NewJSONHandler(&buf, &slog.HandlerOptions{Level: slog.LevelInfo})) + + ctx := state.WithAmbientMiddleware(t.Context(), middleware.Log(logger)) + state.Run(ctx, state.AmbientDispatch( + state.Named("myStep", state.Do(func(ctx context.Context) context.Context { return ctx })), + )) + + var entry map[string]any + require.NoError(t, json.Unmarshal(buf.Bytes(), &entry)) + require.Equal(t, "myStep", entry["step"]) + _, hasDuration := entry["duration_ms"] + require.True(t, hasDuration, "log entry should contain duration_ms") +} + +func TestLogUnnamedStepLogsEmptyName(t *testing.T) { + var buf bytes.Buffer + logger := slog.New(slog.NewJSONHandler(&buf, &slog.HandlerOptions{Level: slog.LevelInfo})) + + ctx := state.WithAmbientMiddleware(t.Context(), middleware.Log(logger)) + state.Run(ctx, state.AmbientDispatch( + state.Do(func(ctx context.Context) context.Context { return ctx }), + )) + + var entry map[string]any + require.NoError(t, json.Unmarshal(buf.Bytes(), &entry)) + name, ok := entry["step"] + require.True(t, ok, "log entry should contain step") + require.Empty(t, name) +} + +func TestLogRecordsNameForTerminalStep(t *testing.T) { + // Named terminal steps (those that never call their continuation) must still + // log the step name correctly. + var buf bytes.Buffer + logger := slog.New(slog.NewJSONHandler(&buf, &slog.HandlerOptions{Level: slog.LevelInfo})) + + ctx := state.WithAmbientMiddleware(t.Context(), middleware.Log(logger)) + state.Run(ctx, state.AmbientDispatch( + state.Named("cleanup", state.Terminal), + )) + + var entry map[string]any + require.NoError(t, json.Unmarshal(buf.Bytes(), &entry)) + require.Equal(t, "cleanup", entry["step"]) + _, hasDuration := entry["duration_ms"] + require.True(t, hasDuration, "log entry should contain duration_ms") +} + +func TestRecoverEnablesBracketRelease(t *testing.T) { + // middleware.Recover marks its dynamic extent as a panic-recovery scope, + // so MakeMiddleware brackets inside it close during unwinding: spans and + // timers are released before the error handler observes the panic. + var trace []string + hook := func(s string) state.ContextFunc { + return func(ctx context.Context) context.Context { + trace = append(trace, s) + return ctx + } + } + var handlerErr error + ctx := state.WithErrorHandler(t.Context(), func(err error) state.Step { + trace = append(trace, "handler") + handlerErr = err + return nil + }) + + mw := state.MakeMiddleware(hook("before"), hook("after")) + panicStep := state.NewStepFunc(func(_ context.Context, _ state.Step) state.Step { + panic("bracketed panic") + }) + + require.NotPanics(t, func() { + state.Run(ctx, middleware.Recover(mw(panicStep))) + }) + require.Equal(t, []string{"before", "after", "handler"}, trace) + var p *state.PanicError + require.ErrorAs(t, handlerErr, &p) + require.Equal(t, "bracketed panic", p.Value) +} diff --git a/state/state.go b/state/state.go index b517024..fcc56d9 100644 --- a/state/state.go +++ b/state/state.go @@ -40,6 +40,9 @@ import ( "fmt" "slices" "sync" + "sync/atomic" + + "github.com/authzed/ctxkey" ) // Step represents a step in a processing pipeline. @@ -64,6 +67,229 @@ type ContextFunc func(context.Context) context.Context // This is the basic building block for composing step pipelines using continuation-passing style. type NewStep func(next Step) Step +// Middleware wraps a NewStep with additional behavior such as logging, metrics, +// or error handling. It follows the same decorator pattern as HTTP middleware: +// it receives the next step and returns a new step that can execute before, +// after, or around the original. +// +// Example: +// +// var loggingMiddleware Middleware = func(step NewStep) NewStep { +// return func(next Step) Step { +// return StepFunc(func(ctx context.Context) Step { +// log.Println("before") +// result := step(next).Run(ctx) +// log.Println("after") +// return result +// }) +// } +// } +// +// middleware.Recover (in state/middleware) is an example of a Middleware. +type Middleware func(NewStep) NewStep + +var ambientMiddlewareKey = ctxkey.New[Middleware]() + +// AmbientMiddleware returns the composed middleware from context. +// Returns nil if no middleware has been registered. +func AmbientMiddleware(ctx context.Context) Middleware { + mw, _ := ambientMiddlewareKey.Value(ctx) + return mw +} + +// WithAmbientMiddleware composes mw into the ambient middleware stack and +// returns an updated context. The first-registered middleware is outermost. +// Passing nil is a no-op. +func WithAmbientMiddleware(ctx context.Context, mw Middleware) context.Context { + if mw == nil { + return ctx + } + if existing := AmbientMiddleware(ctx); existing != nil { + outer := existing + inner := mw + mw = func(step NewStep) NewStep { return outer(inner(step)) } + } + return ambientMiddlewareKey.Set(ctx, mw) +} + +// AmbientDispatch wraps each step so that ambient middleware registered in +// context fires around every step as the pipeline executes. It is the opt-in +// mechanism for ambient middleware support: +// +// state.Run(ctx, state.AmbientDispatch(a, b, c)) +// +// Each step argument is wrapped individually — middleware fires once per step, +// not once for the whole group. Middleware registered via WithAmbientMiddleware +// before Run applies to all steps. Middleware registered mid-pipeline (inside a +// step) applies to all subsequent steps in the same AmbientDispatch call. +// +// AmbientDispatch works for any NewStep, including raw StepFunc closures and +// struct method steps — no special constructor is required. +func AmbientDispatch(steps ...NewStep) NewStep { + return func(outerNext Step) Step { + // Build the chain right-to-left. Each step's "next" is the already-dispatch- + // wrapped Step for the following step, so middleware fires exactly once per step. + current := outerNext + for _, step := range slices.Backward(steps) { + s := step + n := current + current = StepFunc(func(ctx context.Context) Step { + mw := AmbientMiddleware(ctx) + if mw == nil { + return s(n).Run(ctx) + } + return mw(s)(n).Run(ctx) + }) + } + return current + } +} + +var panicRecoveryKey = ctxkey.NewWithDefault[bool](false) + +// WithPanicRecovery marks ctx as running under middleware that recovers +// panics (middleware.Recover does this automatically). Within a marked scope, +// brackets built by MakeMiddleware run their after hooks during panic +// unwinding — without calling recover, so the panic continues to the recovery +// middleware untouched. Outside a marked scope, a panicking step skips all +// hooks and the panic propagates unmodified. +// +// Custom panic-recovering middleware should set this on the context it passes +// to wrapped steps. Integrators whose framework recovers panics outside the +// pipeline entirely can set it before Run. +func WithPanicRecovery(ctx context.Context) context.Context { + return panicRecoveryKey.Set(ctx, true) +} + +// PanicRecoveryActive reports whether ctx is within a panic-recovery scope +// established by WithPanicRecovery. Parallel clears the marker for branch +// goroutines, since panic recovery cannot cross goroutine boundaries. +func PanicRecoveryActive(ctx context.Context) bool { + return panicRecoveryKey.Value(ctx) +} + +// MakeMiddleware builds a Middleware from before and after ContextFunc hooks. +// Either hook may be nil. before runs before the wrapped step; after runs once +// the step completes — including when the step is terminal or cancels the +// context, neither of which invokes the continuation. This guarantees +// after-hooks (stop a timer, end a span, record a metric) always fire. +// +// On panic, the after hook runs only within a panic-recovery scope (see +// WithPanicRecovery): if recovery middleware is going to keep the program +// alive, brackets release during unwinding; otherwise the panic is fatal and +// propagates with no hooks fired and its value untouched. +// +// after runs for its side effects: the context it returns is threaded to +// subsequent steps only when the wrapped step continued the pipeline. +// +// Use MakeMiddleware for logging, metrics, tracing, and similar concerns. +// For middleware requiring defer semantics (e.g. panic recovery), write a +// Middleware function directly (see middleware.Recover in state/middleware) +// or use a custom StepFunc. +// +// MakeMiddleware is a bracket, not a Kleisli composition: a pure +// Sequence(Do(before), step, Do(after)) provably drops the after hook on +// terminal steps. The laws it satisfies instead (pairing, nesting, +// transparency, qualified identity) are stated in FORMAL.md under +// "Bracket Middleware" and verified in formal_test.go. +func MakeMiddleware(before, after ContextFunc) Middleware { + return func(step NewStep) NewStep { + return func(next Step) Step { + return StepFunc(func(ctx context.Context) Step { + if before != nil { + ctx = before(ctx) + } + var continued, afterRan bool + fwd := ctx + runAfter := func() { + if afterRan { + return + } + afterRan = true + if after != nil { + fwd = after(fwd) + } + } + if PanicRecoveryActive(ctx) { + // Recovery middleware above will keep the program alive, + // so close the bracket during panic unwinding. recover is + // not called here — the panic continues to propagate. On + // the normal path the flag makes this call a no-op. + defer runAfter() + } + // Intercept the step's continuation with a sentinel so we + // regain control after the step finishes even if it is + // terminal or cancels the context. The sentinel records + // whether the step continued and the context it threaded + // forward. + sentinel := StepFunc(func(c context.Context) Step { + continued, fwd = true, c + return nil + }) + step(sentinel).Run(ctx) + runAfter() + if continued { + return Continue(fwd, next) + } + return nil + }) + } + } +} + +var ( + stepNameKey = ctxkey.NewWithDefault[string]("") + stepNameCaptureKey = ctxkey.New[*atomic.Pointer[string]]() +) + +// StepName returns the name of the currently-executing step from context. +// Returns the name set by Named if present, otherwise "". +func StepName(ctx context.Context) string { + return stepNameKey.Value(ctx) +} + +// WithStepNameCapture allocates a capture slot in context that the first +// (outermost) Named step fills with its name. This lets observability +// middleware read the step name even for terminal steps that never invoke +// their continuation. The slot is a pointer to an atomic, so concurrent Named +// branches (e.g. under Parallel) fill it without racing and the first writer +// wins. Pair with CapturedStepName to read the result. +func WithStepNameCapture(ctx context.Context) context.Context { + return stepNameCaptureKey.Set(ctx, &atomic.Pointer[string]{}) +} + +// CapturedStepName returns the step name written into the capture slot set up +// by WithStepNameCapture. Returns "" if no slot was allocated or no Named step +// ran. +func CapturedStepName(ctx context.Context) string { + if slot, ok := stepNameCaptureKey.Value(ctx); ok && slot != nil { + if name := slot.Load(); name != nil { + return *name + } + } + return "" +} + +// Named annotates a step with a human-readable name for observability. +// The name is stored in context before the step executes, making it +// available to middleware via StepName(ctx). +// +// Named is entirely optional — pipelines behave identically without it. +func Named(name string, step NewStep) NewStep { + return func(next Step) Step { + return StepFunc(func(ctx context.Context) Step { + ctx = stepNameKey.Set(ctx, name) + // Fill the capture slot (if a middleware allocated one). First + // writer wins, so an outer Named is recorded over any nested ones, + // and concurrent Named branches do not race. + if slot, ok := stepNameCaptureKey.Value(ctx); ok && slot != nil { + slot.CompareAndSwap(nil, &name) + } + return step(next).Run(ctx) + }) + } +} + // Step converts a NewStep to a Step by calling it with nil as the next step. func (ns NewStep) Step() Step { return ns(nil) @@ -78,17 +304,6 @@ func (ns NewStep) Step() Step { // return Continue(ctx, next) // }) // } -// -// Instead of the more verbose: -// -// func MyStep() NewStep { -// return func(next Step) Step { -// return StepFunc(func(ctx context.Context) Step { -// // your logic here -// return Continue(ctx, next) -// }) -// } -// } func NewStepFunc(fn func(ctx context.Context, next Step) Step) NewStep { return func(next Step) Step { return StepFunc(func(ctx context.Context) Step { @@ -122,13 +337,13 @@ func Run(ctx context.Context, newStep NewStep) { } } -type errorHandlerKey struct{} +var errorHandlerKey = ctxkey.New[func(error) Step]() // WithErrorHandler adds an error handler to the context. // When Continue encounters a cancelled context, it will call this handler // instead of stopping the pipeline. func WithErrorHandler(ctx context.Context, handler func(error) Step) context.Context { - return context.WithValue(ctx, errorHandlerKey{}, handler) + return errorHandlerKey.Set(ctx, handler) } // Continue runs the next step if it exists and context is not cancelled. @@ -148,7 +363,7 @@ func WithErrorHandler(ctx context.Context, handler func(error) Step) context.Con func Continue(ctx context.Context, next Step) Step { if ctx.Err() != nil { err := context.Cause(ctx) - if handler, ok := ctx.Value(errorHandlerKey{}).(func(error) Step); ok && handler != nil { + if handler, ok := errorHandlerKey.Value(ctx); ok && handler != nil { return handler(err) } return nil @@ -193,24 +408,19 @@ func Sequence(steps ...NewStep) NewStep { } } -// ParallelWith composes multiple NewStep functions to run in parallel, -// applying wrapper to each branch before execution. This is useful for -// applying a uniform policy to all branches, such as panic recovery: +// Map applies wrapper to each step and returns the resulting slice. +// This is useful for applying a uniform policy to a set of steps before +// passing them to Parallel or other combinators: // -// ParallelWith(Recover, step1, step2, step3) +// Parallel(Map(Recover, step1, step2, step3)...) // -// The wrapper is called once per branch at instantiation time (when the -// returned NewStep is called), not at execution time. -// -// If using Recover as the wrapper and multiple branches panic concurrently, -// the error handler may be called multiple times — once per panicking branch. -// The context cause will reflect whichever panic cancelled the context first. -func ParallelWith(wrapper func(NewStep) NewStep, steps ...NewStep) NewStep { - wrapped := make([]NewStep, len(steps)) +// The wrapper is called once per step when Map is called. +func Map(wrapper Middleware, steps ...NewStep) []NewStep { + result := make([]NewStep, len(steps)) for i, s := range steps { - wrapped[i] = wrapper(s) + result[i] = wrapper(s) } - return Parallel(wrapped...) + return result } // Parallel composes multiple NewStep functions to run in parallel, @@ -233,10 +443,20 @@ type ParallelStep struct { func (p *ParallelStep) Run(ctx context.Context) Step { var wg sync.WaitGroup + // Panic recovery cannot cross goroutine boundaries, so branches must not + // observe a recovery scope established outside the Parallel. A branch + // with its own recovery middleware re-establishes the marker itself. The + // continuation after Parallel runs inline in the enclosing frame and + // keeps the original context (and marker). + branchCtx := ctx + if PanicRecoveryActive(ctx) { + branchCtx = panicRecoveryKey.Set(ctx, false) + } + for _, step := range p.steps { wg.Go(func() { if s := step.Step(); s != nil { - s.Run(ctx) + s.Run(branchCtx) } }) } @@ -246,9 +466,10 @@ func (p *ParallelStep) Run(ctx context.Context) Step { } // PanicError wraps a recovered panic value as an error. -// It is set as the cause on the context passed to WithErrorHandler when Recover -// catches a panic, so callers can distinguish panics from normal cancellations -// and recover the original panic value via errors.As. +// It is set as the cause on the context passed to WithErrorHandler when +// middleware.Recover (in state/middleware) catches a panic, so callers can +// distinguish panics from normal cancellations and recover the original panic +// value via errors.As. type PanicError struct { Value any } @@ -257,32 +478,6 @@ func (p *PanicError) Error() string { return fmt.Sprintf("panic: %v", p.Value) } -// Recover wraps a step with panic recovery. On panic, it cancels a child -// context with a *PanicError cause and routes through Continue, so any -// WithErrorHandler registered on the context will be called with the -// *PanicError. The pipeline does not continue past the recovered panic. -// In most cases, panics should propagate naturally — only use Recover when -// you explicitly need to handle a panic from a specific step. -// -// Note: Recover cannot catch panics from goroutines spawned by the wrapped -// step (e.g. a Parallel step). Go's runtime does not allow cross-goroutine -// panic recovery; those panics will still crash the program. -func Recover(step NewStep) NewStep { - return func(next Step) Step { - return StepFunc(func(ctx context.Context) (result Step) { - childCtx, cancel := context.WithCancelCause(ctx) - defer cancel(nil) - defer func() { - if r := recover(); r != nil { - cancel(&PanicError{Value: r}) - result = Continue(childCtx, next) - } - }() - return step(next).Run(childCtx) - }) - } -} - // Decision creates a conditional step that chooses between two paths. func Decision(predicate func(context.Context) bool, trueHandler, falseHandler NewStep) NewStep { return func(next Step) Step { diff --git a/state/state_test.go b/state/state_test.go index 5b32e36..7bc77a4 100644 --- a/state/state_test.go +++ b/state/state_test.go @@ -990,6 +990,74 @@ func TestEnumWithoutDefault(t *testing.T) { } } +func TestWhenPredicateTrue(t *testing.T) { + ctx := t.Context() + var handlerExecuted, afterExecuted bool + + pipeline := Sequence( + When( + func(_ context.Context) bool { return true }, + Do(func(ctx context.Context) context.Context { + handlerExecuted = true + return ctx + }), + ), + Do(func(ctx context.Context) context.Context { + afterExecuted = true + return ctx + }), + ) + + Run(ctx, pipeline) + + require.True(t, handlerExecuted) + require.True(t, afterExecuted, "pipeline continues after a non-terminal When handler") +} + +func TestWhenPredicateFalse(t *testing.T) { + ctx := t.Context() + var handlerExecuted, afterExecuted bool + + pipeline := Sequence( + When( + func(_ context.Context) bool { return false }, + Do(func(ctx context.Context) context.Context { + handlerExecuted = true + return ctx + }), + ), + Do(func(ctx context.Context) context.Context { + afterExecuted = true + return ctx + }), + ) + + Run(ctx, pipeline) + + require.False(t, handlerExecuted) + require.True(t, afterExecuted, "pipeline continues when predicate is false") +} + +func TestWhenWithTerminalHandlerStopsPipeline(t *testing.T) { + ctx := t.Context() + var afterExecuted bool + + pipeline := Sequence( + When( + func(_ context.Context) bool { return true }, + Terminal, + ), + Do(func(ctx context.Context) context.Context { + afterExecuted = true + return ctx + }), + ) + + Run(ctx, pipeline) + + require.False(t, afterExecuted, "terminal handler prevents continuation to subsequent steps") +} + func TestDecisionRespectsCancellation(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() @@ -1080,7 +1148,7 @@ func TestParallelDoesNotRecoverPanics(t *testing.T) { // // Verified by inspection: ParallelStep.Run spawns goroutines with no recover, // and Go's runtime does not allow cross-goroutine panic recovery. - t.Log("Parallel panics cannot be caught by Recover (cross-goroutine limitation, verified by inspection)") + t.Log("Parallel panics cannot be caught by middleware.Recover (cross-goroutine limitation, verified by inspection)") } func TestParallelRespectsErrorHandler(t *testing.T) { @@ -1102,174 +1170,26 @@ func TestParallelRespectsErrorHandler(t *testing.T) { require.True(t, handlerCalled, "Parallel should invoke WithErrorHandler on cancellation, not bypass it") } -func TestRecoverStopsPipelineOnPanic(t *testing.T) { - ctx := t.Context() - var afterPanicExecuted bool - - pipeline := Sequence( - Recover(Do(func(_ context.Context) context.Context { - panic("intentional panic") - })), - Do(func(ctx context.Context) context.Context { - afterPanicExecuted = true - return ctx - }), - ) - - require.NotPanics(t, func() { - Run(ctx, pipeline) - }, "Recover should suppress the panic") - require.False(t, afterPanicExecuted, "pipeline should stop after recovered panic") -} - -func TestRecoverContinuesNormally(t *testing.T) { - ctx := t.Context() - var executed []string - - pipeline := Sequence( - Recover(Do(func(ctx context.Context) context.Context { - executed = append(executed, "step1") - return ctx - })), - Do(func(ctx context.Context) context.Context { - executed = append(executed, "step2") - return ctx - }), - ) - - Run(ctx, pipeline) - - require.Equal(t, []string{"step1", "step2"}, executed) -} - -func TestRecoverRoutesToErrorHandler(t *testing.T) { - var handlerErr error - ctx := WithErrorHandler(t.Context(), func(err error) Step { - handlerErr = err - return nil - }) - - pipeline := Recover(Do(func(_ context.Context) context.Context { - panic("intentional panic") - })) - - require.NotPanics(t, func() { - Run(ctx, pipeline) - }) - require.Error(t, handlerErr, "error handler should be called after recovered panic") -} - -func TestRecoverPanicValueAvailableViaCause(t *testing.T) { - type sentinelType struct{ msg string } - panicValue := sentinelType{"something went wrong"} - - var causeErr error - ctx := WithErrorHandler(t.Context(), func(err error) Step { - causeErr = err - return nil - }) - - pipeline := Recover(Do(func(_ context.Context) context.Context { - panic(panicValue) - })) - - Run(ctx, pipeline) - - var p *PanicError - require.ErrorAs(t, causeErr, &p, "error should be a *PanicError") - require.Equal(t, panicValue, p.Value, "PanicError.Value should be the original panic value") -} - -func TestRecoverPanicErrorMessage(t *testing.T) { +func TestPanicErrorMessage(t *testing.T) { p := &PanicError{Value: "boom"} require.Equal(t, "panic: boom", p.Error()) } -func TestRecoverWrappingSequence(t *testing.T) { - var afterPanicExecuted bool - var handlerErr error - - ctx := WithErrorHandler(t.Context(), func(err error) Step { - handlerErr = err - return nil - }) - - pipeline := Sequence( - Recover(Sequence( - Do(func(ctx context.Context) context.Context { return ctx }), - Do(func(_ context.Context) context.Context { panic("mid-sequence panic") }), - Do(func(ctx context.Context) context.Context { - afterPanicExecuted = true - return ctx - }), - )), - Do(func(ctx context.Context) context.Context { - afterPanicExecuted = true - return ctx - }), - ) - - require.NotPanics(t, func() { Run(ctx, pipeline) }) - require.False(t, afterPanicExecuted, "steps after panic should not execute") - var p *PanicError - require.ErrorAs(t, handlerErr, &p) - require.Equal(t, "mid-sequence panic", p.Value) -} - -func TestParallelWithRunsAllBranches(t *testing.T) { - ctx := t.Context() - var counter int32 - - pipeline := ParallelWith(Recover, - Do(func(ctx context.Context) context.Context { - atomic.AddInt32(&counter, 1) - return ctx - }), - Do(func(ctx context.Context) context.Context { - atomic.AddInt32(&counter, 1) - return ctx - }), - ) - - Run(ctx, pipeline) - require.Equal(t, int32(2), atomic.LoadInt32(&counter)) -} - -func TestParallelWithRecoverRoutesPanicToErrorHandler(t *testing.T) { - var handlerErr error - ctx := WithErrorHandler(t.Context(), func(err error) Step { - handlerErr = err - return nil - }) - - pipeline := ParallelWith(Recover, - Do(func(ctx context.Context) context.Context { return ctx }), - Do(func(_ context.Context) context.Context { panic("branch panic") }), - Do(func(ctx context.Context) context.Context { return ctx }), - ) - - require.NotPanics(t, func() { Run(ctx, pipeline) }) - - var p *PanicError - require.ErrorAs(t, handlerErr, &p, "error handler should receive a *PanicError from the panicking branch") - require.Equal(t, "branch panic", p.Value) -} - -func TestParallelWithCustomWrapper(t *testing.T) { - // Verify that the wrapper is actually applied to each branch, not just once. +func TestMapAppliesWrapperToEach(t *testing.T) { + // Map applies the wrapper to each step. var wrappedCount int32 countingWrapper := func(step NewStep) NewStep { atomic.AddInt32(&wrappedCount, 1) return step } - ParallelWith(countingWrapper, + Map(countingWrapper, Do(func(ctx context.Context) context.Context { return ctx }), Do(func(ctx context.Context) context.Context { return ctx }), Do(func(ctx context.Context) context.Context { return ctx }), - )(nil) // instantiate to trigger wrapping + ) - require.Equal(t, int32(3), atomic.LoadInt32(&wrappedCount), "wrapper should be applied to each branch") + require.Equal(t, int32(3), atomic.LoadInt32(&wrappedCount), "wrapper should be applied to each step") } func TestParallelCancellation(t *testing.T) { @@ -2135,6 +2055,131 @@ func TestIntegrationContextThreadingShowcase(t *testing.T) { } } +// ============================================================================ +// AMBIENT DISPATCH TESTS +// ============================================================================ + +func TestAmbientDispatchNoMiddleware(t *testing.T) { + // With no middleware registered, AmbientDispatch is transparent + var executed bool + pipeline := AmbientDispatch(Do(func(ctx context.Context) context.Context { + executed = true + return ctx + })) + Run(t.Context(), pipeline) + require.True(t, executed) +} + +func TestAmbientDispatchAppliesMiddleware(t *testing.T) { + var log []string + mw := MakeMiddleware( + func(ctx context.Context) context.Context { log = append(log, "before"); return ctx }, + func(ctx context.Context) context.Context { log = append(log, "after"); return ctx }, + ) + ctx := WithAmbientMiddleware(t.Context(), mw) + + pipeline := AmbientDispatch(Do(func(ctx context.Context) context.Context { + log = append(log, "step") + return ctx + })) + Run(ctx, pipeline) + require.Equal(t, []string{"before", "step", "after"}, log) +} + +func TestAmbientDispatchPropagatesAcrossSteps(t *testing.T) { + // Middleware registered before Run fires for every step in the pipeline + var log []string + mw := MakeMiddleware( + func(ctx context.Context) context.Context { log = append(log, "before"); return ctx }, + func(ctx context.Context) context.Context { log = append(log, "after"); return ctx }, + ) + ctx := WithAmbientMiddleware(t.Context(), mw) + + pipeline := AmbientDispatch( + Do(func(ctx context.Context) context.Context { log = append(log, "step1"); return ctx }), + Do(func(ctx context.Context) context.Context { log = append(log, "step2"); return ctx }), + ) + Run(ctx, pipeline) + require.Equal(t, []string{ + "before", "step1", "after", + "before", "step2", "after", + }, log) +} + +func TestAmbientDispatchMidPipelineInjection(t *testing.T) { + // Middleware registered inside a step applies to subsequent steps only + var log []string + mw := MakeMiddleware( + func(ctx context.Context) context.Context { log = append(log, "mw-before"); return ctx }, + func(ctx context.Context) context.Context { log = append(log, "mw-after"); return ctx }, + ) + + pipeline := AmbientDispatch( + Do(func(ctx context.Context) context.Context { + log = append(log, "step1") + return ctx + }), + Do(func(ctx context.Context) context.Context { + log = append(log, "inject") + return WithAmbientMiddleware(ctx, mw) + }), + Do(func(ctx context.Context) context.Context { + log = append(log, "step3") + return ctx + }), + Do(func(ctx context.Context) context.Context { + log = append(log, "step4") + return ctx + }), + ) + Run(t.Context(), pipeline) + require.Equal(t, []string{ + "step1", + "inject", + "mw-before", "step3", "mw-after", + "mw-before", "step4", "mw-after", + }, log) +} + +func TestAmbientDispatchWorksWithStructStep(t *testing.T) { + // AmbientDispatch works for steps that are raw StepFuncs (not NewStepFunc), + // simulating the struct method pattern. + var log []string + mw := MakeMiddleware( + func(ctx context.Context) context.Context { log = append(log, "before"); return ctx }, + func(ctx context.Context) context.Context { log = append(log, "after"); return ctx }, + ) + ctx := WithAmbientMiddleware(t.Context(), mw) + + // Simulate a struct-based step: raw NewStep returning a raw StepFunc + structStep := NewStep(func(next Step) Step { + return StepFunc(func(ctx context.Context) Step { + log = append(log, "struct-step") + return Continue(ctx, next) + }) + }) + + pipeline := AmbientDispatch(structStep) + Run(ctx, pipeline) + require.Equal(t, []string{"before", "struct-step", "after"}, log) +} + +func TestAmbientDispatchNamedStepNameVisibleInStepBody(t *testing.T) { + // Named sets the step name in ctx before the inner step executes. + // Middleware wrapping the outer Named step fires before the name is set, + // so middleware "before" hooks see "" for StepName. The name is visible + // within the step body and to any middleware applied to the inner step. + var nameInBody string + pipeline := AmbientDispatch( + Named("myStep", Do(func(ctx context.Context) context.Context { + nameInBody = StepName(ctx) + return ctx + })), + ) + Run(t.Context(), pipeline) + require.Equal(t, "myStep", nameInBody) +} + // ============================================================================ // BENCHMARKS // ============================================================================ @@ -2302,3 +2347,496 @@ func ExampleWithErrorHandler() { // step 1 // handling cancellation } + +// ============================================================================ +// MIDDLEWARE TESTS +// ============================================================================ + +func TestNewStepFuncIsConvenienceWrapper(t *testing.T) { + // NewStepFunc(fn) should be equivalent to: + // func(next Step) Step { return StepFunc(func(ctx context.Context) Step { return fn(ctx, next) }) } + var called bool + step := NewStepFunc(func(ctx context.Context, next Step) Step { + called = true + return Continue(ctx, next) + }) + Run(t.Context(), step) + require.True(t, called) +} + +func TestMiddlewareIsAssignableFromFunc(t *testing.T) { + // Middleware should be a named type for func(NewStep) NewStep, + // assignable from a plain function literal. + var mw Middleware = func(step NewStep) NewStep { + return step + } + require.NotNil(t, mw) +} + +func TestMiddlewareWrapsStep(t *testing.T) { + // A Middleware should be able to intercept execution. + var called bool + mw := Middleware(func(step NewStep) NewStep { + return func(next Step) Step { + called = true + return step(next) + } + }) + + pipeline := mw(Do(func(ctx context.Context) context.Context { return ctx })) + Run(t.Context(), pipeline) + require.True(t, called) +} + +func TestMapAcceptsMiddleware(t *testing.T) { + // Map should accept a Middleware value directly (not just a raw func). + var wrappedCount int32 + var mw Middleware = func(step NewStep) NewStep { + atomic.AddInt32(&wrappedCount, 1) + return step + } + + Map(mw, + Do(func(ctx context.Context) context.Context { return ctx }), + Do(func(ctx context.Context) context.Context { return ctx }), + ) + + require.Equal(t, int32(2), atomic.LoadInt32(&wrappedCount)) +} + +func TestAmbientMiddlewareNilByDefault(t *testing.T) { + require.Nil(t, AmbientMiddleware(t.Context())) +} + +func TestWithAmbientMiddlewareSingle(t *testing.T) { + var called bool + mw := Middleware(func(step NewStep) NewStep { + called = true + return step + }) + ctx := WithAmbientMiddleware(t.Context(), mw) + got := AmbientMiddleware(ctx) + require.NotNil(t, got) + got(Noop) // trigger it + require.True(t, called) +} + +func TestWithAmbientMiddlewareComposes(t *testing.T) { + // Middleware registered first is outermost at pipeline run time. + // "Outermost" means its before-logic runs first, after-logic runs last. + var log []string + + mw1 := Middleware(func(step NewStep) NewStep { + return func(next Step) Step { + return StepFunc(func(ctx context.Context) Step { + log = append(log, "mw1-before") + result := step(next).Run(ctx) + log = append(log, "mw1-after") + return result + }) + } + }) + mw2 := Middleware(func(step NewStep) NewStep { + return func(next Step) Step { + return StepFunc(func(ctx context.Context) Step { + log = append(log, "mw2-before") + result := step(next).Run(ctx) + log = append(log, "mw2-after") + return result + }) + } + }) + + ctx := WithAmbientMiddleware(t.Context(), mw1) + ctx = WithAmbientMiddleware(ctx, mw2) + + // AmbientDispatch is required to apply ambient middleware to each step. + Run(ctx, AmbientDispatch(Do(func(ctx context.Context) context.Context { + log = append(log, "step") + return ctx + }))) + + // mw1 registered first = outermost: mw1(mw2(step)). + require.Equal(t, []string{"mw1-before", "mw2-before", "step", "mw2-after", "mw1-after"}, log) +} + +func TestWithAmbientMiddlewareComposesDirectly(t *testing.T) { + var log []string + mw1 := Middleware(func(step NewStep) NewStep { + return func(next Step) Step { + return StepFunc(func(ctx context.Context) Step { + log = append(log, "mw1-before") + result := step(next).Run(ctx) + log = append(log, "mw1-after") + return result + }) + } + }) + mw2 := Middleware(func(step NewStep) NewStep { + return func(next Step) Step { + return StepFunc(func(ctx context.Context) Step { + log = append(log, "mw2-before") + result := step(next).Run(ctx) + log = append(log, "mw2-after") + return result + }) + } + }) + + ctx := WithAmbientMiddleware(t.Context(), mw1) + ctx = WithAmbientMiddleware(ctx, mw2) + + // Apply composed middleware directly to a step and run it + composed := AmbientMiddleware(ctx) + require.NotNil(t, composed) + + pipeline := composed(Do(func(ctx context.Context) context.Context { + log = append(log, "step") + return ctx + })) + Run(t.Context(), pipeline) + + // mw1 registered first = outermost + require.Equal(t, []string{"mw1-before", "mw2-before", "step", "mw2-after", "mw1-after"}, log) +} + +func TestWithAmbientMiddlewareNilIsNoOp(t *testing.T) { + // Passing nil should be a no-op, not panic + mw := Middleware(func(step NewStep) NewStep { return step }) + ctx := WithAmbientMiddleware(t.Context(), mw) + ctx2 := WithAmbientMiddleware(ctx, nil) + // nil is no-op — same middleware still present + require.NotNil(t, AmbientMiddleware(ctx2)) +} + +func TestAmbientMiddlewareMidPipelineInjection(t *testing.T) { + // Middleware registered inside a step during execution should + // affect all subsequent steps. + var log []string + + mw := Middleware(func(step NewStep) NewStep { + return func(next Step) Step { + return StepFunc(func(ctx context.Context) Step { + log = append(log, "mw-before") + result := step(next).Run(ctx) + log = append(log, "mw-after") + return result + }) + } + }) + + pipeline := AmbientDispatch( + Do(func(ctx context.Context) context.Context { + log = append(log, "step1") + return ctx + }), + Do(func(ctx context.Context) context.Context { + // Inject mid-pipeline + log = append(log, "inject") + return WithAmbientMiddleware(ctx, mw) + }), + Do(func(ctx context.Context) context.Context { + log = append(log, "step3") + return ctx + }), + Do(func(ctx context.Context) context.Context { + log = append(log, "step4") + return ctx + }), + ) + + Run(t.Context(), pipeline) + // step1 and inject are before middleware registration — no wrapping. + // step3 and step4 get middleware applied. CPS runs continuations inline: + // step4's dispatch is nested inside step3's middleware scope, giving + // continuation-threading ordering where afters fire after all remaining steps. + require.Equal(t, []string{ + "step1", + "inject", + "mw-before", "step3", "mw-before", "step4", "mw-after", "mw-after", + }, log) +} + +func TestMakeMiddlewareFiresBeforeAndAfter(t *testing.T) { + var log []string + + mw := MakeMiddleware( + func(ctx context.Context) context.Context { + log = append(log, "before") + return ctx + }, + func(ctx context.Context) context.Context { + log = append(log, "after") + return ctx + }, + ) + + pipeline := mw(Do(func(ctx context.Context) context.Context { + log = append(log, "step") + return ctx + })) + + Run(t.Context(), pipeline) + require.Equal(t, []string{"before", "step", "after"}, log) +} + +func TestMakeMiddlewareNilBeforeOrAfter(t *testing.T) { + // nil before or after should not panic + mw := MakeMiddleware(nil, nil) + pipeline := mw(Do(func(ctx context.Context) context.Context { return ctx })) + require.NotPanics(t, func() { Run(t.Context(), pipeline) }) +} + +func TestMakeMiddlewareWrapsSequenceAsUnit(t *testing.T) { + // mw(Sequence(a, b)) wraps the whole sequence with one before/after pair. + // This differs from Sequence(mw(a), mw(b)) which wraps each step individually. + // MakeMiddleware treats its step argument as a single unit. + var log []string + mw := MakeMiddleware( + func(ctx context.Context) context.Context { log = append(log, "before"); return ctx }, + func(ctx context.Context) context.Context { log = append(log, "after"); return ctx }, + ) + a := Do(func(ctx context.Context) context.Context { log = append(log, "a"); return ctx }) + b := Do(func(ctx context.Context) context.Context { log = append(log, "b"); return ctx }) + + Run(t.Context(), mw(Sequence(a, b))) + require.Equal(t, []string{"before", "a", "b", "after"}, log) +} + +func TestAmbientMiddlewareFiresAroundSubsequentSteps(t *testing.T) { + // Middleware injected at step N fires before/after steps N+1, N+2, etc. + var log []string + + loggingMW := MakeMiddleware( + func(ctx context.Context) context.Context { + log = append(log, "before") + return ctx + }, + func(ctx context.Context) context.Context { + log = append(log, "after") + return ctx + }, + ) + + pipeline := AmbientDispatch( + Do(func(ctx context.Context) context.Context { + return WithAmbientMiddleware(ctx, loggingMW) + }), + Do(func(ctx context.Context) context.Context { + log = append(log, "step2") + return ctx + }), + Do(func(ctx context.Context) context.Context { + log = append(log, "step3") + return ctx + }), + ) + + Run(t.Context(), pipeline) + require.Equal(t, []string{"before", "step2", "after", "before", "step3", "after"}, log) +} + +func TestAmbientMiddlewareNotAppliedBeforeInjection(t *testing.T) { + // Steps before the injection point are NOT wrapped. + var log []string + + loggingMW := MakeMiddleware( + func(ctx context.Context) context.Context { + log = append(log, "mw") + return ctx + }, + nil, + ) + + pipeline := AmbientDispatch( + Do(func(ctx context.Context) context.Context { + log = append(log, "step1") + return ctx + }), + Do(func(ctx context.Context) context.Context { + return WithAmbientMiddleware(ctx, loggingMW) + }), + Do(func(ctx context.Context) context.Context { + log = append(log, "step3") + return ctx + }), + ) + + Run(t.Context(), pipeline) + // mw should appear only around step3, not step1 + require.Equal(t, []string{"step1", "mw", "step3"}, log) +} + +func TestNamedStepNameAvailableInStepBody(t *testing.T) { + // Named sets the step name in ctx before the inner step body executes. + // The name is visible within the step body via StepName(ctx). + // Middleware wrapping the outer Named step fires before the name is set, + // so outer middleware "before" hooks see "". Use Named to annotate steps + // for structured logging within the step body, not for middleware observability. + var nameInBody string + + pipeline := AmbientDispatch( + Do(func(ctx context.Context) context.Context { + return WithAmbientMiddleware(ctx, Middleware(func(step NewStep) NewStep { + return func(next Step) Step { + return StepFunc(func(ctx context.Context) Step { + return step(next).Run(ctx) + }) + } + })) + }), + Named("myStep", Do(func(ctx context.Context) context.Context { + nameInBody = StepName(ctx) + return ctx + })), + ) + + Run(t.Context(), pipeline) + require.Equal(t, "myStep", nameInBody) +} + +func TestUnnamedStepHasEmptyName(t *testing.T) { + // Unnamed steps have StepName == "" — no reflection-based fallback. + // Use Named() to provide an explicit name for observability. + var observedName string + + namingMW := Middleware(func(step NewStep) NewStep { + return func(next Step) Step { + return StepFunc(func(ctx context.Context) Step { + observedName = StepName(ctx) + return step(next).Run(ctx) + }) + } + }) + + ctx := WithAmbientMiddleware(t.Context(), namingMW) + Run(ctx, AmbientDispatch(Do(func(ctx context.Context) context.Context { return ctx }))) + require.Empty(t, observedName) +} + +func TestNamedPipelineWorksWithoutMiddleware(t *testing.T) { + // Named should be transparent — pipeline works identically with or without middleware + var executed bool + pipeline := Named("myStep", Do(func(ctx context.Context) context.Context { + executed = true + return ctx + })) + Run(t.Context(), pipeline) + require.True(t, executed) +} + +func TestNamedSetsNameInContext(t *testing.T) { + var observedName string + step := Named("myStep", NewStepFunc(func(ctx context.Context, next Step) Step { + observedName = StepName(ctx) + return Continue(ctx, next) + })) + Run(t.Context(), step) + require.Equal(t, "myStep", observedName) +} + +func TestNamedTransparentWithoutMiddleware(t *testing.T) { + var executed bool + pipeline := Named("myStep", Do(func(ctx context.Context) context.Context { + executed = true + return ctx + })) + Run(t.Context(), pipeline) + require.True(t, executed) +} + +// ============================================================================ +// MIDDLEWARE AMBIENT INTEGRATION TESTS +// ============================================================================ + +func TestMiddlewareAmbient(t *testing.T) { + // Middleware registered as ambient applies to each subsequent step. + // Since CPS runs continuations inline, step3's dispatch is nested inside + // step2's middleware scope — afters fire after the entire remaining pipeline. + var log []string + + mw := Middleware(func(step NewStep) NewStep { + return func(next Step) Step { + return StepFunc(func(ctx context.Context) Step { + log = append(log, "before") + result := step(next).Run(ctx) + log = append(log, "after") + return result + }) + } + }) + + pipeline := AmbientDispatch( + Do(func(ctx context.Context) context.Context { + return WithAmbientMiddleware(ctx, mw) + }), + Do(func(ctx context.Context) context.Context { + log = append(log, "step2") + return ctx + }), + Do(func(ctx context.Context) context.Context { + log = append(log, "step3") + return ctx + }), + ) + + Run(t.Context(), pipeline) + require.Equal(t, []string{"before", "step2", "before", "step3", "after", "after"}, log) +} + +// ============================================================================ +// STEP NAME CAPTURE TESTS (regression: nested clobber + parallel race) +// ============================================================================ + +// Regression for the capture-clobber bug: when an outer Named step composes an +// inner Named step, observability middleware reading the capture slot must see +// the OUTER (the step it wrapped) name, not the innermost one. +func TestCapturedStepNameUsesOutermostName(t *testing.T) { + ctx := WithStepNameCapture(t.Context()) + pipeline := Named("outer", Sequence( + Named("inner", Do(func(c context.Context) context.Context { return c })), + )) + pipeline.Step().Run(ctx) + require.Equal(t, "outer", CapturedStepName(ctx)) +} + +// Regression for the capture-race bug: parallel Named branches sharing one +// capture slot must not race when writing the name. Run under -race. +func TestCapturedStepNameNoRaceUnderParallel(t *testing.T) { + ctx := WithStepNameCapture(t.Context()) + branches := make([]NewStep, 0, 8) + for i := 0; i < 8; i++ { + branches = append(branches, Named("branch", Do(func(c context.Context) context.Context { return c }))) + } + Parallel(branches...).Step().Run(ctx) + require.Equal(t, "branch", CapturedStepName(ctx)) +} + +// Regression for the MakeMiddleware after-drop bug: the after hook must fire +// even when the wrapped step is terminal (never invokes its continuation). +func TestMakeMiddlewareAfterFiresOnTerminalStep(t *testing.T) { + var ran []string + mw := MakeMiddleware( + func(ctx context.Context) context.Context { ran = append(ran, "before"); return ctx }, + func(ctx context.Context) context.Context { ran = append(ran, "after"); return ctx }, + ) + Run(t.Context(), mw(Terminal)) + require.Equal(t, []string{"before", "after"}, ran) +} + +// Regression: the after hook must fire even when the wrapped step cancels the +// context (the failure path), so metrics/tracing record the step. +func TestMakeMiddlewareAfterFiresOnCancelledStep(t *testing.T) { + var ran []string + mw := MakeMiddleware( + func(ctx context.Context) context.Context { ran = append(ran, "before"); return ctx }, + func(ctx context.Context) context.Context { ran = append(ran, "after"); return ctx }, + ) + canceller := Do(func(ctx context.Context) context.Context { + c, cancel := context.WithCancel(ctx) + cancel() + return c + }) + Run(t.Context(), mw(canceller)) + require.Equal(t, []string{"before", "after"}, ran) +}