resource control: fix unsafe usage of timer.Reset

client/resource_group/controller/controller.go

@@ -32,6 +32,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/timerutils"
 	atomicutil "go.uber.org/atomic"
 	"go.uber.org/zap"
 	"golang.org/x/exp/slices"
@@ -291,6 +292,7 @@ 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))
+						timerutils.DrainAndStopTimer(watchRetryTimer)
 						watchRetryTimer.Reset(watchRetryInterval)
 						failpoint.Inject("watchStreamError", func() {
 							watchRetryTimer.Reset(20 * time.Millisecond)
@@ -301,6 +303,7 @@ func (c *ResourceGroupsController) Start(ctx context.Context) {
 					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))
+						timerutils.DrainAndStopTimer(watchRetryTimer)
 						watchRetryTimer.Reset(watchRetryInterval)
 					}
 				}
@@ -335,6 +338,7 @@ func (c *ResourceGroupsController) Start(ctx context.Context) {
 				})
 				if !ok {
 					watchMetaChannel = nil
+					timerutils.DrainAndStopTimer(watchRetryTimer)
 					watchRetryTimer.Reset(watchRetryInterval)
 					failpoint.Inject("watchStreamError", func() {
 						watchRetryTimer.Reset(20 * time.Millisecond)
@@ -371,6 +375,7 @@ func (c *ResourceGroupsController) Start(ctx context.Context) {
 			case resp, ok := <-watchConfigChannel:
 				if !ok {
 					watchConfigChannel = nil
+					timerutils.DrainAndStopTimer(watchRetryTimer)
 					watchRetryTimer.Reset(watchRetryInterval)
 					failpoint.Inject("watchStreamError", func() {
 						watchRetryTimer.Reset(20 * time.Millisecond)

client/timerpool/pool.go → client/timerutils/pool.go

@@ -4,7 +4,7 @@
 
 // Note: This file is copied from https://go-review.googlesource.com/c/go/+/276133
 
-package timerpool
+package timerutils
 
 import (
 	"sync"

client/timerpool/pool_test.go → client/timerutils/pool_test.go

@@ -4,7 +4,7 @@
 
 // Note: This file is copied from https://go-review.googlesource.com/c/go/+/276133
 
-package timerpool
+package timerutils
 
 import (
 	"testing"

client/timerutils/utils.go

@@ -0,0 +1,30 @@
+// 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 timerutils
+
+import "time"
+
+// DrainAndStopTimer is used to drain and stop timer.
+// We need be careful here, see more details in the comments of Timer.Reset.
+// https://pkg.go.dev/time@master#Timer.Reset
+func DrainAndStopTimer(t *time.Timer) {
+	// Stop the timer if it's not stopped.
+	if !t.Stop() {
+		select {
+		case <-t.C: // try to drain from the channel
+		default:
+		}
+	}
+}

client/tso_dispatcher.go

@@ -29,7 +29,7 @@ import (
 	"github.com/tikv/pd/client/errs"
 	"github.com/tikv/pd/client/grpcutil"
 	"github.com/tikv/pd/client/retry"
-	"github.com/tikv/pd/client/timerpool"
+	"github.com/tikv/pd/client/timerutils"
 	"github.com/tikv/pd/client/tsoutil"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
@@ -151,7 +151,7 @@ func newTSDeadline(
 	done chan struct{},
 	cancel context.CancelFunc,
 ) *deadline {
-	timer := timerpool.GlobalTimerPool.Get(timeout)
+	timer := timerutils.GlobalTimerPool.Get(timeout)
 	return &deadline{
 		timer:  timer,
 		done:   done,
@@ -197,11 +197,11 @@ func (c *tsoClient) watchTSDeadline(ctx context.Context, dcLocation string) {
 					case <-d.timer.C:
 						log.Error("[tso] tso request is canceled due to timeout", zap.String("dc-location", dc), errs.ZapError(errs.ErrClientGetTSOTimeout))
 						d.cancel()
-						timerpool.GlobalTimerPool.Put(d.timer)
+						timerutils.GlobalTimerPool.Put(d.timer)
 					case <-d.done:
-						timerpool.GlobalTimerPool.Put(d.timer)
+						timerutils.GlobalTimerPool.Put(d.timer)
 					case <-ctx.Done():
-						timerpool.GlobalTimerPool.Put(d.timer)
+						timerutils.GlobalTimerPool.Put(d.timer)
 						return
 					}
 				case <-ctx.Done():
@@ -432,15 +432,7 @@ tsoBatchLoop:
 		if maxBatchWaitInterval >= 0 {
 			tbc.adjustBestBatchSize()
 		}
-		// Stop the timer if it's not stopped.
-		if !streamLoopTimer.Stop() {
-			select {
-			case <-streamLoopTimer.C: // try to drain from the channel
-			default:
-			}
-		}
-		// We need be careful here, see more details in the comments of Timer.Reset.
-		// https://pkg.go.dev/time@master#Timer.Reset
+		timerutils.DrainAndStopTimer(streamLoopTimer)
 		streamLoopTimer.Reset(c.option.timeout)
 		// Choose a stream to send the TSO gRPC request.
 	streamChoosingLoop: