niki/vendor/github.com/redis/go-redis/v9/ring.go

package redis

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cespare/xxhash/v2"
	"github.com/dgryski/go-rendezvous" //nolint
	"github.com/redis/go-redis/v9/internal"
	"github.com/redis/go-redis/v9/internal/hashtag"
	"github.com/redis/go-redis/v9/internal/pool"
	"github.com/redis/go-redis/v9/internal/rand"
)

var errRingShardsDown = errors.New("redis: all ring shards are down")

//------------------------------------------------------------------------------

type ConsistentHash interface {
	Get(string) string
}

type rendezvousWrapper struct {
	*rendezvous.Rendezvous
}

func (w rendezvousWrapper) Get(key string) string {

	return w.Lookup(key)

}

func newRendezvous(shards []string) ConsistentHash {

	return rendezvousWrapper{rendezvous.New(shards, xxhash.Sum64String)}

}

//------------------------------------------------------------------------------

// RingOptions are used to configure a ring client and should be

// passed to NewRing.

type RingOptions struct {

	// Map of name => host:port addresses of ring shards.

	Addrs map[string]string

	// NewClient creates a shard client with provided options.

	NewClient func(opt *Options) *Client

	// ClientName will execute the `CLIENT SETNAME ClientName` command for each conn.

	ClientName string

	// Frequency of PING commands sent to check shards availability.

	// Shard is considered down after 3 subsequent failed checks.

	HeartbeatFrequency time.Duration

	// NewConsistentHash returns a consistent hash that is used

	// to distribute keys across the shards.

	//

	// See https://medium.com/@dgryski/consistent-hashing-algorithmic-tradeoffs-ef6b8e2fcae8

	// for consistent hashing algorithmic tradeoffs.

	NewConsistentHash func(shards []string) ConsistentHash

	// Following options are copied from Options struct.

	Dialer func(ctx context.Context, network, addr string) (net.Conn, error)

	OnConnect func(ctx context.Context, cn *Conn) error

	Protocol int

	Username string

	Password string

	DB int

	MaxRetries int

	MinRetryBackoff time.Duration

	MaxRetryBackoff time.Duration

	DialTimeout time.Duration

	ReadTimeout time.Duration

	WriteTimeout time.Duration

	ContextTimeoutEnabled bool

	// PoolFIFO uses FIFO mode for each node connection pool GET/PUT (default LIFO).

	PoolFIFO bool

	PoolSize int

	PoolTimeout time.Duration

	MinIdleConns int

	MaxIdleConns int

	MaxActiveConns int

	ConnMaxIdleTime time.Duration

	ConnMaxLifetime time.Duration

	TLSConfig *tls.Config

	Limiter Limiter

	DisableIndentity bool

	IdentitySuffix string
}

func (opt *RingOptions) init() {

	if opt.NewClient == nil {

		opt.NewClient = func(opt *Options) *Client {

			return NewClient(opt)

		}

	}

	if opt.HeartbeatFrequency == 0 {

		opt.HeartbeatFrequency = 500 * time.Millisecond

	}

	if opt.NewConsistentHash == nil {

		opt.NewConsistentHash = newRendezvous

	}

	if opt.MaxRetries == -1 {

		opt.MaxRetries = 0

	} else if opt.MaxRetries == 0 {

		opt.MaxRetries = 3

	}

	switch opt.MinRetryBackoff {

	case -1:

		opt.MinRetryBackoff = 0

	case 0:

		opt.MinRetryBackoff = 8 * time.Millisecond

	}

	switch opt.MaxRetryBackoff {

	case -1:

		opt.MaxRetryBackoff = 0

	case 0:

		opt.MaxRetryBackoff = 512 * time.Millisecond

	}

}

func (opt *RingOptions) clientOptions() *Options {

	return &Options{

		ClientName: opt.ClientName,

		Dialer: opt.Dialer,

		OnConnect: opt.OnConnect,

		Protocol: opt.Protocol,

		Username: opt.Username,

		Password: opt.Password,

		DB: opt.DB,

		MaxRetries: -1,

		DialTimeout: opt.DialTimeout,

		ReadTimeout: opt.ReadTimeout,

		WriteTimeout: opt.WriteTimeout,

		ContextTimeoutEnabled: opt.ContextTimeoutEnabled,

		PoolFIFO: opt.PoolFIFO,

		PoolSize: opt.PoolSize,

		PoolTimeout: opt.PoolTimeout,

		MinIdleConns: opt.MinIdleConns,

		MaxIdleConns: opt.MaxIdleConns,

		MaxActiveConns: opt.MaxActiveConns,

		ConnMaxIdleTime: opt.ConnMaxIdleTime,

		ConnMaxLifetime: opt.ConnMaxLifetime,

		TLSConfig: opt.TLSConfig,

		Limiter: opt.Limiter,

		DisableIndentity: opt.DisableIndentity,

		IdentitySuffix: opt.IdentitySuffix,
	}

}

