From bdaa629310d1dfe391513b7d3217d09553135a84 Mon Sep 17 00:00:00 2001
From: Hugon Sknadaj <hugon.sknadaj@msales.com>
Date: Wed, 30 Mar 2022 12:09:12 +0200
Subject: [PATCH] Improve race detection. (#4)

---
 clock.go            | 244 +++++++++++++++++++++++++++++++++++---------
 clock_bench_test.go |  29 ++++++
 clock_test.go       |  73 +++++++++++--
 go.mod              |   2 +-
 go.sum              |   4 +-
 race_test.go        |  22 +++-
 6 files changed, 312 insertions(+), 62 deletions(-)
 create mode 100644 clock_bench_test.go

diff --git a/clock.go b/clock.go
index 9cdbb9a..aa66725 100644
--- a/clock.go
+++ b/clock.go
@@ -1,17 +1,13 @@
 package goclock
 
 import (
+	"context"
 	"sync"
 	"time"
 
 	"github.com/benbjohnson/clock"
 )
 
-// init initializes the Clock variable with a real Clock.
-func init() {
-	Restore()
-}
-
 const (
 	// Day represents full day.
 	Day = 24 * time.Hour
@@ -20,95 +16,245 @@ const (
 )
 
 var (
-	// Clock represents a global clock.
-	Clock clock.Clock
-
-	// mutex is used to sync package mocking and restoring.
-	mutex sync.Mutex
+	// def is the package-level clock.Clock instance.
+	def = newClk(clock.New())
 )
 
+func newClk(clk clock.Clock) *localClock {
+	lClk := &localClock{
+		clock: clk,
+	}
+	lClk.mutex.Disable()
+	return lClk
+}
+
+type localClock struct {
+	// Used to sync overriding clock. Locking is disabled by Default
+	mutex mutexWrap
+	clock clock.Clock
+}
+
+func (l *localClock) After(d time.Duration) <-chan time.Time {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	return l.clock.After(d)
+}
+
+func (l *localClock) AfterFunc(d time.Duration, f func()) *clock.Timer {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	return l.clock.AfterFunc(d, f)
+}
+
+func (l *localClock) Now() time.Time {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	return l.clock.Now()
+}
+
+func (l *localClock) Since(t time.Time) time.Duration {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	return l.clock.Since(t)
+}
+
+func (l *localClock) Sleep(d time.Duration) {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	l.clock.Sleep(d)
+}
+
+func (l *localClock) Tick(d time.Duration) <-chan time.Time {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	return l.clock.Tick(d)
+}
+
+func (l *localClock) Ticker(d time.Duration) *clock.Ticker {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	return l.clock.Ticker(d)
+}
+
+func (l *localClock) Timer(d time.Duration) *clock.Timer {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	return l.clock.Timer(d)
+}
+
+func (l *localClock) Until(t time.Time) time.Duration {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	return l.clock.Until(t)
+}
+
+func (l *localClock) WithDeadline(parent context.Context, d time.Time) (context.Context, context.CancelFunc) {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	return l.clock.WithDeadline(parent, d)
+}
+
+func (l *localClock) WithTimeout(parent context.Context, t time.Duration) (context.Context, context.CancelFunc) {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	return l.clock.WithTimeout(parent, t)
+}
+
+func (l *localClock) set(clk clock.Clock) {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	l.clock = clk
+}
+
+func (l *localClock) get() clock.Clock {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	return l.clock
+}
+
+func (l *localClock) restore() {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+
+	l.clock = clock.New()
+}
+
+func (l *localClock) setLock() {
+	l.mutex.Enable()
+}
+
+func (l *localClock) setNoLock() {
+	l.mutex.Disable()
+}
+
 // After waits for the duration to elapse and then sends the current time
 func After(d time.Duration) <-chan time.Time {
-	mutex.Lock()
-	defer mutex.Unlock()
-
-	return Clock.After(d)
+	return def.After(d)
 }
 
 // AfterFunc waits for the duration to elapse and then calls f in its own goroutine.
 func AfterFunc(d time.Duration, f func()) *clock.Timer {
-	mutex.Lock()
-	defer mutex.Unlock()
-
-	return Clock.AfterFunc(d, f)
+	return def.AfterFunc(d, f)
 }
 
 // Now returns the current local time.
 func Now() time.Time {
-	mutex.Lock()
-	defer mutex.Unlock()
-
-	return Clock.Now()
+	return def.Now()
 }
 
 // Since returns the time elapsed since t.
 func Since(t time.Time) time.Duration {
-	mutex.Lock()
-	defer mutex.Unlock()
-
-	return Clock.Since(t)
+	return def.Since(t)
 }
 
 // Sleep pauses the current goroutine for at least the duration d.
 func Sleep(d time.Duration) {
-	mutex.Lock()
-	defer mutex.Unlock()
-
-	Clock.Sleep(d)
+	def.Sleep(d)
 }
 
 // Tick is a convenience wrapper for NewTicker providing access to the ticking channel only.
 func Tick(d time.Duration) <-chan time.Time {
-	mutex.Lock()
-	defer mutex.Unlock()
-
-	return Clock.Tick(d)
+	return def.Tick(d)
 }
 
 // Ticker returns a new Ticker containing a channel that will send the
 // time with a period specified by the duration argument.
 func Ticker(d time.Duration) *clock.Ticker {
-	mutex.Lock()
-	defer mutex.Unlock()
-
-	return Clock.Ticker(d)
+	return def.Ticker(d)
 }
 
 // Timer creates a new Timer that will send the current time on its channel after at least duration d.
 func Timer(d time.Duration) *clock.Timer {
-	mutex.Lock()
-	defer mutex.Unlock()
+	return def.Timer(d)
+}
 
-	return Clock.Timer(d)
+func Until(t time.Time) time.Duration {
+	return def.Until(t)
 }
 
-// Mock replaces the Clock with a mock frozen at the given time and returns it.
-func Mock(now time.Time) *clock.Mock {
-	mutex.Lock()
-	defer mutex.Unlock()
+func WithDeadline(parent context.Context, d time.Time) (context.Context, context.CancelFunc) {
+	return def.WithDeadline(parent, d)
+}
+
+func WithTimeout(parent context.Context, t time.Duration) (context.Context, context.CancelFunc) {
+	return def.WithTimeout(parent, t)
+}
+
+// Current returns current instance of clock
+func Current() clock.Clock {
+	return def.get()
+}
 
+// Mock sets the Clock with a mock frozen at the given time and returns it.
+func Mock(now time.Time) *clock.Mock {
 	mock := clock.NewMock()
 	mock.Set(now)
 
-	Clock = mock
+	def.set(mock)
 
 	return mock
 }
 
+// Set sets the clock.
+func Set(clk clock.Clock) {
+	def.set(clk)
+}
+
 // Restore replaces the Clock with the real clock.
 func Restore() {
-	mutex.Lock()
-	defer mutex.Unlock()
+	def.restore()
+}
+
+// UseLock adds locking mechanism back on clock.
+func UseLock() {
+	def.setLock()
+}
+
+// NoLock removes locking mechanism usage on clock.
+func NoLock() {
+	def.setNoLock()
+}
+
+type mutexWrap struct {
+	lock     sync.Mutex
+	disabled bool
+}
+
+func (mw *mutexWrap) Lock() {
+	if !mw.disabled {
+		mw.lock.Lock()
+	}
+}
+
+func (mw *mutexWrap) Unlock() {
+	if !mw.disabled {
+		mw.lock.Unlock()
+	}
+}
+
+func (mw *mutexWrap) Enable() {
+	mw.lock.Lock()
+	defer mw.lock.Unlock()
+
+	mw.disabled = false
+}
+
+func (mw *mutexWrap) Disable() {
+	mw.lock.Lock()
+	defer mw.lock.Unlock()
 
-	Clock = clock.New()
+	mw.disabled = true
 }
diff --git a/clock_bench_test.go b/clock_bench_test.go
new file mode 100644
index 0000000..f0d9ac6
--- /dev/null
+++ b/clock_bench_test.go
@@ -0,0 +1,29 @@
+package goclock_test
+
+import (
+	"testing"
+
+	goclock "github.com/msales/go-clock/v2"
+)
+
+func BenchmarkClock_Now_Lock(b *testing.B) {
+	goclock.UseLock()
+
+	b.SetParallelism(100)
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			goclock.Now()
+		}
+	})
+}
+
+func BenchmarkClock_Now_NoLock(b *testing.B) {
+	goclock.NoLock()
+
+	b.SetParallelism(100)
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			goclock.Now()
+		}
+	})
+}
diff --git a/clock_test.go b/clock_test.go
index 58991f8..c2dbdd1 100644
--- a/clock_test.go
+++ b/clock_test.go
@@ -1,6 +1,7 @@
 package goclock_test
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -39,7 +40,7 @@ func TestAfter(t *testing.T) {
 
 	clk := new(mockClock)
 	clk.On("After", d).Return(ch)
-	goclock.Clock = clk
+	goclock.Set(clk)
 
 	got := goclock.After(d)
 
@@ -54,7 +55,7 @@ func TestAfterFunc(t *testing.T) {
 
 	clk := new(mockClock)
 	clk.On("AfterFunc", d, mock.AnythingOfType("func()")).Return(timer)
-	goclock.Clock = clk
+	goclock.Set(clk)
 
 	got := goclock.AfterFunc(d, f)
 
@@ -67,7 +68,7 @@ func TestNow(t *testing.T) {
 
 	clk := new(mockClock)
 	clk.On("Now").Return(now)
-	goclock.Clock = clk
+	goclock.Set(clk)
 
 	got := goclock.Now()
 
@@ -81,7 +82,7 @@ func TestSince(t *testing.T) {
 
 	clk := new(mockClock)
 	clk.On("Since", now).Return(d)
-	goclock.Clock = clk
+	goclock.Set(clk)
 
 	got := goclock.Since(now)
 
@@ -94,7 +95,7 @@ func TestSleep(t *testing.T) {
 
 	clk := new(mockClock)
 	clk.On("Sleep", d)
-	goclock.Clock = clk
+	goclock.Set(clk)
 
 	goclock.Sleep(d)
 
@@ -107,7 +108,7 @@ func TestTick(t *testing.T) {
 
 	clk := new(mockClock)
 	clk.On("Tick", d).Return(ch)
-	goclock.Clock = clk
+	goclock.Set(clk)
 
 	got := goclock.Tick(d)
 
@@ -121,7 +122,7 @@ func TestTicker(t *testing.T) {
 
 	clk := new(mockClock)
 	clk.On("Ticker", d).Return(ticker)
-	goclock.Clock = clk
+	goclock.Set(clk)
 
 	got := goclock.Ticker(d)
 
@@ -135,7 +136,7 @@ func TestTimer(t *testing.T) {
 
 	clk := new(mockClock)
 	clk.On("Timer", d).Return(timer)
-	goclock.Clock = clk
+	goclock.Set(clk)
 
 	got := goclock.Timer(d)
 
@@ -143,10 +144,66 @@ func TestTimer(t *testing.T) {
 	assert.Equal(t, timer, got)
 }
 
+func TestUntil(t *testing.T) {
+	now := time.Date(2019, time.September, 30, 14, 30, 00, 00, time.UTC)
+	dur := 5 * time.Second
+
+	clk := new(mockClock)
+	clk.On("Until", now).Return(dur)
+	goclock.Set(clk)
+
+	got := goclock.Until(now)
+
+	clk.AssertExpectations(t)
+	assert.Equal(t, dur, got)
+}
+
+func TestWithDeadline(t *testing.T) {
+	ctx := context.Background()
+	now := time.Now().Add(5 * time.Second)
+
+	clk := new(mockClock)
+	clk.On("WithDeadline", ctx, now).Return(ctx, context.CancelFunc(func() {}))
+	goclock.Set(clk)
+
+	gotCtx, cancelFn := goclock.WithDeadline(ctx, now)
+
+	clk.AssertExpectations(t)
+	assert.Equal(t, ctx, gotCtx)
+	assert.IsType(t, context.CancelFunc(func() {}), cancelFn)
+}
+
+func TestWithTimeout(t *testing.T) {
+	ctx := context.Background()
+	timeout := 5 * time.Second
+
+	clk := new(mockClock)
+	clk.On("WithTimeout", ctx, timeout).Return(ctx, context.CancelFunc(func() {}))
+	goclock.Set(clk)
+
+	gotCtx, cancelFn := goclock.WithTimeout(ctx, timeout)
+
+	clk.AssertExpectations(t)
+	assert.Equal(t, ctx, gotCtx)
+	assert.IsType(t, context.CancelFunc(func() {}), cancelFn)
+}
+
 type mockClock struct {
 	mock.Mock
 }
 
+func (m *mockClock) Until(t time.Time) time.Duration {
+	return m.Called(t).Get(0).(time.Duration)
+}
+
+func (m *mockClock) WithDeadline(parent context.Context, d time.Time) (context.Context, context.CancelFunc) {
+	return m.Called(parent, d).Get(0).(context.Context), m.Called(parent, d).Get(1).(context.CancelFunc)
+}
+
+func (m *mockClock) WithTimeout(parent context.Context, t time.Duration) (context.Context, context.CancelFunc) {
+	return m.Called(parent, t).Get(0).(context.Context), m.Called(parent, t).Get(1).(context.CancelFunc)
+}
+
 func (m *mockClock) After(d time.Duration) <-chan time.Time {
 	return m.Called(d).Get(0).(<-chan time.Time)
 }
diff --git a/go.mod b/go.mod
index 8f9b4d2..f710b2f 100644
--- a/go.mod
+++ b/go.mod
@@ -3,6 +3,6 @@ module github.com/msales/go-clock/v2
 go 1.13
 
 require (
-	github.com/benbjohnson/clock v0.0.0-20161215174838-7dc76406b6d3
+	github.com/benbjohnson/clock v1.3.0
 	github.com/stretchr/testify v1.4.0
 )
diff --git a/go.sum b/go.sum
index 3342f69..e69f8fd 100644
--- a/go.sum
+++ b/go.sum
@@ -1,5 +1,5 @@
-github.com/benbjohnson/clock v0.0.0-20161215174838-7dc76406b6d3 h1:wOysYcIdqv3WnvwqFFzrYCFALPED7qkUGaLXu359GSc=
-github.com/benbjohnson/clock v0.0.0-20161215174838-7dc76406b6d3/go.mod h1:UMqtWQTnOe4byzwe7Zhwh8f8s+36uszN51sJrSIZlTE=
+github.com/benbjohnson/clock v1.3.0 h1:ip6w0uFQkncKQ979AypyG0ER7mqUSBdKLOgAle/AT8A=
+github.com/benbjohnson/clock v1.3.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
diff --git a/race_test.go b/race_test.go
index 6d03317..69ff419 100644
--- a/race_test.go
+++ b/race_test.go
@@ -1,6 +1,7 @@
 package goclock_test
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -51,13 +52,28 @@ func TestRace_Timer(t *testing.T) {
 	})
 }
 
-func TestRace_Mock(t *testing.T) {
+func TestRace_Until(t *testing.T) {
 	testMockInRace(func() {
-		goclock.Mock(time.Now())
+		now := time.Now().Add(5 * time.Second)
+		goclock.Until(now)
+	})
+}
+
+func TestRace_WithDeadline(t *testing.T) {
+	testMockInRace(func() {
+		now := time.Now().Add(5 * time.Second)
+		goclock.WithDeadline(context.Background(), now)
+	})
+}
+
+func TestRace_WithTimeout(t *testing.T) {
+	testMockInRace(func() {
+		goclock.WithTimeout(context.Background(), 5*time.Second)
 	})
 }
 
 func testMockInRace(runFunc func()) {
+	goclock.UseLock()
 	now := time.Date(2019, time.September, 30, 14, 30, 00, 00, time.UTC)
 
 	wait1 := make(chan struct{}, 1)
@@ -80,6 +96,8 @@ func testMockInRace(runFunc func()) {
 }
 
 func testInRace(runFunc func()) {
+	goclock.UseLock()
+	goclock.Restore()
 	wait1 := make(chan struct{}, 1)
 	wait2 := make(chan struct{}, 1)