Fix ExpireNano overflow (#58)

* fix ExpireNano overflow

* fix flaky test

internal/clock/clock.go

@@ -1,6 +1,7 @@
 package clock
 
 import (
+	"math"
 	"sync/atomic"
 	"time"
 )
@@ -28,9 +29,28 @@ func (c *Clock) SetNowCache(n int64) {
 }
 
 func (c *Clock) ExpireNano(ttl time.Duration) int64 {
-	return c.NowNano() + ttl.Nanoseconds()
+	// Both `ttl` and `nano + ttl` can overflow, but we only handle the overflow of `nano + ttl` here.
+	// An overflowed `ttl` can be either positive or negative. If it's positive, we won't detect it since it behaves
+	// like a regular `ttl`. Users of Theine should ensure that `ttl` does not overflow (this should be the case in most scenarios
+	// unless the value is directly manipulated via math operations).
+	// When `nano + ttl` overflows, we cap the returned expiration time at `math.MaxInt64`.
+	return saturatingAdd(c.NowNano(), ttl.Nanoseconds())
 }
 
 func (c *Clock) SetStart(ts int64) {
 	c.Start = time.Unix(0, ts)
 }
+
+func saturatingAdd(a, b int64) int64 {
+	var max int64 = math.MaxInt64
+	var min int64 = math.MinInt64
+	if b > 0 && a > max-b {
+		return max
+	}
+
+	if b < 0 && a < min-b {
+		return min
+	}
+
+	return a + b
+}

internal/clock/clock_test.go

@@ -0,0 +1,59 @@
+package clock_test
+
+import (
+	"math"
+	"testing"
+	"time"
+
+	"github.com/Yiling-J/theine-go/internal/clock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestClock_NowNano(t *testing.T) {
+	c := &clock.Clock{Start: time.Now()}
+	start := c.NowNano()
+	time.Sleep(5 * time.Millisecond)
+	end := c.NowNano()
+
+	require.Greater(t, end, start)
+}
+
+func TestClock_ExpireNano(t *testing.T) {
+	c := &clock.Clock{Start: time.Now()}
+	nano := c.NowNano()
+
+	ttl := 1 * time.Second
+	expireNano := c.ExpireNano(ttl)
+	lower := nano + ttl.Nanoseconds()
+	upper := c.NowNano() + ttl.Nanoseconds()
+	require.Greater(t, expireNano, lower)
+	require.Less(t, expireNano, upper)
+
+	overflowTTL := time.Duration(math.MaxInt64)
+	expireNano = c.ExpireNano(overflowTTL)
+	require.Equal(t, int64(math.MaxInt64), expireNano)
+
+}
+
+func TestClock_RefreshNowCache(t *testing.T) {
+	c := &clock.Clock{Start: time.Now()}
+	now := c.NowNanoCached()
+	time.Sleep(5 * time.Millisecond)
+	require.Equal(t, now, c.NowNanoCached())
+
+	c.RefreshNowCache()
+	require.NotEqual(t, c.NowNanoCached(), now)
+}
+
+func TestClock_SetNowCache(t *testing.T) {
+	c := &clock.Clock{}
+	c.SetNowCache(123456789)
+	require.Equal(t, int64(123456789), c.NowNanoCached())
+}
+
+func TestClock_SetStart(t *testing.T) {
+	c := &clock.Clock{}
+	ts := time.Now().UnixNano()
+	c.SetStart(ts)
+	require.Equal(t, ts, c.Start.UnixNano())
+}

internal/persistence_test.go

@@ -188,6 +188,17 @@ func TestStorePersistence_Resize(t *testing.T) {
 
 }
 
+type DelayWriter struct {
+	f           *os.File
+	beforeWrite func()
+}
+
+func (dw *DelayWriter) Write(p []byte) (n int, err error) {
+	dw.beforeWrite()
+	time.Sleep(2 * time.Second)
+	return dw.f.Write(p)
+}
+
 func TestStorePersistence_Readonly(t *testing.T) {
 	store := NewStore[int, int](1000, false, true, nil, nil, nil, 0, 0, nil)
 	for i := 0; i < 1000; i++ {
@@ -204,8 +215,11 @@ func TestStorePersistence_Readonly(t *testing.T) {
 	require.True(t, ok)
 	require.Equal(t, 100, v)
 
+	startCounter := make(chan bool)
+	started := false
 	go func() {
 		done := false
+		<-startCounter
 		for !done {
 			select {
 			case <-persistDone:
@@ -234,8 +248,12 @@ func TestStorePersistence_Readonly(t *testing.T) {
 	f, err := os.Create("stest")
 	defer os.Remove("stest")
 	require.Nil(t, err)
-	start := counter.Load()
-	err = store.Persist(0, f)
+	err = store.Persist(0, &DelayWriter{f: f, beforeWrite: func() {
+		if !started {
+			close(startCounter)
+			started = true
+		}
+	}})
 	require.Nil(t, err)
 	f.Close()
 	persistDone <- true
@@ -247,7 +265,8 @@ func TestStorePersistence_Readonly(t *testing.T) {
 	require.Nil(t, err)
 	f.Close()
 
-	require.True(t, counter.Load()-start > 10)
+	// read should not be blocked during persistence
+	require.Greater(t, counter.Load(), uint64(100))
 
 	oldv, ok := store.Get(100)
 	require.True(t, ok)