//------------------------------------------------------------------------------

type ringShard struct {
	Client *Client

	down int32

	addr string
}

func newRingShard(opt *RingOptions, addr string) *ringShard {

	clopt := opt.clientOptions()

	clopt.Addr = addr

	return &ringShard{

		Client: opt.NewClient(clopt),

		addr: addr,
	}

}

func (shard *ringShard) String() string {

	var state string

	if shard.IsUp() {

		state = "up"

	} else {

		state = "down"

	}

	return fmt.Sprintf("%s is %s", shard.Client, state)

}

func (shard *ringShard) IsDown() bool {

	const threshold = 3

	return atomic.LoadInt32(&shard.down) >= threshold

}

func (shard *ringShard) IsUp() bool {

	return !shard.IsDown()

}

// Vote votes to set shard state and returns true if state was changed.

func (shard *ringShard) Vote(up bool) bool {

	if up {

		changed := shard.IsDown()

		atomic.StoreInt32(&shard.down, 0)

		return changed

	}

	if shard.IsDown() {

		return false

	}

	atomic.AddInt32(&shard.down, 1)

	return shard.IsDown()

}

//------------------------------------------------------------------------------

type ringSharding struct {
	opt *RingOptions

	mu sync.RWMutex

	shards *ringShards

	closed bool

	hash ConsistentHash

	numShard int

	onNewNode []func(rdb *Client)

	// ensures exclusive access to SetAddrs so there is no need

	// to hold mu for the duration of potentially long shard creation

	setAddrsMu sync.Mutex
}

type ringShards struct {
	m map[string]*ringShard

	list []*ringShard
}

func newRingSharding(opt *RingOptions) *ringSharding {

	c := &ringSharding{

		opt: opt,
	}

	c.SetAddrs(opt.Addrs)

	return c

}

func (c *ringSharding) OnNewNode(fn func(rdb *Client)) {

	c.mu.Lock()

	c.onNewNode = append(c.onNewNode, fn)

	c.mu.Unlock()

}

// SetAddrs replaces the shards in use, such that you can increase and

// decrease number of shards, that you use. It will reuse shards that

// existed before and close the ones that will not be used anymore.

func (c *ringSharding) SetAddrs(addrs map[string]string) {

	c.setAddrsMu.Lock()

	defer c.setAddrsMu.Unlock()

	cleanup := func(shards map[string]*ringShard) {

		for addr, shard := range shards {

			if err := shard.Client.Close(); err != nil {

				internal.Logger.Printf(context.Background(), "shard.Close %s failed: %s", addr, err)

			}

		}

	}

	c.mu.RLock()

	if c.closed {

		c.mu.RUnlock()

		return

	}

	existing := c.shards

	c.mu.RUnlock()

	shards, created, unused := c.newRingShards(addrs, existing)

	c.mu.Lock()

	if c.closed {

		cleanup(created)

		c.mu.Unlock()

		return

	}

	c.shards = shards

	c.rebalanceLocked()

	c.mu.Unlock()

	cleanup(unused)

}

func (c *ringSharding) newRingShards(

	addrs map[string]string, existing *ringShards,

) (shards *ringShards, created, unused map[string]*ringShard) {

	shards = &ringShards{m: make(map[string]*ringShard, len(addrs))}

	created = make(map[string]*ringShard) // indexed by addr

	unused = make(map[string]*ringShard) // indexed by addr

	if existing != nil {

		for _, shard := range existing.list {

			unused[shard.addr] = shard

		}

	}

	for name, addr := range addrs {

		if shard, ok := unused[addr]; ok {

			shards.m[name] = shard

			delete(unused, addr)

		} else {

			shard := newRingShard(c.opt, addr)

			shards.m[name] = shard

			created[addr] = shard

			for _, fn := range c.onNewNode {

				fn(shard.Client)

			}

		}

	}

	for _, shard := range shards.m {

		shards.list = append(shards.list, shard)

	}

	return

}

