Improve race detection. (#4)

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
 }

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()
+		}
+	})
+}