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

mafredri · mafredri · commit e6e458208c9d · 2025-12-05T13:00:31.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/cli/ssh_test.go b/cli/ssh_test.go
@@ -2052,7 +2052,6 @@ func TestSSH_Container(t *testing.T) {
 		t.Parallel()
 
 		client, workspace, agentToken := setupWorkspaceForAgent(t)
-		ctx := testutil.Context(t, testutil.WaitLong)
 		pool, err := dockertest.NewPool("")
 		require.NoError(t, err, "Could not connect to docker")
 		ct, err := pool.RunWithOptions(&dockertest.RunOptions{
@@ -2083,6 +2082,7 @@ func TestSSH_Container(t *testing.T) {
 		})
 		_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
 
+		ctx := testutil.Context(t, testutil.WaitLong)
 		inv, root := clitest.New(t, "ssh", workspace.Name, "-c", ct.Container.ID)
 		clitest.SetupConfig(t, client, root)
 		ptty := ptytest.New(t).Attach(inv)
diff --git a/testutil/ctx.go b/testutil/ctx.go
@@ -6,7 +6,33 @@ import (
 	"time"
 )
 
-func Context(t testing.TB, dur time.Duration) context.Context {
+// Context returns a context that resets its timeout when accessed from new
+// locations in the test file. The timeout does not begin until the context is
+// first used.
+//
+// This is useful for integration tests that pass contexts through many
+// subsystems, where each subsystem should get a fresh timeout window.
+//
+// Note: Each call to Done(), Deadline(), or Err() from a new line in the test
+// file resets the timeout. If you need to prevent resets (e.g., to test actual
+// timeout behavior), store the channel:
+//
+//	done := ctx.Done()  // Timeout starts, channel stored
+//	// ... do work ...
+//	select {
+//	case <-done:  // No reset, using stored channel
+//	    // handle timeout
+//	}
+//
+// Wrapping with a child context (e.g., context.WithCancel) will also prevent
+// resets since the child's methods don't call through to the parent.
+func Context(t testing.TB, timeout time.Duration) context.Context {
+	return newLazyTimeoutContext(t, timeout)
+}
+
+// ContextFixed returns a context with a fixed timeout that starts immediately.
+// Use Context() instead for contexts that should reset on new package access.
+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,175 @@
+package testutil
+
+import (
+	"context"
+	"fmt"
+	"runtime"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+)
+
+var _ context.Context = (*lazyTimeoutContext)(nil)
+
+// lazyTimeoutContext is a context.Context that resets its timeout when accessed
+// from new locations in the test file. The timeout does not begin until the
+// context is first used.
+type lazyTimeoutContext struct {
+	t       testing.TB
+	timeout time.Duration
+
+	mu            sync.Mutex
+	started       bool
+	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, if any. The deadline is set lazily
+// on first access and may be extended when accessed from new locations.
+func (c *lazyTimeoutContext) Deadline() (deadline time.Time, ok bool) {
+	c.maybeResetForLocation()
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if !c.started {
+		return time.Time{}, false
+	}
+	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; this context carries no values.
+// Note: Value() does NOT trigger lazy initialization or timeout reset.
+func (*lazyTimeoutContext) Value(any) any {
+	return nil
+}
+
+// maybeResetForLocation starts the timer on first access and resets it when
+// accessed from a new location in the test file.
+func (c *lazyTimeoutContext) maybeResetForLocation() {
+	loc := callerLocation()
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	// Don't reset if already canceled.
+	if c.err != nil {
+		return
+	}
+
+	// Always start the timer on first access, regardless of location.
+	if !c.started {
+		c.startLocked()
+		if loc != "" {
+			c.seenLocations[loc] = struct{}{}
+		}
+		if testing.Verbose() {
+			c.t.Logf("lazyTimeoutContext: started timeout for location: %s", loc)
+		}
+		return
+	}
+
+	// Only reset for known test file locations.
+	if loc == "" {
+		return
+	}
+
+	if _, seen := c.seenLocations[loc]; seen {
+		return
+	}
+	c.seenLocations[loc] = struct{}{}
+
+	// Reset deadline.
+	c.deadline = time.Now().Add(c.timeout)
+	if c.timer != nil && c.timer.Stop() {
+		c.timer.Reset(c.timeout)
+	}
+
+	if testing.Verbose() {
+		c.t.Logf("lazyTimeoutContext: reset timeout for new location: %s", loc)
+	}
+}
+
+// startLocked initializes the timer. Must be called with mu held.
+func (c *lazyTimeoutContext) startLocked() {
+	c.started = true
+	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.
+func (c *lazyTimeoutContext) cancel() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.timer != nil {
+		c.timer.Stop()
+	}
+	if c.err == nil {
+		c.err = context.Canceled
+		close(c.done)
+	}
+}
+
+// callerLocation walks the stack to find the line in a test file that
+// initiated the call. Returns empty string if not called from a test file.
+func callerLocation() string {
+	// Skip: runtime.Callers, callerLocation, maybeResetForLocation,
+	// Done/Deadline/Err, and we want to find the caller of those.
+	pc := make([]uintptr, 50)
+	n := runtime.Callers(4, pc)
+	if n == 0 {
+		return ""
+	}
+
+	frames := runtime.CallersFrames(pc[:n])
+	for {
+		frame, more := frames.Next()
+
+		// Look for frames in _test.go files.
+		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
