refactor(timer): centralize timer management with utility functions

Author: sandy2008

pkg/ring/basic_lifecycler.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/cortexproject/cortex/pkg/ring/kv"
 	"github.com/cortexproject/cortex/pkg/util/services"
+	utiltimer "github.com/cortexproject/cortex/pkg/util/timer"
 )
 
 type BasicLifecyclerDelegate interface {
@@ -328,7 +329,7 @@ func (l *BasicLifecycler) waitStableTokens(ctx context.Context, period time.Dura
 	// The first observation will occur after the specified period.
 	level.Info(l.logger).Log("msg", "waiting stable tokens", "ring", l.ringName)
 	observeTimer := time.NewTimer(period)
-	defer stopAndDrainTimer(observeTimer)
+	defer utiltimer.StopAndDrainTimer(observeTimer)
 	observeChan := observeTimer.C
 
 	for {
@@ -337,7 +338,7 @@ func (l *BasicLifecycler) waitStableTokens(ctx context.Context, period time.Dura
 			if !l.verifyTokens(ctx) {
 				// The verification has failed
 				level.Info(l.logger).Log("msg", "tokens verification failed, keep observing", "ring", l.ringName)
-				resetTimer(observeTimer, period)
+				utiltimer.ResetTimer(observeTimer, period)
 				break
 			}
 

pkg/ring/kv/dynamodb/client.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/cortexproject/cortex/pkg/ring/kv/codec"
 	"github.com/cortexproject/cortex/pkg/util/backoff"
+	utiltimer "github.com/cortexproject/cortex/pkg/util/timer"
 )
 
 const (
@@ -267,7 +268,7 @@ func (c *Client) WatchKey(ctx context.Context, key string, f func(any) bool) {
 		}
 
 		bo.Reset()
-		resetTimer(syncTimer, c.pullerSyncTime)
+		utiltimer.ResetTimer(syncTimer, c.pullerSyncTime)
 		select {
 		case <-ctx.Done():
 			return
@@ -305,7 +306,7 @@ func (c *Client) WatchPrefix(ctx context.Context, prefix string, f func(string,
 		}
 
 		bo.Reset()
-		resetTimer(syncTimer, c.pullerSyncTime)
+		utiltimer.ResetTimer(syncTimer, c.pullerSyncTime)
 		select {
 		case <-ctx.Done():
 			return
@@ -314,16 +315,6 @@ func (c *Client) WatchPrefix(ctx context.Context, prefix string, f func(string,
 	}
 }
 
-func resetTimer(timer *time.Timer, d time.Duration) {
-	if !timer.Stop() {
-		select {
-		case <-timer.C:
-		default:
-		}
-	}
-	timer.Reset(d)
-}
-
 func (c *Client) decodeMultikey(data map[string]dynamodbItem) (codec.MultiKey, error) {
 	multiKeyData := make(map[string][]byte, len(data))
 	for key, ddbItem := range data {

pkg/ring/kv/dynamodb/client_timer_benchmark_test.go

@@ -3,6 +3,8 @@ package dynamodb
 import (
 	"testing"
 	"time"
+
+	utiltimer "github.com/cortexproject/cortex/pkg/util/timer"
 )
 
 func BenchmarkWatchLoopWaitWithTimeAfter(b *testing.B) {
@@ -30,7 +32,7 @@ func BenchmarkWatchLoopWaitWithReusableTimer(b *testing.B) {
 	b.ReportAllocs()
 
 	for b.Loop() {
-		resetTimer(timer, interval)
+		utiltimer.ResetTimer(timer, interval)
 
 		select {
 		case <-ctx.Done():

pkg/ring/lifecycler.go

@@ -21,6 +21,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/ring/kv"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
 	"github.com/cortexproject/cortex/pkg/util/services"
+	utiltimer "github.com/cortexproject/cortex/pkg/util/timer"
 )
 
 var (
@@ -535,7 +536,7 @@ func (i *Lifecycler) loop(ctx context.Context) error {
 		if autoJoinTimer == nil {
 			autoJoinTimer = time.NewTimer(d)
 		} else {
-			resetTimer(autoJoinTimer, d)
+			utiltimer.ResetTimer(autoJoinTimer, d)
 		}
 		autoJoinAfter = autoJoinTimer.C
 	}
@@ -544,13 +545,13 @@ func (i *Lifecycler) loop(ctx context.Context) error {
 		if observeTimer == nil {
 			observeTimer = time.NewTimer(d)
 		} else {
-			resetTimer(observeTimer, d)
+			utiltimer.ResetTimer(observeTimer, d)
 		}
 		observeChan = observeTimer.C
 	}
 
-	defer stopAndDrainTimer(autoJoinTimer)
-	defer stopAndDrainTimer(observeTimer)
+	defer utiltimer.StopAndDrainTimer(autoJoinTimer)
+	defer utiltimer.StopAndDrainTimer(observeTimer)
 
 	if i.autoJoinOnStartup {
 		setAutoJoinAfter(i.cfg.JoinAfter)
@@ -616,7 +617,7 @@ func (i *Lifecycler) loop(ctx context.Context) error {
 			// When observing is done, observeChan is set to nil.
 
 			observeChan = nil
-			stopAndDrainTimer(observeTimer)
+			utiltimer.StopAndDrainTimer(observeTimer)
 			if s := i.GetState(); s != JOINING {
 				level.Error(i.logger).Log("msg", "unexpected state while observing tokens", "state", s, "ring", i.RingName)
 			}

pkg/ring/ticker.go

@@ -12,21 +12,3 @@ func newDisableableTicker(interval time.Duration) (func(), <-chan time.Time) {
 	tick := time.NewTicker(interval)
 	return func() { tick.Stop() }, tick.C
 }
-
-func stopAndDrainTimer(timer *time.Timer) {
-	if timer == nil {
-		return
-	}
-
-	if !timer.Stop() {
-		select {
-		case <-timer.C:
-		default:
-		}
-	}
-}
-
-func resetTimer(timer *time.Timer, d time.Duration) {
-	stopAndDrainTimer(timer)
-	timer.Reset(d)
-}

pkg/util/backoff/backoff.go

@@ -6,6 +6,8 @@ import (
 	"fmt"
 	"math/rand"
 	"time"
+
+	utiltimer "github.com/cortexproject/cortex/pkg/util/timer"
 )
 
 // Config configures a Backoff
@@ -81,12 +83,12 @@ func (b *Backoff) Wait() {
 		if b.waitTimer == nil {
 			b.waitTimer = time.NewTimer(sleepTime)
 		} else {
-			resetTimer(b.waitTimer, sleepTime)
+			utiltimer.ResetTimer(b.waitTimer, sleepTime)
 		}
 
 		select {
 		case <-b.ctx.Done():
-			stopAndDrainTimer(b.waitTimer)
+			utiltimer.StopAndDrainTimer(b.waitTimer)
 		case <-b.waitTimer.C:
 		}
 	}
@@ -123,21 +125,3 @@ func doubleDuration(value time.Duration, max time.Duration) time.Duration {
 
 	return max
 }
-
-func stopAndDrainTimer(timer *time.Timer) {
-	if timer == nil {
-		return
-	}
-
-	if !timer.Stop() {
-		select {
-		case <-timer.C:
-		default:
-		}
-	}
-}
-
-func resetTimer(timer *time.Timer, d time.Duration) {
-	stopAndDrainTimer(timer)
-	timer.Reset(d)
-}

pkg/util/backoff/backoff_test.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"testing"
 	"time"
+
+	utiltimer "github.com/cortexproject/cortex/pkg/util/timer"
 )
 
 func TestBackoff_NextDelay(t *testing.T) {
@@ -100,3 +102,61 @@ func TestBackoff_NextDelay(t *testing.T) {
 		})
 	}
 }
+
+func TestBackoff_WaitReusesTimer(t *testing.T) {
+	t.Parallel()
+
+	b := New(context.Background(), Config{
+		MinBackoff: time.Nanosecond,
+		MaxBackoff: time.Nanosecond,
+		MaxRetries: 0,
+	})
+
+	b.Wait()
+	if b.waitTimer == nil {
+		t.Fatal("expected wait timer to be initialized")
+	}
+
+	firstTimer := b.waitTimer
+
+	b.Wait()
+	if b.waitTimer != firstTimer {
+		t.Fatal("expected wait timer to be reused")
+	}
+}
+
+func TestBackoff_WaitReturnsWhenContextCancelled(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	b := New(ctx, Config{
+		MinBackoff: time.Second,
+		MaxBackoff: time.Second,
+		MaxRetries: 0,
+	})
+
+	go func() {
+		time.Sleep(10 * time.Millisecond)
+		cancel()
+	}()
+
+	startedAt := time.Now()
+	b.Wait()
+
+	if time.Since(startedAt) >= 900*time.Millisecond {
+		t.Fatal("expected Wait to return quickly after context cancellation")
+	}
+
+	if b.waitTimer == nil {
+		t.Fatal("expected wait timer to be initialized")
+	}
+
+	utiltimer.ResetTimer(b.waitTimer, time.Nanosecond)
+	select {
+	case <-b.waitTimer.C:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("expected wait timer to be reusable after cancellation")
+	}
+}

pkg/util/timer/timer.go

@@ -0,0 +1,23 @@
+package timer
+
+import "time"
+
+// StopAndDrainTimer stops the timer and drains its channel if a tick was already queued.
+func StopAndDrainTimer(timer *time.Timer) {
+	if timer == nil {
+		return
+	}
+
+	if !timer.Stop() {
+		select {
+		case <-timer.C:
+		default:
+		}
+	}
+}
+
+// ResetTimer safely resets timer, handling the required stop+drain sequence first.
+func ResetTimer(timer *time.Timer, d time.Duration) {
+	StopAndDrainTimer(timer)
+	timer.Reset(d)
+}

pkg/util/timer/timer_test.go

@@ -0,0 +1,85 @@
+package timer
+
+import (
+	"testing"
+	"time"
+)
+
+func TestStopAndDrainTimer_NilTimer(t *testing.T) {
+	// Should not panic on nil timer.
+	StopAndDrainTimer(nil)
+}
+
+func TestStopAndDrainTimer_UnfiredTimer(t *testing.T) {
+	timer := time.NewTimer(time.Hour)
+	StopAndDrainTimer(timer)
+
+	// Channel should be empty after stop+drain.
+	select {
+	case <-timer.C:
+		t.Fatal("expected timer channel to be drained")
+	default:
+	}
+}
+
+func TestStopAndDrainTimer_FiredTimer(t *testing.T) {
+	timer := time.NewTimer(time.Nanosecond)
+	// Wait for it to fire.
+	time.Sleep(time.Millisecond)
+
+	StopAndDrainTimer(timer)
+
+	// Channel should be empty after stop+drain.
+	select {
+	case <-timer.C:
+		t.Fatal("expected timer channel to be drained")
+	default:
+	}
+}
+
+func TestResetTimer(t *testing.T) {
+	timer := time.NewTimer(time.Hour)
+
+	// Reset to a very short duration.
+	ResetTimer(timer, time.Nanosecond)
+
+	select {
+	case <-timer.C:
+		// Expected.
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("expected timer to fire after reset")
+	}
+}
+
+func TestResetTimer_AfterFired(t *testing.T) {
+	timer := time.NewTimer(time.Nanosecond)
+	// Wait for it to fire.
+	time.Sleep(time.Millisecond)
+	<-timer.C
+
+	// Reset after consuming the fired event.
+	ResetTimer(timer, time.Nanosecond)
+
+	select {
+	case <-timer.C:
+		// Expected.
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("expected timer to fire after reset")
+	}
+}
+
+func TestResetTimer_MultipleTimes(t *testing.T) {
+	timer := time.NewTimer(time.Hour)
+	defer timer.Stop()
+
+	for i := 0; i < 10; i++ {
+		ResetTimer(timer, time.Nanosecond)
+
+		select {
+		case <-timer.C:
+			// Expected.
+		case <-time.After(100 * time.Millisecond):
+			t.Fatalf("iteration %d: expected timer to fire after reset", i)
+		}
+	}
+}