@@ -22,12 +22,20 @@ import (
2222
2323var errRingShardsDown = errors .New ("redis: all ring shards are down" )
2424
25+ // defaultShardHealthCheckFn is the default function used to check the shard liveness
26+ var defaultShardHealthCheckFn = func (ctx context.Context , client * Client ) bool {
27+ err := client .Ping (ctx ).Err ()
28+ return err == nil || err == pool .ErrPoolTimeout
29+ }
30+
2531//------------------------------------------------------------------------------
2632
2733type ConsistentHash interface {
2834 Get (string ) string
2935}
3036
37+ type ShardHealthCheckFn func (ctx context.Context , client * Client ) bool
38+
3139type rendezvousWrapper struct {
3240 * rendezvous.Rendezvous
3341}
@@ -54,10 +62,14 @@ type RingOptions struct {
5462 // ClientName will execute the `CLIENT SETNAME ClientName` command for each conn.
5563 ClientName string
5664
57- // Frequency of PING commands sent to check shards availability.
65+ // Frequency of executing ShardHealthCheckFn to check shards availability.
5866 // Shard is considered down after 3 subsequent failed checks.
5967 HeartbeatFrequency time.Duration
6068
69+ // A function used to check the shard liveness
70+ // if not set, defaults to defaultShardHealthCheckFn
71+ ShardHealthCheckFn ShardHealthCheckFn
72+
6173 // NewConsistentHash returns a consistent hash that is used
6274 // to distribute keys across the shards.
6375 //
@@ -113,6 +125,10 @@ func (opt *RingOptions) init() {
113125 opt .HeartbeatFrequency = 500 * time .Millisecond
114126 }
115127
128+ if opt .ShardHealthCheckFn == nil {
129+ opt .ShardHealthCheckFn = defaultShardHealthCheckFn
130+ }
131+
116132 if opt .NewConsistentHash == nil {
117133 opt .NewConsistentHash = newRendezvous
118134 }
@@ -408,8 +424,7 @@ func (c *ringSharding) Heartbeat(ctx context.Context, frequency time.Duration) {
408424 var rebalance bool
409425
410426 for _ , shard := range c .List () {
411- err := shard .Client .Ping (ctx ).Err ()
412- isUp := err == nil || err == pool .ErrPoolTimeout
427+ isUp := c .opt .ShardHealthCheckFn (ctx , shard .Client )
413428 if shard .Vote (isUp ) {
414429 internal .Logger .Printf (ctx , "ring shard state changed: %s" , shard )
415430 rebalance = true
0 commit comments