op-service: Add SleepCtx util method

Makes it easy to wait for either a specified duration or a context being done

op-service/clock/clock.go

 // Package clock provides an abstraction for time to enable testing of functionality that uses time as an input.
 package clock
 
-import "time"
+import (
+	"context"
+	"time"
+)
 
 // Clock represents time in a way that can be provided by varying implementations.
-// Methods are designed to be direct replacements for methods in the time package.
+// Methods are designed to be direct replacements for methods in the time package,
+// with some new additions to make common patterns simple.
 type Clock interface {
 	// Now provides the current local time. Equivalent to time.Now
 	Now() time.Time
@@ -26,6 +30,10 @@ type Clock interface {
 	// NewTimer creates a new Timer that will send
 	// the current time on its channel after at least duration d.
 	NewTimer(d time.Duration) Timer
+
+	// SleepCtx sleeps until either ctx is done or the specified duration has elapsed.
+	// Returns the ctx.Err if it returns because the context is done.
+	SleepCtx(ctx context.Context, d time.Duration) error
 }
 
 // A Ticker holds a channel that delivers "ticks" of a clock at intervals
@@ -104,3 +112,14 @@ func (t *SystemTimer) Ch() <-chan time.Time {
 func (s systemClock) AfterFunc(d time.Duration, f func()) Timer {
 	return &SystemTimer{time.AfterFunc(d, f)}
 }
+
+func (s systemClock) SleepCtx(ctx context.Context, d time.Duration) error {
+	timer := s.NewTimer(d)
+	defer timer.Stop()
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.Ch():
+		return nil
+	}
+}

op-service/clock/clock_test.go

+package clock
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSystemClock_SleepCtx(t *testing.T) {
+	t.Run("ReturnWhenContextDone", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		start := time.Now()
+		err := SystemClock.SleepCtx(ctx, 5*time.Minute)
+		end := time.Now()
+		require.ErrorIs(t, err, context.Canceled)
+		// The call shouldn't block for the 5 minutes, but use a high tolerance as test servers can be slow
+		// and clocks are inaccurate.
+		require.Less(t, end.Sub(start), time.Minute)
+	})
+
+	t.Run("ReturnAfterDuration", func(t *testing.T) {
+		start := time.Now()
+		err := SystemClock.SleepCtx(context.Background(), 100*time.Millisecond)
+		end := time.Now()
+		require.NoError(t, err)
+		// Require the call to sleep for at least a little. Use a high tolerance since clocks can be quite inaccurate.
+		require.Greater(t, end.Sub(start), 5*time.Millisecond, "should sleep at least a bit")
+	})
+}

op-service/clock/common.go

+package clock
+
+import (
+	"context"
+	"time"
+)
+
+func sleepCtx(ctx context.Context, d time.Duration, c Clock) error {
+	timer := c.NewTimer(d)
+	defer timer.Stop()
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.Ch():
+		return nil
+	}
+}

op-service/clock/deterministic.go

@@ -195,6 +195,10 @@ func (s *DeterministicClock) NewTimer(d time.Duration) Timer {
 	return t
 }
 
+func (s *DeterministicClock) SleepCtx(ctx context.Context, d time.Duration) error {
+	return sleepCtx(ctx, d, s)
+}
+
 func (s *DeterministicClock) addPending(t action) {
 	s.pending = append(s.pending, t)
 	select {

op-service/clock/deterministic_test.go

@@ -315,3 +315,38 @@ func TestWaitForPending(t *testing.T) {
 		require.False(t, clock.WaitForNewPendingTask(ctx), "should have reset new pending task flag")
 	})
 }
+
+func TestSleepCtx(t *testing.T) {
+	t.Run("ReturnWhenContextComplete", func(t *testing.T) {
+		clock := NewDeterministicClock(time.UnixMilli(1000))
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		err := clock.SleepCtx(ctx, 5*time.Minute)
+		require.ErrorIs(t, err, context.Canceled)
+	})
+
+	t.Run("ReturnWhenDurationComplete", func(t *testing.T) {
+		clock := NewDeterministicClock(time.UnixMilli(1000))
+		var wg sync.WaitGroup
+		var result atomic.Value
+		wg.Add(1)
+		go func() {
+			err := clock.SleepCtx(context.Background(), 5*time.Minute)
+			if err != nil {
+				result.Store(err)
+			}
+			wg.Done()
+		}()
+
+		ctx, cancelFunc := context.WithTimeout(context.Background(), 30*time.Second)
+		defer cancelFunc()
+		// Wait until the SleepCtx is called and schedules a pending task
+		clock.WaitForNewPendingTask(ctx)
+
+		clock.AdvanceTime(5 * time.Minute)
+
+		// Wait for the call to return
+		wg.Wait()
+		require.Nil(t, result.Load())
+	})
+}