perf: watch the call's context with AfterFunc instead of a goroutine

hoglet.go

@@ -2,7 +2,6 @@ package hoglet
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"sync"
 	"sync/atomic"
@@ -276,31 +275,24 @@ func Wrap[IN, OUT any](c *Circuit, f WrappableFunc[IN, OUT]) WrappableFunc[IN, O
 		if err != nil {
 			// Note: any errors here are not "observed" and do not count towards the breaker's failure rate.
 			// This includes:
-			// - ErrCircuitOpen
-			// - ErrConcurrencyLimit (for blocking limited circuits)
-			// - context timeouts while blocked on concurrency limit
+			// - [ErrCircuitOpen]
+			// - [ErrConcurrencyLimitReached] (for non-blocking limited circuits)
+			// - [ErrWaitingForSlot] (for blocking limited circuits: context errors while waiting for a slot)
 			// And any other errors that may be returned by optional breaker wrappers.
 			return out, err
 		}
 
-		// The watchdog goroutine exists to record a context cancellation/deadline as a failure promptly, even if the
-		// wrapped function ignores its context and blocks. If the context can never be canceled (no deadline and no
-		// cancellation, e.g. [context.Background]), the watchdog can never fire usefully, so we skip it and the
-		// associated context allocation entirely, relying solely on the deferred observation below.
-		//
-		// TODO: allow skipping the watchdog via an option for callers that guarantee their wrapped function respects
-		// its context, trading prompt cancellation detection for one less goroutine + context allocation per call.
+		// The watchdog is a callback on the context (see [Wrap] for why), so it only costs a goroutine if the context is
+		// canceled during the call. A context that can never be canceled (e.g. [context.Background]) needs none.
 		if ctx.Done() != nil {
 			// Only here can the watchdog race the deferred observe, so dedup to ensure the - potentially wrapped -
 			// observer is observed exactly once. Without a watchdog the deferred func below is the sole observer
 			// (normal return and panic go through the same defer and are mutually exclusive), so no dedup is needed.
 			// This relies on breaker middleware observing synchronously; an async middleware observer must dedup itself.
 			obs = dedupObservableCall(obs)
 
-			obsCtx, cancel := context.WithCancelCause(ctx)
-			defer cancel(errWrappedFunctionDone)
-
-			go c.observeCtx(obs, obsCtx)
+			stop := c.watch(ctx, obs)
+			defer stop()
 		}
 
 		defer func() {
@@ -316,25 +308,16 @@ func Wrap[IN, OUT any](c *Circuit, f WrappableFunc[IN, OUT]) WrappableFunc[IN, O
 	}
 }
 
-// errWrappedFunctionDone is used to distinguish between internal and external (to the lib) context cancellations.
-var errWrappedFunctionDone = errors.New("wrapped function done")
-
-// observeCtx observes the given context for cancellation and records it as a failure.
-// It assumes [Observer] is idempotent and deduplicates calls itself.
-func (c *Circuit) observeCtx(obs Observer, ctx context.Context) {
-	// We want to observe a context error as soon as possible to open the breaker, but at the same time we want to
-	// keep the call to the wrapped function synchronous to avoid all pitfalls that come with asynchronicity.
-	<-ctx.Done()
-
-	if context.Cause(ctx) == errWrappedFunctionDone {
-		return // internal cancellation: the wrapped function returned already and is observed by its caller
-	}
-
-	// Only a failure is worth recording early. Anything else is left to the wrapped function's actual result:
-	// recording a success here would mask whatever the function ends up returning.
-	if c.isFailure(ctx.Err()) {
-		obs.Observe(true)
-	}
+// watch registers the watchdog recording a context error on ctx as a failure, and returns the function unregistering
+// it. It is kept out of the generic [Wrap], where the callback would be allocated one size class bigger.
+func (c *Circuit) watch(ctx context.Context, obs Observer) (stop func() bool) {
+	return context.AfterFunc(ctx, func() {
+		// Only a failure is worth recording early. Anything else is left to the wrapped function's actual result:
+		// recording a success here would mask whatever the function ends up returning.
+		if c.isFailure(ctx.Err()) {
+			obs.Observe(true)
+		}
+	})
 }
 
 // State represents the state of a circuit.

hoglet_test.go

@@ -3,6 +3,7 @@ package hoglet
 import (
 	"context"
 	"errors"
+	"sync/atomic"
 	"testing"
 	"testing/synctest"
 	"time"
@@ -55,6 +56,29 @@ func BenchmarkHoglet_Do_EWMA(b *testing.B) {
 	})
 }
 
+func BenchmarkHoglet_Do_EWMA_cancellable(b *testing.B) {
+	// Each goroutine calls with its own cancellable context, like concurrent requests each carrying theirs, so the
+	// circuit watches it for cancellation.
+	noop := func(context.Context, struct{}) (out struct{}, err error) { return }
+	h, err := NewCircuit(
+		NewEWMABreaker(10, 0.9),
+		WithHalfOpenDelay(time.Second),
+	)
+	require.NoError(b, err)
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	f := Wrap(h, noop)
+	b.RunParallel(func(pb *testing.PB) {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		for pb.Next() {
+			_, _ = f(ctx, struct{}{})
+		}
+	})
+}
+
 func BenchmarkHoglet_Do_SlidingWindow(b *testing.B) {
 	noop := func(context.Context, struct{}) (out struct{}, err error) { return }
 
@@ -326,6 +350,38 @@ func TestCircuit_context_error_is_observed_before_wrapped_function_returns(t *te
 	})
 }
 
+// watchedContext is a never canceled context counting the callbacks [context.AfterFunc] registers on it, and how many
+// of them it unregisters again.
+type watchedContext struct {
+	context.Context
+	done                chan struct{}
+	registered, stopped atomic.Int64
+}
+
+func (w *watchedContext) Done() <-chan struct{} { return w.done }
+
+// AfterFunc is used by [context.AfterFunc] instead of its own registration.
+func (w *watchedContext) AfterFunc(func()) (stop func() bool) {
+	w.registered.Add(1)
+	return func() bool {
+		w.stopped.Add(1)
+		return true
+	}
+}
+
+func TestCircuit_watchdog_is_unregistered_when_the_call_returns(t *testing.T) {
+	c, err := NewCircuit(&mockBreaker{})
+	require.NoError(t, err)
+
+	ctx := &watchedContext{Context: context.Background(), done: make(chan struct{})}
+	_, err = Wrap(c, noop)(ctx, noopInSuccess)
+	require.NoError(t, err)
+
+	// A watchdog left registered would pile up on long-lived contexts until they end.
+	require.EqualValues(t, 1, ctx.registered.Load(), "a cancellable context must be watched")
+	assert.EqualValues(t, 1, ctx.stopped.Load(), "the watchdog must be unregistered once the call returned")
+}
+
 func TestCircuit_failure_condition_never_called_with_nil_error(t *testing.T) {
 	conditionCalled := false
 	condition := func(err error) bool {