resource control: fix unsafe usage of timer.Reset (#8877)

client/pd_service_discovery.go

@@ -188,14 +188,16 @@ func (c *pdServiceDiscovery) Init() error {
 
 func (c *pdServiceDiscovery) initRetry(f func() error) error {
 	var err error
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
 	for i := 0; i < c.option.maxRetryTimes; i++ {
 		if err = f(); err == nil {
 			return nil
 		}
 		select {
 		case <-c.ctx.Done():
 			return err
-		case <-time.After(time.Second):
+		case <-ticker.C:
 		}
 	}
 	return errors.WithStack(err)

client/resource_group/controller/controller.go

@@ -31,6 +31,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	pd "github.com/tikv/pd/client"
 	"github.com/tikv/pd/client/errs"
+	"github.com/tikv/pd/client/timerutil"
 	atomicutil "go.uber.org/atomic"
 	"go.uber.org/zap"
 )
@@ -286,17 +287,17 @@ func (c *ResourceGroupsController) Start(ctx context.Context) {
 					watchMetaChannel, err = c.provider.Watch(ctx, pd.GroupSettingsPathPrefixBytes, pd.WithRev(metaRevision), pd.WithPrefix(), pd.WithPrevKV())
 					if err != nil {
 						log.Warn("watch resource group meta failed", zap.Error(err))
-						watchRetryTimer.Reset(watchRetryInterval)
+						timerutil.SafeResetTimer(watchRetryTimer, watchRetryInterval)
 						failpoint.Inject("watchStreamError", func() {
-							watchRetryTimer.Reset(20 * time.Millisecond)
+							timerutil.SafeResetTimer(watchRetryTimer, 20*time.Millisecond)
 						})
 					}
 				}
 				if watchConfigChannel == nil {
 					watchConfigChannel, err = c.provider.Watch(ctx, pd.ControllerConfigPathPrefixBytes, pd.WithRev(cfgRevision), pd.WithPrefix())
 					if err != nil {
 						log.Warn("watch resource group config failed", zap.Error(err))
-						watchRetryTimer.Reset(watchRetryInterval)
+						timerutil.SafeResetTimer(watchRetryTimer, watchRetryInterval)
 					}
 				}
 			case <-emergencyTokenAcquisitionTicker.C:
@@ -330,9 +331,9 @@ func (c *ResourceGroupsController) Start(ctx context.Context) {
 				})
 				if !ok {
 					watchMetaChannel = nil
-					watchRetryTimer.Reset(watchRetryInterval)
+					timerutil.SafeResetTimer(watchRetryTimer, watchRetryInterval)
 					failpoint.Inject("watchStreamError", func() {
-						watchRetryTimer.Reset(20 * time.Millisecond)
+						timerutil.SafeResetTimer(watchRetryTimer, 20*time.Millisecond)
 					})
 					continue
 				}
@@ -366,9 +367,9 @@ func (c *ResourceGroupsController) Start(ctx context.Context) {
 			case resp, ok := <-watchConfigChannel:
 				if !ok {
 					watchConfigChannel = nil
-					watchRetryTimer.Reset(watchRetryInterval)
+					timerutil.SafeResetTimer(watchRetryTimer, watchRetryInterval)
 					failpoint.Inject("watchStreamError", func() {
-						watchRetryTimer.Reset(20 * time.Millisecond)
+						timerutil.SafeResetTimer(watchRetryTimer, 20*time.Millisecond)
 					})
 					continue
 				}
@@ -521,7 +522,7 @@ func (c *ResourceGroupsController) sendTokenBucketRequests(ctx context.Context,
 		ClientUniqueId:        c.clientUniqueID,
 	}
 	if c.ruConfig.DegradedModeWaitDuration > 0 && c.responseDeadlineCh == nil {
-		c.run.responseDeadline.Reset(c.ruConfig.DegradedModeWaitDuration)
+		timerutil.SafeResetTimer(c.run.responseDeadline, c.ruConfig.DegradedModeWaitDuration)
 		c.responseDeadlineCh = c.run.responseDeadline.C
 	}
 	go func() {

client/resource_manager_client.go

@@ -348,6 +348,8 @@ func (c *client) tryResourceManagerConnect(ctx context.Context, connection *reso
 		err    error
 		stream rmpb.ResourceManager_AcquireTokenBucketsClient
 	)
+	ticker := time.NewTicker(retryInterval)
+	defer ticker.Stop()
 	for i := 0; i < maxRetryTimes; i++ {
 		cc, err := c.resourceManagerClient()
 		if err != nil {
@@ -365,7 +367,7 @@ func (c *client) tryResourceManagerConnect(ctx context.Context, connection *reso
 		select {
 		case <-ctx.Done():
 			return err
-		case <-time.After(retryInterval):
+		case <-ticker.C:
 		}
 	}
 	return err

client/retry/backoff.go

@@ -34,9 +34,11 @@ func (bo *BackOffer) Exec(
 	fn func() error,
 ) error {
 	if err := fn(); err != nil {
+		timer := time.NewTimer(bo.nextInterval())
+		defer timer.Stop()
 		select {
 		case <-ctx.Done():
-		case <-time.After(bo.nextInterval()):
+		case <-timer.C:
 			failpoint.Inject("backOffExecute", func() {
 				testBackOffExecuteFlag = true
 			})

client/timerutil/pool.go

@@ -0,0 +1,43 @@
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Note: This file is copied from https://go-review.googlesource.com/c/go/+/276133
+
+package timerutil
+
+import (
+	"sync"
+	"time"
+)
+
+// GlobalTimerPool is a global pool for reusing *time.Timer.
+var GlobalTimerPool TimerPool
+
+// TimerPool is a wrapper of sync.Pool which caches *time.Timer for reuse.
+type TimerPool struct {
+	pool sync.Pool
+}
+
+// Get returns a timer with a given duration.
+func (tp *TimerPool) Get(d time.Duration) *time.Timer {
+	if v := tp.pool.Get(); v != nil {
+		timer := v.(*time.Timer)
+		timer.Reset(d)
+		return timer
+	}
+	return time.NewTimer(d)
+}
+
+// Put tries to call timer.Stop() before putting it back into pool,
+// if the timer.Stop() returns false (it has either already expired or been stopped),
+// have a shot at draining the channel with residual time if there is one.
+func (tp *TimerPool) Put(timer *time.Timer) {
+	if !timer.Stop() {
+		select {
+		case <-timer.C:
+		default:
+		}
+	}
+	tp.pool.Put(timer)
+}

client/timerutil/pool_test.go

@@ -0,0 +1,70 @@
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Note: This file is copied from https://go-review.googlesource.com/c/go/+/276133
+
+package timerutil
+
+import (
+	"testing"
+	"time"
+)
+
+func TestTimerPool(t *testing.T) {
+	var tp TimerPool
+
+	for i := 0; i < 100; i++ {
+		timer := tp.Get(20 * time.Millisecond)
+
+		select {
+		case <-timer.C:
+			t.Errorf("timer expired too early")
+			continue
+		default:
+		}
+
+		select {
+		case <-time.After(100 * time.Millisecond):
+			t.Errorf("timer didn't expire on time")
+		case <-timer.C:
+		}
+
+		tp.Put(timer)
+	}
+}
+
+const timeout = 10 * time.Millisecond
+
+func BenchmarkTimerUtilization(b *testing.B) {
+	b.Run("TimerWithPool", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			t := GlobalTimerPool.Get(timeout)
+			GlobalTimerPool.Put(t)
+		}
+	})
+	b.Run("TimerWithoutPool", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			t := time.NewTimer(timeout)
+			t.Stop()
+		}
+	})
+}
+
+func BenchmarkTimerPoolParallel(b *testing.B) {
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			t := GlobalTimerPool.Get(timeout)
+			GlobalTimerPool.Put(t)
+		}
+	})
+}
+
+func BenchmarkTimerNativeParallel(b *testing.B) {
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			t := time.NewTimer(timeout)
+			t.Stop()
+		}
+	})
+}

client/timerutil/util.go

@@ -0,0 +1,32 @@
+// Copyright 2024 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package timerutil
+
+import "time"
+
+// SafeResetTimer is used to reset timer safely.
+// Before Go 1.23, the only safe way to use Reset was to call Timer.Stop and explicitly drain the timer first.
+// We need be careful here, see more details in the comments of Timer.Reset.
+// https://pkg.go.dev/time@master#Timer.Reset
+func SafeResetTimer(t *time.Timer, d time.Duration) {
+	// Stop the timer if it's not stopped.
+	if !t.Stop() {
+		select {
+		case <-t.C: // try to drain from the channel
+		default:
+		}
+	}
+	t.Reset(d)
+}