func (c *ringSharding) List() []*ringShard {

	var list []*ringShard

	c.mu.RLock()

	if !c.closed {

		list = c.shards.list

	}

	c.mu.RUnlock()

	return list

}

func (c *ringSharding) Hash(key string) string {

	key = hashtag.Key(key)

	var hash string

	c.mu.RLock()

	defer c.mu.RUnlock()

	if c.numShard > 0 {

		hash = c.hash.Get(key)

	}

	return hash

}

func (c *ringSharding) GetByKey(key string) (*ringShard, error) {

	key = hashtag.Key(key)

	c.mu.RLock()

	defer c.mu.RUnlock()

	if c.closed {

		return nil, pool.ErrClosed

	}

	if c.numShard == 0 {

		return nil, errRingShardsDown

	}

	shardName := c.hash.Get(key)

	if shardName == "" {

		return nil, errRingShardsDown

	}

	return c.shards.m[shardName], nil

}

func (c *ringSharding) GetByName(shardName string) (*ringShard, error) {

	if shardName == "" {

		return c.Random()

	}

	c.mu.RLock()

	defer c.mu.RUnlock()

	return c.shards.m[shardName], nil

}

func (c *ringSharding) Random() (*ringShard, error) {

	return c.GetByKey(strconv.Itoa(rand.Int()))

}

// Heartbeat monitors state of each shard in the ring.

func (c *ringSharding) Heartbeat(ctx context.Context, frequency time.Duration) {

	ticker := time.NewTicker(frequency)

	defer ticker.Stop()

	for {

		select {

		case <-ticker.C:

			var rebalance bool

			for _, shard := range c.List() {

				err := shard.Client.Ping(ctx).Err()

				isUp := err == nil || err == pool.ErrPoolTimeout

				if shard.Vote(isUp) {

					internal.Logger.Printf(ctx, "ring shard state changed: %s", shard)

					rebalance = true

				}

			}

			if rebalance {

				c.mu.Lock()

				c.rebalanceLocked()

				c.mu.Unlock()

			}

		case <-ctx.Done():

			return

		}

	}

}

// rebalanceLocked removes dead shards from the Ring.

// Requires c.mu locked.

func (c *ringSharding) rebalanceLocked() {

	if c.closed {

		return

	}

	if c.shards == nil {

		return

	}

	liveShards := make([]string, 0, len(c.shards.m))

	for name, shard := range c.shards.m {

		if shard.IsUp() {

			liveShards = append(liveShards, name)

		}

	}

	c.hash = c.opt.NewConsistentHash(liveShards)

	c.numShard = len(liveShards)

}

func (c *ringSharding) Len() int {

	c.mu.RLock()

	defer c.mu.RUnlock()

	return c.numShard

}

func (c *ringSharding) Close() error {

	c.mu.Lock()

	defer c.mu.Unlock()

	if c.closed {

		return nil

	}

	c.closed = true

	var firstErr error

	for _, shard := range c.shards.list {

		if err := shard.Client.Close(); err != nil && firstErr == nil {

			firstErr = err

		}

	}

	c.hash = nil

	c.shards = nil

	c.numShard = 0

	return firstErr

}

//------------------------------------------------------------------------------

// Ring is a Redis client that uses consistent hashing to distribute

// keys across multiple Redis servers (shards). It's safe for

// concurrent use by multiple goroutines.

//

// Ring monitors the state of each shard and removes dead shards from

// the ring. When a shard comes online it is added back to the ring. This

// gives you maximum availability and partition tolerance, but no

// consistency between different shards or even clients. Each client

// uses shards that are available to the client and does not do any

// coordination when shard state is changed.

//

// Ring should be used when you need multiple Redis servers for caching

// and can tolerate losing data when one of the servers dies.

// Otherwise you should use Redis Cluster.

type Ring struct {
	cmdable

	hooksMixin

	opt *RingOptions

	sharding *ringSharding

	cmdsInfoCache *cmdsInfoCache

	heartbeatCancelFn context.CancelFunc
}

