feat(testutil): add lazy timeout context with location-based reset

mafredri · mafredri · commit c9238e240f91 · 2025-12-08T15:59:04.000Z
It's common to create a context early in a test body, then do setup work unrelated to that context. By the time the context is actually used, it may have already timed out. This was detected as test failures in #21091. The new Context() function returns a context that resets its timeout when accessed from new lines in the test file. The timeout does not begin until the context is first used (lazy initialization). This is useful for integration tests that pass contexts through many subsystems, where each subsystem should get a fresh timeout window. Key behaviors: - Timer starts on first Done(), Deadline(), or Err() call - Value() does not trigger initialization (used for tracing/logging) - Each unique line in a _test.go file gets a fresh timeout window - Same-line access (e.g., in loops) does not reset - Expired contexts cannot be resurrected Limitations: - Wrapping with a child context (e.g., context.WithCancel) prevents resets since the child's methods don't call through to the parent - Storing the Done() channel prevents resets on subsequent accesses The original fixed-timeout behavior is available via ContextFixed().
diff --git a/testutil/ctx.go b/testutil/ctx.go
@@ -6,7 +6,21 @@ import (
 	"time"
 )
 
-func Context(t testing.TB, dur time.Duration) context.Context {
+// Context returns a context with a timeout that starts on first use and resets
+// when accessed from new lines in test files. Each call to Done, Deadline, or
+// Err from a new line resets the deadline.
+//
+// To prevent resets, store the Done channel or wrap with a child context:
+//
+//	done := ctx.Done()
+//	<-done // Uses stored channel, no reset.
+func Context(t testing.TB, timeout time.Duration) context.Context {
+	return newLazyTimeoutContext(t, timeout)
+}
+
+// ContextFixed returns a context with a timeout that starts immediately and
+// does not reset.
+func ContextFixed(t testing.TB, dur time.Duration) context.Context {
 	ctx, cancel := context.WithTimeout(context.Background(), dur)
 	t.Cleanup(cancel)
 	return ctx
diff --git a/testutil/lazy_ctx.go b/testutil/lazy_ctx.go
@@ -0,0 +1,170 @@
+package testutil
+
+import (
+	"context"
+	"fmt"
+	"runtime"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+)
+
+var _ context.Context = (*lazyTimeoutContext)(nil)
+
+// lazyTimeoutContext implements context.Context with a timeout that starts on
+// first use and resets when accessed from new locations in test files.
+type lazyTimeoutContext struct {
+	t       testing.TB
+	timeout time.Duration
+
+	mu            sync.Mutex // Protects following fields.
+	testDone      bool       // True after cancel, prevents post-test logging.
+	deadline      time.Time
+	timer         *time.Timer
+	done          chan struct{}
+	err           error
+	seenLocations map[string]struct{}
+}
+
+func newLazyTimeoutContext(t testing.TB, timeout time.Duration) context.Context {
+	ctx := &lazyTimeoutContext{
+		t:             t,
+		timeout:       timeout,
+		done:          make(chan struct{}),
+		seenLocations: make(map[string]struct{}),
+	}
+	t.Cleanup(ctx.cancel)
+	return ctx
+}
+
+// Deadline returns the current deadline. The deadline is set on first access
+// and may be extended when accessed from new locations in test files.
+func (c *lazyTimeoutContext) Deadline() (deadline time.Time, ok bool) {
+	c.maybeResetForLocation()
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.deadline, true
+}
+
+// Done returns a channel that's closed when the context is canceled.
+func (c *lazyTimeoutContext) Done() <-chan struct{} {
+	c.maybeResetForLocation()
+	return c.done
+}
+
+// Err returns the error indicating why this context was canceled.
+func (c *lazyTimeoutContext) Err() error {
+	c.maybeResetForLocation()
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.err
+}
+
+// Value returns nil. It does not trigger initialization or reset.
+func (*lazyTimeoutContext) Value(any) any {
+	return nil
+}
+
+// maybeResetForLocation starts the timer on first access and resets the
+// deadline when called from a previously unseen location in a test file.
+func (c *lazyTimeoutContext) maybeResetForLocation() {
+	loc := callerLocation()
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	// Already canceled.
+	if c.err != nil {
+		return
+	}
+
+	// First access, start timer.
+	if c.timer == nil {
+		c.startLocked()
+		if loc != "" {
+			c.seenLocations[loc] = struct{}{}
+		}
+		if testing.Verbose() && !c.testDone {
+			c.t.Logf("lazyTimeoutContext: started timeout for location: %s", loc)
+		}
+		return
+	}
+
+	// Non-test location, ignore.
+	if loc == "" {
+		return
+	}
+
+	if _, seen := c.seenLocations[loc]; seen {
+		return
+	}
+	c.seenLocations[loc] = struct{}{}
+
+	// New location, reset deadline.
+	c.deadline = time.Now().Add(c.timeout)
+	if c.timer.Stop() {
+		c.timer.Reset(c.timeout)
+	}
+
+	if testing.Verbose() && !c.testDone {
+		c.t.Logf("lazyTimeoutContext: reset timeout for new location: %s", loc)
+	}
+}
+
+// startLocked initializes the deadline and timer. It must be called with mu held.
+func (c *lazyTimeoutContext) startLocked() {
+	c.deadline = time.Now().Add(c.timeout)
+	c.timer = time.AfterFunc(c.timeout, func() {
+		c.mu.Lock()
+		defer c.mu.Unlock()
+		if c.err == nil {
+			c.err = context.DeadlineExceeded
+			close(c.done)
+		}
+	})
+}
+
+// cancel stops the timer and marks the context as canceled. It is called by
+// t.Cleanup when the test ends.
+func (c *lazyTimeoutContext) cancel() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.testDone = true
+	if c.timer != nil {
+		c.timer.Stop()
+	}
+	if c.err == nil {
+		c.err = context.Canceled
+		close(c.done)
+	}
+}
+
+// callerLocation returns the file:line of the first caller in a _test.go file,
+// or the empty string if none is found.
+func callerLocation() string {
+	// Skip runtime.Callers, callerLocation, maybeResetForLocation, and the
+	// context method (Done/Deadline/Err).
+	pc := make([]uintptr, 50)
+	n := runtime.Callers(4, pc)
+	if n == 0 {
+		return ""
+	}
+
+	frames := runtime.CallersFrames(pc[:n])
+	for {
+		frame, more := frames.Next()
+
+		if strings.HasSuffix(frame.File, "_test.go") {
+			return fmt.Sprintf("%s:%d", frame.File, frame.Line)
+		}
+
+		if !more {
+			break
+		}
+	}
+
+	return ""
+}
diff --git a/testutil/lazy_ctx_test.go b/testutil/lazy_ctx_test.go
