Last tweaks of implimentation before testing.

clockwork.go

@@ -119,14 +119,6 @@ func (fc *FakeClock) After(d time.Duration) <-chan time.Time {
 	return fc.NewTimer(d).Chan()
 }
 
-// afterTime is like After, but uses a time instead of a duration.
-//
-// It is used to ensure FakeClock's lock is held constant through calling
-// fc.After(t.Sub(fc.Now())). It should not be exposed externally.
-func (fc *FakeClock) afterTime(t time.Time) <-chan time.Time {
-	return fc.newTimerAtTime(t, nil).Chan()
-}
-
 // Sleep blocks until the given duration has passed on the fakeClock.
 func (fc *FakeClock) Sleep(d time.Duration) {
 	<-fc.After(d)
@@ -165,34 +157,48 @@ func (fc *FakeClock) NewTicker(d time.Duration) Ticker {
 	ft = &fakeTicker{
 		firer: newFirer(),
 		d:     d,
-		reset: func(d time.Duration) { fc.set(ft, d) },
-		stop:  func() { fc.stop(ft) },
+		reset: func(d time.Duration) {
+			fc.l.Lock()
+			defer fc.l.Unlock()
+			fc.setExpirer(ft, d)
+		},
+		stop: func() { fc.stop(ft) },
 	}
-	fc.set(ft, d)
+	fc.l.Lock()
+	defer fc.l.Unlock()
+	fc.setExpirer(ft, d)
 	return ft
 }
 
 // NewTimer returns a Timer that will fire only after calls to
 // fakeClock.Advance() have moved the clock past the given duration.
 func (fc *FakeClock) NewTimer(d time.Duration) Timer {
-	return fc.newTimer(d, nil)
+	t, _ := fc.newTimer(d, nil)
+	return t
 }
 
 // AfterFunc mimics [time.AfterFunc]; it returns a Timer that will invoke the
 // given function only after calls to fakeClock.Advance() have moved the clock
 // past the given duration.
 func (fc *FakeClock) AfterFunc(d time.Duration, f func()) Timer {
-	return fc.newTimer(d, f)
+	t, _ := fc.newTimer(d, f)
+	return t
 }
 
-// newTimer returns a new timer, using an optional afterFunc.
-func (fc *FakeClock) newTimer(d time.Duration, afterfunc func()) *fakeTimer {
+// newTimer returns a new timer using an optional afterFunc and the time that
+// timer expires.
+func (fc *FakeClock) newTimer(d time.Duration, afterfunc func()) (*fakeTimer, time.Time) {
 	ft := newFakeTimer(fc, afterfunc)
-	fc.set(ft, d)
-	return ft
+	fc.l.Lock()
+	defer fc.l.Unlock()
+	fc.setExpirer(ft, d)
+	return ft, fc.time.Add(d)
 }
 
 // newTimerAtTime is like newTimer, but uses a time instead of a duration.
+//
+// It is used to ensure FakeClock's lock is held constant through calling
+// fc.After(t.Sub(fc.Now())). It should not be exposed externally.
 func (fc *FakeClock) newTimerAtTime(t time.Time, afterfunc func()) *fakeTimer {
 	ft := newFakeTimer(fc, afterfunc)
 	fc.l.Lock()
@@ -292,13 +298,6 @@ func (fc *FakeClock) stopExpirer(e expirer) bool {
 	return true
 }
 
-// set sets an expirer to expire after a duration.
-func (fc *FakeClock) set(e expirer, d time.Duration) {
-	fc.l.Lock()
-	defer fc.l.Unlock()
-	fc.setExpirer(e, d)
-}
-
 // setExpirer sets an expirer to expire at a future point in time.
 //
 // The caller must hold fc.l.
@@ -319,16 +318,14 @@ func (fc *FakeClock) setExpirer(e expirer, d time.Duration) {
 	})
 
 	// Notify blockers of our new waiter.
-	var blocked []*blocker
 	count := len(fc.waiters)
-	for _, b := range fc.blockers {
+	fc.blockers = slices.DeleteFunc(fc.blockers, func(b *blocker) bool {
 		if b.count <= count {
 			close(b.ch)
-			continue
+			return true
 		}
-		blocked = append(blocked, b)
-	}
-	fc.blockers = blocked
+		return false
+	})
 }
 
 // firer is used by fakeTimer and fakeTicker used to help implement expirer.

context.go

@@ -34,17 +34,6 @@ func FromContext(ctx context.Context) Clock {
 	return NewRealClock()
 }
 
-type fakeClockContext struct {
-	parent context.Context
-	clock  *FakeClock
-
-	deadline time.Time
-
-	mu   sync.Mutex
-	done chan struct{}
-	err  error
-}
-
 // WithDeadline returns a context with a deadline based on a [FakeClock].
 //
 // The returned context ignores parent cancelation if the parent was cancelled
@@ -55,37 +44,55 @@ type fakeClockContext struct {
 // way to then cancel the returned context is by calling the returned
 // context.CancelFunc.
 func WithDeadline(parent context.Context, clock *FakeClock, t time.Time) (context.Context, context.CancelFunc) {
-	ctx := &fakeClockContext{
-		parent: parent,
-	}
-	cancelOnce := ctx.runCancel(clock.afterTime(t))
-	return &fakeClockContext{}, cancelOnce
+	return newFakeClockContext(parent, t, clock.newTimerAtTime(t, nil).Chan())
 }
 
 // WithTimeout returns a context with a timeout based on a [FakeClock].
 //
 // The returned context follows the same behaviors as [WithDeadline].
 func WithTimeout(parent context.Context, clock *FakeClock, d time.Duration) (context.Context, context.CancelFunc) {
-	ctx := &fakeClockContext{
-		parent: parent,
-	}
-	cancelOnce := ctx.runCancel(clock.After(d))
-	return &fakeClockContext{}, cancelOnce
+	t, deadline := clock.newTimer(d, nil)
+	return newFakeClockContext(parent, deadline, t.Chan())
 }
 
-func (c *fakeClockContext) setError(err error) {
-	c.mu.Lock()
-	defer c.mu.Unlock()
-	c.err = err
-	close(c.done)
+// fakeClockContext implements context.Context, using a fake clock for its
+// deadline.
+//
+// It ignores parent cancellation if the parent is cancelled with
+// context.DeadlineExceeded.
+type fakeClockContext struct {
+	parent context.Context
+
+	// The done channel used to track the context status based on the fake clock.
+	timerDone <-chan time.Time
+
+	deadline time.Time // The deadline, based on the fake clock's time.
+	cancel   func()
+
+	mu      sync.Mutex
+	ctxDone chan struct{} // Returned by Done() per context.Context interface.
+	err     error
+}
+
+func newFakeClockContext(parent context.Context, deadline time.Time, done <-chan time.Time) (context.Context, context.CancelFunc) {
+	ctxDone := make(chan struct{})
+	ctx := &fakeClockContext{
+		parent:    parent,
+		deadline:  deadline,
+		timerDone: done,
+		ctxDone:   ctxDone,
+		cancel:    sync.OnceFunc(func() { close(ctxDone) }),
+	}
+	go ctx.runCancel()
+	return ctx, ctx.cancel
 }
 
 func (c *fakeClockContext) Deadline() (time.Time, bool) {
-	return time.Time{}, false
+	return c.deadline, true
 }
 
 func (c *fakeClockContext) Done() <-chan struct{} {
-	return c.done
+	return c.ctxDone
 }
 
 func (c *fakeClockContext) Err() error {
@@ -98,36 +105,43 @@ func (c *fakeClockContext) Value(key any) any {
 	return c.parent.Value(key)
 }
 
-func (c *fakeClockContext) runCancel(clockExpCh <-chan time.Time) context.CancelFunc {
-	cancelCh := make(chan struct{})
-	result := sync.OnceFunc(func() {
-		close(cancelCh)
-	})
-
-	go func() {
-		select {
-		case <-clockExpCh:
-			c.setError(context.DeadlineExceeded)
-			return
-		case <-cancelCh:
-			c.setError(context.DeadlineExceeded)
+// runCancel runs the fakeClockContext's cancel goroutine and returns the
+// fakeClockContext's cancel function.
+//
+// fakeClockContext is then cancelled when any of the following occur:
+//
+//   - The fakeClockContext.done channel is closed by its timer.
+//   - The returned CancelFunc is executed.
+//   - The fakeClockContext's parent context is cancelled with an error other
+//     than context.DeadlineExceeded.
+func (c *fakeClockContext) runCancel() {
+	select {
+	case <-c.timerDone:
+		c.setError(context.DeadlineExceeded)
+		return
+	case <-c.ctxDone:
+		c.setError(context.Canceled)
+		return
+	case <-c.parent.Done():
+		if err := c.parent.Err(); !errors.Is(err, context.DeadlineExceeded) {
+			c.setError(err)
 			return
-		case <-c.parent.Done():
-			if err := c.parent.Err(); !errors.Is(err, context.DeadlineExceeded) {
-				c.setError(err)
-				return
-			}
 		}
+	}
 
-		// The parent context has hit its deadline, but because we are using a fake
-		// clock we ignore it.
-		select {
-		case <-clockExpCh:
-			c.setError(context.DeadlineExceeded)
-		case <-cancelCh:
-			c.setError(context.DeadlineExceeded)
-		}
-	}()
+	// The parent context has hit its deadline, but because we are using a fake
+	// clock we ignore it.
+	select {
+	case <-c.timerDone:
+		c.setError(context.DeadlineExceeded)
+	case <-c.ctxDone:
+		c.setError(context.Canceled)
+	}
+}
 
-	return result
+func (c *fakeClockContext) setError(err error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.err = err
+	c.cancel()
 }