Add Context cancellation to clock.

common/clock/testclock/testclock.go

 // Copyright 2014 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package testclock
 
 import (
         "sync"
         "time"
 
         "github.com/luci/luci-go/common/clock"
+        "golang.org/x/net/context"
 )
 
 // TestClock is a Clock interface with additional methods to help instrument it.
 type TestClock interface {
         clock.Clock
+
+        // Set sets the test clock's time.
         Set(time.Time)
+
+        // Add advances the test clock's time.
         Add(time.Duration)
+
+        // SetTimerCallback is a goroutine-safe method to set an instance-wide
+        // callback that is invoked when any timer begins.
         SetTimerCallback(TimerCallback)
 }
 
 // TimerCallback that can be invoked when a timer has been set. This is useful
 // for sychronizing state when testing.
 type TimerCallback func(time.Duration, clock.Timer)
 
-type cancelFunc func()
-
 // testClock is a test-oriented implementation of the 'Clock' interface.
 //
 // This implementation's Clock responses are configurable by modifying its
 // member variables. Time-based events are explicitly triggered by sending on a
 // Timer instance's channel.
 type testClock struct {
         sync.Mutex
 
         now       time.Time  // The current clock time.
         timerCond *sync.Cond // Condition used to manage timer blocking.
@@ skipping 12 matching lines @@
         return &c
 }
 
 func (c *testClock) Now() time.Time {
         c.Lock()
         defer c.Unlock()
 
         return c.now
 }
 
-func (c *testClock) Sleep(d time.Duration) {
-        <-c.After(d)
+func (c *testClock) Sleep(ctx context.Context, d time.Duration) error {
+        ar := <-c.After(ctx, d)
+        return ar.Err
 }
 
-func (c *testClock) NewTimer() clock.Timer {
-        return newTimer(c)
+func (c *testClock) NewTimer(ctx context.Context) clock.Timer {
+        t := newTimer(ctx, c)
+        return t
 }
 
-func (c *testClock) After(d time.Duration) <-chan time.Time {
-        t := c.NewTimer()
+func (c *testClock) After(ctx context.Context, d time.Duration) <-chan clock.TimerResult {
+        t := newTimer(ctx, c)
         t.Reset(d)
-        return t.GetC()
+        return t.afterC
 }
 
 func (c *testClock) Set(t time.Time) {
         c.Lock()
         defer c.Unlock()
 
         c.setTimeLocked(t)
 }
 
 func (c *testClock) Add(d time.Duration) {
@@ skipping 31 matching lines @@
         if callback := c.getTimerCallback(); callback != nil {
                 callback(d, t)
         }
 }
 
 // invokeAt invokes the specified callback when the Clock has advanced at
 // or after the specified threshold.
 //
 // The returned cancelFunc can be used to cancel the blocking. If the cancel
 // function is invoked before the callback, the callback will not be invoked.
-func (c *testClock) invokeAt(threshold time.Time, callback func(time.Time)) cancelFunc {
+func (c *testClock) invokeAt(ctx context.Context, threshold time.Time, callback func(time.Time, error)) {
+        finishedC := make(chan struct{})
         stopC := make(chan struct{})
-        finishedC := make(chan struct{})
 
         // Our control goroutine will monitor both time and stop signals. It will
         // terminate when either the current time has exceeded our time threshold or
-        // it has been signalled to terminate via stopC.
+        // it has been signalled to terminate via Context.
         //
         // The lock that we take our here is owned by the following goroutine.
         c.Lock()
         go func() {
                 defer close(finishedC)
+
                 defer func() {
                         now := c.now
                         c.Unlock()
 
-                        if callback != nil {
-                                callback(now)
+                        // If we finished naturally, but our Context is Done, prefer the Context
+                        // error.
+                        var err error
+                        select {
+                        case <-ctx.Done():
+                                err = ctx.Err()
+                        default:
+                                break

[iannucci, 2016/02/10 22:30:28]

merrrr

[dnj (Google), 2016/02/11 01:26:54]

Done.

                         }
+                        callback(now, err)

[iannucci, 2016/02/10 22:30:28]

call with TimerResult?

[dnj (Google), 2016/02/11 01:26:54]

Done.

                 }()
 
                 for {
                         // Determine if we are past our signalling threshold. We can safely access
                         // our clock's time member directly, since we hold its lock from the
                         // condition firing.
                         if !c.now.Before(threshold) {
                                 return
                         }
 
                         // Wait for a signal from our clock's condition.
                         c.timerCond.Wait()
 
                         // Have we been stopped?
                         select {
                         case <-stopC:
-                                callback = nil
                                 return
 
                         default:
                                 // Nope.
                         }
                 }
         }()
 
-        return func() {
-                // Signal our goroutine to stop.
-                close(stopC)
+        // This goroutine will monitor our Context and unblock if its Done channel
+        // has an error.

[iannucci, 2016/02/10 22:30:28]

s/error/be all closed up in here.

[dnj (Google), 2016/02/11 01:26:54]

Done.

+        go func() {
+                select {
+                case <-finishedC:
+                        // Our timer has expired, so no need to further monitor.
+                        return
 
-                // Lock, then Broadcast. This ensures that if the goroutine is running, it
-                // will be at a Wait() and therefore this Broadcast will wake it.
-                func() {
-                        c.Lock()
-                        defer c.Unlock()
+                case <-ctx.Done():
+                        // If we finish at the same time our Context is canceled, we don't need to
+                        // wake our monitor (determinism).
+                        select {
+                        case <-finishedC:
+                                return
+                        default:
+                                break
+                        }
+
+                        // Our Context has been canceled. Forward this to our monitor goroutine
+                        // and wake our condition.
+                        close(stopC)
                         c.timerCond.Broadcast()
-                }()
-
-                // Block until it has finished.
-                <-finishedC
-        }
+                }
+        }()
 }