fix: let done mock contexts be garbage collected

README.md

@@ -361,9 +361,6 @@ case <-testCtx.Done():
 }
 ```
 
-Contexts still live when the test ends are canceled with cause `quartz.ErrMockClockTestEnded`, and
-contexts created after the test ended are returned already canceled with that cause.
-
 Call the returned cancel function as you would for `context.WithTimeout`. It is idempotent and
 removes the pending mock event. Cause variants preserve the cause for `context.Cause`.
 
@@ -375,6 +372,29 @@ parent's deadline, or a time in year 9999 if the parent has none. Libraries such
 as already expired. If your code needs the remaining mock time, compute it from the deadline or
 timeout you passed in.
 
+#### At test end
+
+When the test ends, each mock context that is still live is canceled with cause
+`quartz.ErrMockClockTestEnded`. Each context does this in a cleanup it registers on the `t` passed
+to `NewMock` when it is created. Cleanups run last registered first, so a cleanup registered before
+the context, such as a `t.Cleanup(srv.Close)` waiting for a handler blocked on the context, sees it
+canceled, unless the cleanup was already running when the context was created. Cleanups registered
+after the context see it live. A context created after the test ended is returned already canceled
+with the same cause.
+
+`srv.Close` still waits forever in two cases:
+
+- The handler creates the context after `srv.Close` started. With `go 1.24` or later in `go.mod`,
+  derive the context from `t.Context()`, which is canceled before any cleanup runs: create the
+  server with `httptest.NewUnstartedServer`, set `srv.Config.BaseContext` to return `t.Context()`,
+  call `srv.Start()`, and derive the handler's context from `r.Context()`.
+- A subtest uses a `*Mock` created with its parent's `t`. The subtest's cleanups run before the
+  parent's, so the context is still live in them. Create the `*Mock` with the subtest's `t`.
+
+If a cleanup waits for code that reads a mock ticker's or timer's `C`, that code may never receive:
+a ticker stops ticking, and a fired timer drops its unread value, in cleanups that can run before
+yours. Have that code also select on a context that is done by then, such as `t.Context()`.
+
 ## Recommended Patterns
 
 ### Options

mock.go

@@ -23,11 +23,14 @@ type TestingT interface {
 	Cleanup(func())
 }
 
-// ErrMockClockTestEnded is the context.Cause of a mock context that was still live when the test
-// ended; its Err is context.Canceled. Each context is canceled by a cleanup registered when the
-// context was created, so it is still live in cleanups registered after that, which run first. A
-// context created after the cleanup registered by NewMock has run, in a cleanup registered before
-// NewMock or by a goroutine that outlived the test, is returned already canceled with this cause.
+// ErrMockClockTestEnded is the context.Cause of a mock context that was still live when the
+// TestingT passed to NewMock ran its cleanups; its Err is context.Canceled. Each context is
+// canceled in a cleanup that it registers on that TestingT when it is created. Cleanups run last
+// registered first, so cleanups registered before the context see it canceled, unless they were
+// already running when it was created, and cleanups registered after it see it live. A subtest's
+// cleanups run before its parent's, so they see the context live when that TestingT is the
+// parent's. A context created after the cleanup registered by NewMock has run is returned already
+// canceled with this cause.
 var ErrMockClockTestEnded = errors.New("quartz: mock clock test ended")
 
 // Mock is the testing implementation of Clock.  It tracks a time that monotonically increases
@@ -201,6 +204,7 @@ func (m *Mock) newMockCtx(parent context.Context, call *apiCall, cause error, de
 		mock:         m,
 		done:         make(chan struct{}),
 	}
+	c.testEndRef = &mockCtxRef{c: c}
 
 	m.mu.Lock()
 	m.matchCallLocked(call)
@@ -212,7 +216,18 @@ func (m *Mock) newMockCtx(parent context.Context, call *apiCall, cause error, de
 	if !pastDeadline {
 		m.addEventLocked(c)
 	}
-	m.tb.Cleanup(func() { c.cancel(context.Canceled, ErrMockClockTestEnded) })
+	// Each context registers its own cleanup so it is canceled before the cleanups registered
+	// before it that have not started yet (see ErrMockClockTestEnded). The cleanup holds only the
+	// ref, which cancel clears, so a done context is not kept reachable until the test ends.
+	ref := c.testEndRef
+	m.tb.Cleanup(func() {
+		m.mu.Lock()
+		live := ref.c
+		m.mu.Unlock()
+		if live != nil {
+			live.cancel(context.Canceled, ErrMockClockTestEnded)
+		}
+	})
 	close(call.complete)
 	m.mu.Unlock()
 
@@ -315,7 +330,8 @@ func (m *Mock) removeEventLocked(e event) {
 	defer m.recomputeNextLocked()
 	for i := range m.all {
 		if m.all[i] == e {
-			m.all = append(m.all[:i], m.all[i+1:]...)
+			// slices.Delete zeroes the vacated slot, so the backing array stops referencing e.
+			m.all = slices.Delete(m.all, i, i+1)
 			return
 		}
 	}
@@ -751,14 +767,18 @@ type mockCtx struct {
 	// Cause recorded when the deadline fires.
 	cause error
 
-	// err, done, afterFuncs, and stopParent are guarded by mock.mu. err and done are written
-	// together.
+	// err, done, afterFuncs, stopParent, and testEndRef.c are guarded by mock.mu. err and done are
+	// written together.
 	err        error
 	done       chan struct{}
 	afterFuncs []*afterFunc
 	stopParent func() bool
+	testEndRef *mockCtxRef
 }
 
+// mockCtxRef refers to a mock context until cancel clears it.
+type mockCtxRef struct{ c *mockCtx }
+
 // afterFunc gives each registration a pointer identity for stop to find; funcs are not comparable.
 type afterFunc struct{ f func() }
 
@@ -960,6 +980,7 @@ func (c *mockCtx) cancelTree(err, cause error, stops *[]func() bool) {
 		return
 	}
 	c.err = err
+	c.testEndRef.c = nil
 	c.mock.removeEventLocked(c)
 	c.innerCancel(cause)
 	close(c.done)

mock_go124_test.go

@@ -0,0 +1,79 @@
+//go:build go1.24
+
+// The weak package was added in Go 1.24.
+
+package quartz_test
+
+import (
+	"context"
+	"runtime"
+	"testing"
+	"time"
+	"weak"
+
+	"github.com/coder/quartz"
+)
+
+// collectSentinel is the cause given to the context under test, so it is reachable only through
+// that context. The pointer field keeps the allocator from batching it with other tiny objects,
+// which would keep it reachable through a neighbor.
+type collectSentinel struct{ _ *int }
+
+func (*collectSentinel) Error() string { return "collect sentinel" }
+
+func TestWithTimeout_DoneContextIsCollectable(t *testing.T) {
+	t.Parallel()
+
+	// Each case ends a context whose cause is the returned sentinel. The Mock and the context's own
+	// test-end cleanup stay reachable through t, so each case checks that neither keeps the
+	// context.
+	tests := map[string]func(t *testing.T) weak.Pointer[collectSentinel]{
+		"Canceled": func(t *testing.T) weak.Pointer[collectSentinel] {
+			mClock := quartz.NewMock(t)
+			cause := &collectSentinel{}
+			_, cancel := mClock.WithTimeoutCause(context.Background(), time.Hour, cause)
+			cancel()
+			return weak.Make(cause)
+		},
+		"Expired": func(t *testing.T) weak.Pointer[collectSentinel] {
+			testCtx, testCancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer testCancel()
+			mClock := quartz.NewMock(t)
+			cause := &collectSentinel{}
+			// The cancel function is dropped: an expired context must not need it to be released.
+			_, _ = mClock.WithTimeoutCause(context.Background(), time.Second, cause)
+			mClock.Advance(time.Second).MustWait(testCtx)
+			return weak.Make(cause)
+		},
+		"CanceledWithLiveMockParent": func(t *testing.T) weak.Pointer[collectSentinel] {
+			mClock := quartz.NewMock(t)
+			parent, parentCancel := mClock.WithTimeout(context.Background(), time.Hour)
+			t.Cleanup(parentCancel)
+			cause := &collectSentinel{}
+			_, cancel := mClock.WithTimeoutCause(parent, time.Hour, cause)
+			cancel()
+			return weak.Make(cause)
+		},
+		"CanceledWithLiveStdlibParent": func(t *testing.T) weak.Pointer[collectSentinel] {
+			mClock := quartz.NewMock(t)
+			parent, parentCancel := context.WithCancel(context.Background())
+			t.Cleanup(parentCancel)
+			cause := &collectSentinel{}
+			_, cancel := mClock.WithTimeoutCause(parent, time.Hour, cause)
+			cancel()
+			return weak.Make(cause)
+		},
+	}
+	for name, endContext := range tests {
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+			cause := endContext(t)
+			// runtime.GC returns after a full cycle, including sweeping, so an unreachable sentinel
+			// is already reclaimed.
+			runtime.GC()
+			if cause.Value() != nil {
+				t.Fatal("expected the done mock context to be unreachable")
+			}
+		})
+	}
+}

mock_test.go

@@ -1024,6 +1024,50 @@ func TestWithTimeout_PastDeadline(t *testing.T) {
 	}
 }
 
+// At test end, a live mock context is canceled before the cleanups registered before it run, and a
+// context that was done earlier keeps its error and cause.
+func TestWithTimeout_TestEndCancelsBeforeEarlierCleanups(t *testing.T) {
+	t.Parallel()
+	testCtx, testCancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer testCancel()
+
+	mClock := quartz.NewMock(t)
+	var live, canceled, expired context.Context
+	// Registered before the contexts are created, like t.Cleanup(srv.Close) for a server whose
+	// goroutines block on a mock context, so it must see them already done.
+	t.Cleanup(func() {
+		if !errors.Is(live.Err(), context.Canceled) {
+			t.Errorf("expected live context Canceled, got %v", live.Err())
+		}
+		if !errors.Is(context.Cause(live), quartz.ErrMockClockTestEnded) {
+			t.Errorf("expected live context cause ErrMockClockTestEnded, got %v",
+				context.Cause(live))
+		}
+		// Contexts that were done before the test ended keep their error and cause.
+		if !errors.Is(canceled.Err(), context.Canceled) {
+			t.Errorf("expected canceled context Canceled, got %v", canceled.Err())
+		}
+		if !errors.Is(context.Cause(canceled), context.Canceled) {
+			t.Errorf("expected canceled context cause Canceled, got %v", context.Cause(canceled))
+		}
+		if !errors.Is(expired.Err(), context.DeadlineExceeded) {
+			t.Errorf("expected expired context DeadlineExceeded, got %v", expired.Err())
+		}
+		if !errors.Is(context.Cause(expired), context.DeadlineExceeded) {
+			t.Errorf("expected expired context cause DeadlineExceeded, got %v",
+				context.Cause(expired))
+		}
+	})
+
+	// The cancel function is not called, so the context is live when the test ends.
+	live, _ = mClock.WithTimeout(context.Background(), time.Hour)
+	var cancel context.CancelFunc
+	canceled, cancel = mClock.WithTimeout(context.Background(), time.Hour)
+	cancel()
+	expired, _ = mClock.WithTimeout(context.Background(), time.Minute)
+	mClock.Advance(time.Minute).MustWait(testCtx)
+}
+
 // A context created after the cleanup registered by NewMock has run might never be canceled by a
 // cleanup, so it must be done when it is returned.
 func TestWithTimeout_CreatedAfterMockCleanup(t *testing.T) {