Skip to content

Commit

Permalink
resource control: fix unsafe usage of timer.Reset (#8877) (#8901)
Browse files Browse the repository at this point in the history
close #8876

Signed-off-by: lhy1024 <admin@liudos.us>

Co-authored-by: lhy1024 <admin@liudos.us>
  • Loading branch information
ti-chi-bot and lhy1024 authored Dec 20, 2024
1 parent bc028fc commit a5ca5e7
Show file tree
Hide file tree
Showing 28 changed files with 470 additions and 82 deletions.
4 changes: 3 additions & 1 deletion client/pd_service_discovery.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
17 changes: 9 additions & 8 deletions client/resource_group/controller/controller.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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() {
Expand Down
4 changes: 3 additions & 1 deletion client/resource_manager_client.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand Down
4 changes: 3 additions & 1 deletion client/retry/backoff.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
})
Expand Down
43 changes: 43 additions & 0 deletions client/timerutil/pool.go
Original file line number Diff line number Diff line change
@@ -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)
}
70 changes: 70 additions & 0 deletions client/timerutil/pool_test.go
Original file line number Diff line number Diff line change
@@ -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()
}
})
}
32 changes: 32 additions & 0 deletions client/timerutil/util.go
Original file line number Diff line number Diff line change
@@ -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)
}
Loading

0 comments on commit a5ca5e7

Please sign in to comment.