func NewRing(opt *RingOptions) *Ring {

	opt.init()

	hbCtx, hbCancel := context.WithCancel(context.Background())

	ring := Ring{

		opt: opt,

		sharding: newRingSharding(opt),

		heartbeatCancelFn: hbCancel,
	}

	ring.cmdsInfoCache = newCmdsInfoCache(ring.cmdsInfo)

	ring.cmdable = ring.Process

	ring.initHooks(hooks{

		process: ring.process,

		pipeline: func(ctx context.Context, cmds []Cmder) error {

			return ring.generalProcessPipeline(ctx, cmds, false)

		},

		txPipeline: func(ctx context.Context, cmds []Cmder) error {

			return ring.generalProcessPipeline(ctx, cmds, true)

		},
	})

	go ring.sharding.Heartbeat(hbCtx, opt.HeartbeatFrequency)

	return &ring

}

func (c *Ring) SetAddrs(addrs map[string]string) {

	c.sharding.SetAddrs(addrs)

}

// Do create a Cmd from the args and processes the cmd.

func (c *Ring) Do(ctx context.Context, args ...interface{}) *Cmd {

	cmd := NewCmd(ctx, args...)

	_ = c.Process(ctx, cmd)

	return cmd

}

func (c *Ring) Process(ctx context.Context, cmd Cmder) error {

	err := c.processHook(ctx, cmd)

	cmd.SetErr(err)

	return err

}

// Options returns read-only Options that were used to create the client.

func (c *Ring) Options() *RingOptions {

	return c.opt

}

func (c *Ring) retryBackoff(attempt int) time.Duration {

	return internal.RetryBackoff(attempt, c.opt.MinRetryBackoff, c.opt.MaxRetryBackoff)

}

// PoolStats returns accumulated connection pool stats.

func (c *Ring) PoolStats() *PoolStats {

	shards := c.sharding.List()

	var acc PoolStats

	for _, shard := range shards {

		s := shard.Client.connPool.Stats()

		acc.Hits += s.Hits

		acc.Misses += s.Misses

		acc.Timeouts += s.Timeouts

		acc.TotalConns += s.TotalConns

		acc.IdleConns += s.IdleConns

	}

	return &acc

}

// Len returns the current number of shards in the ring.

func (c *Ring) Len() int {

	return c.sharding.Len()

}

// Subscribe subscribes the client to the specified channels.

func (c *Ring) Subscribe(ctx context.Context, channels ...string) *PubSub {

	if len(channels) == 0 {

		panic("at least one channel is required")

	}

	shard, err := c.sharding.GetByKey(channels[0])

	if err != nil {

		// TODO: return PubSub with sticky error

		panic(err)

	}

	return shard.Client.Subscribe(ctx, channels...)

}

// PSubscribe subscribes the client to the given patterns.

func (c *Ring) PSubscribe(ctx context.Context, channels ...string) *PubSub {

	if len(channels) == 0 {

		panic("at least one channel is required")

	}

	shard, err := c.sharding.GetByKey(channels[0])

	if err != nil {

		// TODO: return PubSub with sticky error

		panic(err)

	}

	return shard.Client.PSubscribe(ctx, channels...)

}

// SSubscribe Subscribes the client to the specified shard channels.

func (c *Ring) SSubscribe(ctx context.Context, channels ...string) *PubSub {

	if len(channels) == 0 {

		panic("at least one channel is required")

	}

	shard, err := c.sharding.GetByKey(channels[0])

	if err != nil {

		// TODO: return PubSub with sticky error

		panic(err)

	}

	return shard.Client.SSubscribe(ctx, channels...)

}

func (c *Ring) OnNewNode(fn func(rdb *Client)) {

	c.sharding.OnNewNode(fn)

}

// ForEachShard concurrently calls the fn on each live shard in the ring.

// It returns the first error if any.

func (c *Ring) ForEachShard(

	ctx context.Context,

	fn func(ctx context.Context, client *Client) error,

) error {

	shards := c.sharding.List()

	var wg sync.WaitGroup

	errCh := make(chan error, 1)

	for _, shard := range shards {

		if shard.IsDown() {

			continue

		}

		wg.Add(1)

		go func(shard *ringShard) {

			defer wg.Done()

			err := fn(ctx, shard.Client)

			if err != nil {

				select {

				case errCh <- err:

				default:

				}

			}

		}(shard)

	}

	wg.Wait()

	select {

	case err := <-errCh:

		return err

	default:

		return nil

	}

}

