Add Context cancellation to clock.

common/clock/systemclock.go

Index: common/clock/systemclock.go
diff --git a/common/clock/systemclock.go b/common/clock/systemclock.go
index b5384a411fed547f99b9067c52abd09f2d2b38b6..3897a6eb41d358e0e867191038b209b19870e21f 100644
--- a/common/clock/systemclock.go
+++ b/common/clock/systemclock.go
@@ -6,6 +6,8 @@ package clock
 
 import (
 	"time"
+
+	"golang.org/x/net/context"
 )
 
 // Implementation of Clock that uses Go's standard library.
@@ -26,14 +28,38 @@ func (systemClock) Now() time.Time {
 	return time.Now()
 }
 
-func (systemClock) Sleep(d time.Duration) {
-	time.Sleep(d)
+func (sc systemClock) Sleep(c context.Context, d time.Duration) error {
+	ar := <-sc.After(c, d)
+	return ar.Err
 }
 
-func (systemClock) NewTimer() Timer {
-	return new(systemTimer)
+func (systemClock) NewTimer(ctx context.Context) Timer {
+	t := systemTimer{
+		ctx: ctx,
+	}
+	t.reset()
+	return &t
 }
 
-func (systemClock) After(d time.Duration) <-chan time.Time {
-	return time.After(d)
+func (systemClock) After(c context.Context, d time.Duration) <-chan TimerResult {
+	ac := make(chan TimerResult)
+	go func() {
+		ar := TimerResult{}
+
+		select {
+		case <-c.Done():
+			ar.Err = c.Err()
+
+		case tm := <-time.After(d):
+			// For determinism, prefer context cancellation over full sleep.
+			select {
+			case <-c.Done():
+				ar.Err = c.Err()
+			default:
+				ar.Time = tm
+			}
+		}
+		ac <- ar
+	}()
+	return ac
 }