Merge pull request #1 from synthesio/fix-data-races

Fix data races

Author: Hugo Chargois

go.mod

@@ -1,6 +1,6 @@
 module github.com/BorisBorshevsky/timemock
 
-go 1.14
+go 1.19
 
 require (
 	github.com/golang/mock v1.4.3

time_mock_suite_test.go

@@ -1,6 +1,7 @@
 package timemock
 
 import (
+	"sync"
 	"testing"
 
 	"time"
@@ -133,3 +134,47 @@ var _ = Describe("clock", func() {
 	})
 
 })
+
+func TestRaces(t *testing.T) {
+	clock := New()
+	refTime := time.Date(2006, 1, 2, 15, 4, 5, 0, time.UTC)
+
+	n := 100
+	var wg sync.WaitGroup
+	wg.Add(n)
+	for i := 0; i < n; i++ {
+		go func() {
+			for j := 0; j < n; j++ {
+				clock.Now()
+				clock.Freeze(refTime)
+				clock.Return()
+				clock.Travel(refTime)
+				clock.Scale(2)
+			}
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+}
+
+func BenchmarkClock(b *testing.B) {
+	b.Run("Freezed", func(b *testing.B) {
+		clock := New()
+		refTime := time.Date(2006, 1, 2, 15, 4, 5, 0, time.UTC)
+		clock.Freeze(refTime)
+		for i := 0; i < b.N; i++ {
+			clock.Now()
+		}
+	})
+	b.Run("Unfreezed", func(b *testing.B) {
+		clock := New()
+		for i := 0; i < b.N; i++ {
+			clock.Now()
+		}
+	})
+	b.Run("Stdlib", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			time.Now()
+		}
+	})
+}

timemock.go

@@ -2,59 +2,64 @@ package timemock
 
 import (
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
-var (
-	now = time.Now
-)
-
 type timemockClock struct {
 	rw         *sync.RWMutex
-	frozen     bool
-	traveled   bool
+	frozen     atomic.Bool
+	traveled   atomic.Bool
 	freezeTime time.Time
 	travelTime time.Time
 	scale      float64
 }
 
 func (c *timemockClock) Scale(scale float64) {
+	c.rw.Lock()
+	defer c.rw.Unlock()
 	c.scale = scale
-	if !c.traveled {
-		c.Travel(now())
+	if !c.traveled.Load() {
+		now := time.Now()
+		c.freezeTime = now
+		c.travelTime = now
+		c.traveled.Store(true)
 	}
 }
 
 func (c *timemockClock) Now() time.Time {
-	if c.frozen || c.traveled {
-		c.rw.RLock()
-		defer c.rw.RUnlock()
+	// fast path
+	if !c.frozen.Load() && !c.traveled.Load() {
+		return time.Now()
 	}
 
-	if c.frozen {
+	c.rw.RLock()
+	defer c.rw.RUnlock()
+
+	if c.frozen.Load() {
 		return c.freezeTime
 	}
 
-	if c.traveled {
+	if c.traveled.Load() {
 		return c.freezeTime.Add(time.Duration(float64(time.Since(c.travelTime)) * c.scale))
 	}
 
-	return now()
+	return time.Now()
 }
 
 func (c *timemockClock) Freeze(t time.Time) {
 	c.rw.Lock()
 	defer c.rw.Unlock()
 	c.freezeTime = t
-	c.frozen = true
+	c.frozen.Store(true)
 }
 
 func (c *timemockClock) Travel(t time.Time) {
 	c.rw.Lock()
 	defer c.rw.Unlock()
 	c.freezeTime = t
-	c.travelTime = now()
-	c.traveled = true
+	c.travelTime = time.Now()
+	c.traveled.Store(true)
 }
 
 func (c *timemockClock) Since(t time.Time) time.Duration {
@@ -68,7 +73,7 @@ func (c *timemockClock) Until(t time.Time) time.Duration {
 func (c *timemockClock) Return() {
 	c.rw.Lock()
 	defer c.rw.Unlock()
-	c.frozen = false
-	c.traveled = false
+	c.frozen.Store(false)
+	c.traveled.Store(false)
 	c.scale = 1
 }