func (c *Ring) cmdsInfo(ctx context.Context) (map[string]*CommandInfo, error) {

	shards := c.sharding.List()

	var firstErr error

	for _, shard := range shards {

		cmdsInfo, err := shard.Client.Command(ctx).Result()

		if err == nil {

			return cmdsInfo, nil

		}

		if firstErr == nil {

			firstErr = err

		}

	}

	if firstErr == nil {

		return nil, errRingShardsDown

	}

	return nil, firstErr

}

func (c *Ring) cmdShard(ctx context.Context, cmd Cmder) (*ringShard, error) {

	pos := cmdFirstKeyPos(cmd)

	if pos == 0 {

		return c.sharding.Random()

	}

	firstKey := cmd.stringArg(pos)

	return c.sharding.GetByKey(firstKey)

}

func (c *Ring) process(ctx context.Context, cmd Cmder) error {

	var lastErr error

	for attempt := 0; attempt <= c.opt.MaxRetries; attempt++ {

		if attempt > 0 {

			if err := internal.Sleep(ctx, c.retryBackoff(attempt)); err != nil {

				return err

			}

		}

		shard, err := c.cmdShard(ctx, cmd)

		if err != nil {

			return err

		}

		lastErr = shard.Client.Process(ctx, cmd)

		if lastErr == nil || !shouldRetry(lastErr, cmd.readTimeout() == nil) {

			return lastErr

		}

	}

	return lastErr

}

func (c *Ring) Pipelined(ctx context.Context, fn func(Pipeliner) error) ([]Cmder, error) {

	return c.Pipeline().Pipelined(ctx, fn)

}

func (c *Ring) Pipeline() Pipeliner {

	pipe := Pipeline{

		exec: pipelineExecer(c.processPipelineHook),
	}

	pipe.init()

	return &pipe

}

func (c *Ring) TxPipelined(ctx context.Context, fn func(Pipeliner) error) ([]Cmder, error) {

	return c.TxPipeline().Pipelined(ctx, fn)

}

func (c *Ring) TxPipeline() Pipeliner {

	pipe := Pipeline{

		exec: func(ctx context.Context, cmds []Cmder) error {

			cmds = wrapMultiExec(ctx, cmds)

			return c.processTxPipelineHook(ctx, cmds)

		},
	}

	pipe.init()

	return &pipe

}

func (c *Ring) generalProcessPipeline(

	ctx context.Context, cmds []Cmder, tx bool,

) error {

	if tx {

		// Trim multi .. exec.

		cmds = cmds[1 : len(cmds)-1]

	}

	cmdsMap := make(map[string][]Cmder)

	for _, cmd := range cmds {

		hash := cmd.stringArg(cmdFirstKeyPos(cmd))

		if hash != "" {

			hash = c.sharding.Hash(hash)

		}

		cmdsMap[hash] = append(cmdsMap[hash], cmd)

	}

	var wg sync.WaitGroup

	for hash, cmds := range cmdsMap {

		wg.Add(1)

		go func(hash string, cmds []Cmder) {

			defer wg.Done()

			// TODO: retry?

			shard, err := c.sharding.GetByName(hash)

			if err != nil {

				setCmdsErr(cmds, err)

				return

			}

			if tx {

				cmds = wrapMultiExec(ctx, cmds)

				_ = shard.Client.processTxPipelineHook(ctx, cmds)

			} else {

				_ = shard.Client.processPipelineHook(ctx, cmds)

			}

		}(hash, cmds)

	}

	wg.Wait()

	return cmdsFirstErr(cmds)

}

func (c *Ring) Watch(ctx context.Context, fn func(*Tx) error, keys ...string) error {

	if len(keys) == 0 {

		return fmt.Errorf("redis: Watch requires at least one key")

	}

	var shards []*ringShard

	for _, key := range keys {

		if key != "" {

			shard, err := c.sharding.GetByKey(hashtag.Key(key))

			if err != nil {

				return err

			}

			shards = append(shards, shard)

		}

	}

	if len(shards) == 0 {

		return fmt.Errorf("redis: Watch requires at least one shard")

	}

	if len(shards) > 1 {

		for _, shard := range shards[1:] {

			if shard.Client != shards[0].Client {

				err := fmt.Errorf("redis: Watch requires all keys to be in the same shard")

				return err

			}

		}

	}

	return shards[0].Client.Watch(ctx, fn, keys...)

}

// Close closes the ring client, releasing any open resources.

//

// 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 {

	c.heartbeatCancelFn()

	return c.sharding.Close()

}