fix: replace heartbeat signal channel with context.WithCancel

This commit is contained in:
Max Riveiro 2022-07-19 22:50:12 +03:00
parent f032c126db
commit 20d0ca235e
1 changed files with 13 additions and 13 deletions

16
ring.go
View File

@ -310,11 +310,10 @@ func (c *ringShards) Random() (*ringShard, error) {
}
// Heartbeat monitors state of each shard in the ring.
func (c *ringShards) Heartbeat(frequency time.Duration, closeCh chan struct{}) {
func (c *ringShards) Heartbeat(ctx context.Context, frequency time.Duration) {
ticker := time.NewTicker(frequency)
defer ticker.Stop()
ctx := context.Background()
for {
select {
case <-ticker.C:
@ -332,7 +331,7 @@ func (c *ringShards) Heartbeat(frequency time.Duration, closeCh chan struct{}) {
if rebalance {
c.rebalance()
}
case <-closeCh:
case <-ctx.Done():
return
}
}
@ -395,7 +394,7 @@ type ring struct {
opt *RingOptions
shards *ringShards
cmdsInfoCache *cmdsInfoCache //nolint:structcheck
hearbeatCloseSignal chan struct{}
heartbeatCancelFn context.CancelFunc
}
// Ring is a Redis client that uses consistent hashing to distribute
@ -421,20 +420,20 @@ type Ring struct {
func NewRing(opt *RingOptions) *Ring {
opt.init()
hearbeatCloseSignal := make(chan struct{})
hbCtx, hbCancel := context.WithCancel(context.Background())
ring := Ring{
ring: &ring{
opt: opt,
shards: newRingShards(opt),
hearbeatCloseSignal: hearbeatCloseSignal,
heartbeatCancelFn: hbCancel,
},
}
ring.cmdsInfoCache = newCmdsInfoCache(ring.cmdsInfo)
ring.cmdable = ring.Process
go ring.shards.Heartbeat(opt.HeartbeatFrequency, hearbeatCloseSignal)
go ring.shards.Heartbeat(hbCtx, opt.HeartbeatFrequency)
return &ring
}
@ -722,6 +721,7 @@ func (c *Ring) Watch(ctx context.Context, fn func(*Tx) error, keys ...string) er
// It is rare to Close a Ring, as the Ring is meant to be long-lived
// and shared between many goroutines.
func (c *Ring) Close() error {
close(c.hearbeatCloseSignal)
c.heartbeatCancelFn()
return c.shards.Close()
}