// Provenance-includes-license: Apache-2.0
// Provenance-includes-copyright: The Thanos Authors.

package cache

import (
	"context"
	"crypto/tls"
	"flag"
	"fmt"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/grafana/gomemcache/memcache"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"golang.org/x/sync/errgroup"

	dstls "github.com/grafana/dskit/crypto/tls"
	"github.com/grafana/dskit/dns"
	"github.com/grafana/dskit/flagext"
	"github.com/grafana/dskit/gate"
)

const (
	dnsProviderUpdateInterval = 30 * time.Second
	maxTTL                    = 30 * 24 * time.Hour
)

var (
	ErrNoMemcachedAddresses                    = errors.New("no memcached addresses provided")
	ErrMemcachedMaxAsyncConcurrencyNotPositive = errors.New("max async concurrency must be positive")

	_ Cache = (*MemcachedClient)(nil)
)

// memcachedClientBackend is an interface used to mock the underlying client in tests.
type memcachedClientBackend interface {
	GetMulti(ctx context.Context, keys []string, opts ...memcache.Option) (map[string]*memcache.Item, error)
	Set(item *memcache.Item) error
	Add(item *memcache.Item) error
	Delete(key string) error
	Decrement(key string, delta uint64) (uint64, error)
	Increment(key string, delta uint64) (uint64, error)
	Touch(key string, seconds int32) error
	Close()
	CompareAndSwap(item *memcache.Item) error
	FlushAll() error
}

// updatableServerSelector extends the interface used for picking a memcached server
// for a key to allow servers to be updated at runtime. It allows the selector used
// by the client to be mocked in tests.
type updatableServerSelector interface {
	memcache.ServerSelector

	// SetServers changes a ServerSelector's set of servers at runtime
	// and is safe for concurrent use by multiple goroutines.
	//
	// SetServers returns an error if any of the server names fail to
	// resolve. No attempt is made to connect to the server. If any
	// error occurs, no changes are made to the internal server list.
	SetServers(servers ...string) error
}

// MemcachedClientConfig is the config accepted by RemoteCacheClient.
type MemcachedClientConfig struct {
	// Addresses specifies the list of memcached addresses. The addresses get
	// resolved with the DNS provider.
	Addresses flagext.StringSliceCSV `yaml:"addresses"`

	// Timeout specifies the socket read/write timeout.
	Timeout time.Duration `yaml:"timeout"`

	// ConnectTimeout specifies the connection timeout.
	ConnectTimeout time.Duration `yaml:"connect_timeout"`

	// MinIdleConnectionsHeadroomPercentage specifies the minimum number of idle connections
	// to keep open as a percentage of the number of recently used idle connections.
	// If negative, idle connections are kept open indefinitely.
	MinIdleConnectionsHeadroomPercentage float64 `yaml:"min_idle_connections_headroom_percentage" category:"advanced"`

	// MaxIdleConnections specifies the maximum number of idle connections that
	// will be maintained per address. For better performances, this should be
	// set to a number higher than your peak parallel requests.
	MaxIdleConnections int `yaml:"max_idle_connections" category:"advanced"`

	// MaxAsyncConcurrency specifies the maximum number of SetAsync goroutines.
	MaxAsyncConcurrency int `yaml:"max_async_concurrency" category:"advanced"`

	// MaxAsyncBufferSize specifies the queue buffer size for SetAsync operations.
	MaxAsyncBufferSize int `yaml:"max_async_buffer_size" category:"advanced"`

	// MaxGetMultiConcurrency specifies the maximum number of concurrent GetMulti() operations.
	// If set to 0, concurrency is unlimited.
	MaxGetMultiConcurrency int `yaml:"max_get_multi_concurrency" category:"advanced"`

	// MaxGetMultiBatchSize specifies the maximum number of keys a single underlying
	// GetMulti() should run. If more keys are specified, internally keys are split
	// into multiple batches and fetched concurrently, honoring MaxGetMultiConcurrency parallelism.
	// If set to 0, the max batch size is unlimited.
	MaxGetMultiBatchSize int `yaml:"max_get_multi_batch_size" category:"advanced"`

	// MaxItemSize specifies the maximum size of an item stored in memcached, in bytes.
	// Items bigger than MaxItemSize are skipped. If set to 0, no maximum size is enforced.
	MaxItemSize int `yaml:"max_item_size" category:"advanced"`

	// TLSEnabled enables connecting to Memcached with TLS.
	TLSEnabled bool `yaml:"tls_enabled" category:"advanced"`

	// TLS to use to connect to the Memcached server.
	TLS dstls.ClientConfig `yaml:",inline"`
}

func (c *MemcachedClientConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
	f.Var(&c.Addresses, prefix+"addresses", "Comma-separated list of memcached addresses. Each address can be an IP address, hostname, or an entry specified in the DNS Service Discovery format.")
	f.DurationVar(&c.Timeout, prefix+"timeout", 200*time.Millisecond, "The socket read/write timeout.")
	f.DurationVar(&c.ConnectTimeout, prefix+"connect-timeout", 200*time.Millisecond, "The connection timeout.")
	f.Float64Var(&c.MinIdleConnectionsHeadroomPercentage, prefix+"min-idle-connections-headroom-percentage", -1, "The minimum number of idle connections to keep open as a percentage (0-100) of the number of recently used idle connections. If negative, idle connections are kept open indefinitely.")
	f.IntVar(&c.MaxIdleConnections, prefix+"max-idle-connections", 100, "The maximum number of idle connections that will be maintained per address.")
	f.IntVar(&c.MaxAsyncConcurrency, prefix+"max-async-concurrency", 50, "The maximum number of concurrent asynchronous operations can occur.")
	f.IntVar(&c.MaxAsyncBufferSize, prefix+"max-async-buffer-size", 25000, "The maximum number of enqueued asynchronous operations allowed.")
	f.IntVar(&c.MaxGetMultiConcurrency, prefix+"max-get-multi-concurrency", 100, "The maximum number of concurrent connections running get operations. If set to 0, concurrency is unlimited.")
	f.IntVar(&c.MaxGetMultiBatchSize, prefix+"max-get-multi-batch-size", 100, "The maximum number of keys a single underlying get operation should run. If more keys are specified, internally keys are split into multiple batches and fetched concurrently, honoring the max concurrency. If set to 0, the max batch size is unlimited.")
	f.IntVar(&c.MaxItemSize, prefix+"max-item-size", 1024*1024, "The maximum size of an item stored in memcached, in bytes. Bigger items are not stored. If set to 0, no maximum size is enforced.")
	f.BoolVar(&c.TLSEnabled, prefix+"tls-enabled", false, "Enable connecting to Memcached with TLS.")
	c.TLS.RegisterFlagsWithPrefix(prefix, f)
}

func (c *MemcachedClientConfig) Validate() error {
	if len(c.Addresses) == 0 {
		return ErrNoMemcachedAddresses
	}

	// Set async only available when MaxAsyncConcurrency > 0.
	if c.MaxAsyncConcurrency <= 0 {
		return ErrMemcachedMaxAsyncConcurrencyNotPositive
	}

	return nil
}

type MemcachedClient struct {
	logger   log.Logger
	config   MemcachedClientConfig
	client   memcachedClientBackend
	selector updatableServerSelector
	metrics  *clientMetrics
	queue    *asyncQueue

	// Name provides an identifier for the instantiated Client
	name string

	// Address provider used to keep the memcached servers list updated.
	addressProvider AddressProvider

	// Channel used to notify internal goroutines when they should quit.
	stop chan struct{}

	// Gate used to enforce the max number of concurrent GetMulti() operations.
	getMultiGate gate.Gate

	// Tracked metrics.
	clientInfo prometheus.GaugeFunc
}

// AddressProvider performs node address resolution given a list of clusters.
type AddressProvider interface {
	// Resolve resolves the provided list of memcached cluster to the actual nodes
	Resolve(context.Context, []string) error

	// Addresses returns the nodes
	Addresses() []string
}

type memcachedGetMultiResult struct {
	items map[string]*memcache.Item
	err   error
}

// NewMemcachedClientWithConfig makes a new MemcachedClient.
func NewMemcachedClientWithConfig(logger log.Logger, name string, config MemcachedClientConfig, reg prometheus.Registerer) (*MemcachedClient, error) {
	if err := config.Validate(); err != nil {
		return nil, err
	}

	// We use a custom servers selector in order to use a jump hash
	// for servers selection.
	selector := &MemcachedJumpHashSelector{}

	client := memcache.NewFromSelector(selector)
	client.Timeout = config.Timeout
	client.ConnectTimeout = config.ConnectTimeout
	client.MinIdleConnsHeadroomPercentage = config.MinIdleConnectionsHeadroomPercentage
	client.MaxIdleConns = config.MaxIdleConnections

	if config.TLSEnabled {
		cfg, err := config.TLS.GetTLSConfig()
		if err != nil {
			return nil, errors.Wrapf(err, "TLS configuration")
		}

		client.DialTimeout = func(network, address string, timeout time.Duration) (net.Conn, error) {
			base := new(net.Dialer)
			base.Timeout = timeout

			return tls.DialWithDialer(base, network, address, cfg)
		}
	}

	if reg != nil {
		reg = prometheus.WrapRegistererWith(prometheus.Labels{labelCacheName: name}, reg)
	}
	mcClient, err := newMemcachedClient(logger, client, selector, config, reg, name)
	if err != nil {
		return nil, err
	}

	// Start the background DNS resolution
	go mcClient.resolveAddrsLoop()

	// Do initial DNS resolution
	if err = mcClient.resolveAddrs(); err != nil {
		mcClient.Stop()
		return nil, err
	}

	return mcClient, nil
}

func newMemcachedClient(
	logger log.Logger,
	client memcachedClientBackend,
	selector updatableServerSelector,
	config MemcachedClientConfig,
	reg prometheus.Registerer,
	name string,
) (*MemcachedClient, error) {
	reg = prometheus.WrapRegistererWith(
		prometheus.Labels{labelCacheBackend: backendValueMemcached},
		prometheus.WrapRegistererWithPrefix(cacheMetricNamePrefix, reg))
	addressProvider := dns.NewProvider(logger, reg, dns.MiekgdnsResolverType)
	metrics := newClientMetrics(reg)

	c := &MemcachedClient{
		metrics:         metrics,
		queue:           newAsyncQueue(config.MaxAsyncBufferSize, config.MaxAsyncConcurrency),
		logger:          log.With(logger, "name", name),
		config:          config,
		client:          client,
		selector:        selector,
		addressProvider: addressProvider,
		stop:            make(chan struct{}, 1),
		getMultiGate: gate.New(
			prometheus.WrapRegistererWithPrefix(getMultiMetricNamePrefix, reg),
			config.MaxGetMultiConcurrency,
		),
	}

	c.clientInfo = promauto.With(reg).NewGaugeFunc(prometheus.GaugeOpts{
		Name: clientInfoMetricName,
		Help: "A metric with a constant '1' value labeled by configuration options from which memcached client was configured.",
		ConstLabels: prometheus.Labels{
			"timeout": config.Timeout.String(),
			"min_idle_connections_headroom_percentage": fmt.Sprintf("%f.2", config.MinIdleConnectionsHeadroomPercentage),
			"max_idle_connections":                     strconv.Itoa(config.MaxIdleConnections),
			"max_async_concurrency":                    strconv.Itoa(config.MaxAsyncConcurrency),
			"max_async_buffer_size":                    strconv.Itoa(config.MaxAsyncBufferSize),
			"max_item_size":                            strconv.FormatUint(uint64(config.MaxItemSize), 10),
			"max_get_multi_concurrency":                strconv.Itoa(config.MaxGetMultiConcurrency),
			"max_get_multi_batch_size":                 strconv.Itoa(config.MaxGetMultiBatchSize),
		},
	},
		func() float64 { return 1 },
	)

	return c, nil
}

func (c *MemcachedClient) Stop() {
	close(c.stop)

	// Stop running async operations.
	c.queue.stop()

	// Stop the underlying client.
	c.client.Close()
}

func (c *MemcachedClient) Name() string {
	return c.name
}

func (c *MemcachedClient) SetMultiAsync(data map[string][]byte, ttl time.Duration) {
	for k, v := range data {
		c.SetAsync(k, v, ttl)
	}
}

func (c *MemcachedClient) SetAsync(key string, value []byte, ttl time.Duration) {
	if !c.itemSizeIsAcceptable(key, value, opSet) {
		return
	}

	err := c.queue.submit(func() {
		// Because this operation is executed in a separate goroutine: We run the operation without
		// a context (it is expected to keep running no matter what happens) and we don't return the
		// error (it will be tracked via metrics instead of being returned to the caller).
		_ = c.storeOperation(context.Background(), key, value, ttl, opSet, func(_ context.Context, key string, value []byte, ttl time.Duration) error {
			return c.setSingleItem(key, value, ttl)
		})
	})

	if err != nil {
		c.metrics.skipped.WithLabelValues(opSet, reasonAsyncBufferFull).Inc()
		level.Debug(c.logger).Log("msg", "failed to store item to cache because the async buffer is full", "key", key, "err", err, "buffer_size", c.config.MaxAsyncBufferSize)
	}
}

func (c *MemcachedClient) itemSizeIsAcceptable(key string, value []byte, op string) bool {
	if c.config.MaxItemSize > 0 && len(value) > c.config.MaxItemSize {
		c.metrics.skipped.WithLabelValues(op, reasonMaxItemSize).Inc()
		c.metrics.skippedDataSize.WithLabelValues(op).Observe(float64(len(value)))
		level.Debug(c.logger).Log("msg", "failed to store item to cache because it is too large", "key", key, "size", len(value), "max_item_size", c.config.MaxItemSize, "op", op)
		return false
	}

	return true
}

func (c *MemcachedClient) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error {
	return c.storeOperation(ctx, key, value, ttl, opSet, func(ctx context.Context, key string, value []byte, ttl time.Duration) error {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
			return c.setSingleItem(key, value, ttl)
		}
	})
}

func (c *MemcachedClient) setSingleItem(key string, value []byte, ttl time.Duration) error {
	ttlSeconds, ok := toSeconds(ttl)
	if !ok {
		return fmt.Errorf("%w: for set operation on %s %s", ErrInvalidTTL, key, ttl)
	}

	return c.client.Set(&memcache.Item{
		Key:        key,
		Value:      value,
		Expiration: ttlSeconds,
	})
}

func (c *MemcachedClient) Add(ctx context.Context, key string, value []byte, ttl time.Duration) error {
	return c.storeOperation(ctx, key, value, ttl, opAdd, func(ctx context.Context, key string, value []byte, ttl time.Duration) error {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
			ttlSeconds, ok := toSeconds(ttl)
			if !ok {
				return fmt.Errorf("%w: for add operation on %s %s", ErrInvalidTTL, key, ttl)
			}

			err := c.client.Add(&memcache.Item{
				Key:        key,
				Value:      value,
				Expiration: ttlSeconds,
			})

			if errors.Is(err, memcache.ErrNotStored) {
				return fmt.Errorf("%w: for add operation on %s", ErrNotStored, key)
			}

			return err
		}
	})
}

// toSeconds converts a time.Duration to seconds as an int32 and returns a boolean
// indicating if the value is valid to be used as a TTL. Durations might not be valid
// to be used for a TTL if they are non-zero but less than a second long (Memcached
// uses seconds for TTL units but "0" to mean infinite TTL) or if they are longer than
// 30 days (Memcached treats TTLs more than 30 days as UNIX timestamps).
func toSeconds(d time.Duration) (int32, bool) {
	if d > maxTTL {
		return 0, false
	}

	secs := int32(d.Seconds())
	if d != 0 && secs <= 0 {
		return 0, false
	}

	return secs, true
}

func toMemcacheOptions(opts ...Option) []memcache.Option {
	if len(opts) == 0 {
		return nil
	}

	base := &Options{}
	for _, opt := range opts {
		opt(base)
	}

	var out []memcache.Option
	if base.Alloc != nil {
		out = append(out, memcache.WithAllocator(base.Alloc))
	}

	return out
}

func (c *MemcachedClient) GetMulti(ctx context.Context, keys []string, opts ...Option) map[string][]byte {
	hits, err := c.GetMultiWithError(ctx, keys, opts...)
	if err != nil {
		if errors.Is(err, context.Canceled) {
			return nil
		}
		level.Warn(c.logger).Log("msg", "failed to fetch items from memcached", "numKeys", len(keys), "firstKey", keys[0], "err", err)
	}
	return hits
}

func (c *MemcachedClient) GetMultiWithError(ctx context.Context, keys []string, opts ...Option) (map[string][]byte, error) {
	if len(keys) == 0 {
		return nil, nil
	}

	c.metrics.requests.Add(float64(len(keys)))
	options := toMemcacheOptions(opts...)
	batches, err := c.getMultiBatched(ctx, keys, options...)

	// Build hits map from successful batches even if there was an error,
	// since some batch requests may have succeeded.
	hits := map[string][]byte{}
	for _, items := range batches {
		for key, item := range items {
			hits[key] = item.Value
		}
	}

	c.metrics.hits.Add(float64(len(hits)))
	return hits, err
}

func (c *MemcachedClient) Delete(ctx context.Context, key string) error {
	start := time.Now()
	c.metrics.operations.WithLabelValues(opDelete).Inc()

	var err error
	select {
	case <-ctx.Done():
		err = ctx.Err()
	default:
		err = c.client.Delete(key)
	}
	if err != nil {
		c.trackError(
			opDelete, err,
			"msg", "failed to delete cache item",
			"key", key,
		)
	} else {
		c.metrics.duration.WithLabelValues(opDelete).Observe(time.Since(start).Seconds())
	}
	return err
}

func (c *MemcachedClient) Increment(ctx context.Context, key string, delta uint64) (uint64, error) {
	return c.incrDecr(ctx, key, opIncrement, func() (uint64, error) {
		return c.client.Increment(key, delta)
	})
}

func (c *MemcachedClient) Decrement(ctx context.Context, key string, delta uint64) (uint64, error) {
	return c.incrDecr(ctx, key, opDecrement, func() (uint64, error) {
		return c.client.Decrement(key, delta)
	})
}

func (c *MemcachedClient) incrDecr(ctx context.Context, key string, operation string, f func() (uint64, error)) (uint64, error) {
	var (
		newValue uint64
		err      error
	)
	start := time.Now()
	c.metrics.operations.WithLabelValues(operation).Inc()

	select {
	case <-ctx.Done():
		err = ctx.Err()
	default:
		newValue, err = f()
	}
	if err != nil {
		c.trackError(
			operation, err,
			"msg", "failed to incr/decr cache item",
			"operation", operation,
			"key", key,
		)
	} else {
		c.metrics.duration.WithLabelValues(operation).Observe(time.Since(start).Seconds())
	}

	return newValue, err
}

func (c *MemcachedClient) Touch(ctx context.Context, key string, ttl time.Duration) error {
	start := time.Now()
	c.metrics.operations.WithLabelValues(opTouch).Inc()

	var err error
	select {
	case <-ctx.Done():
		err = ctx.Err()
	default:
		err = c.client.Touch(key, int32(ttl.Seconds()))
	}
	if err != nil {
		c.trackError(
			opTouch, err,
			"msg", "failed to touch cache item",
			"key", key,
		)
	} else {
		c.metrics.duration.WithLabelValues(opTouch).Observe(time.Since(start).Seconds())
	}
	return err
}

func (c *MemcachedClient) CompareAndSwap(ctx context.Context, key string, value []byte, ttl time.Duration) error {
	return c.storeOperation(ctx, key, value, ttl, opCompareAndSwap, func(ctx context.Context, key string, value []byte, ttl time.Duration) error {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
			ttlSeconds, ok := toSeconds(ttl)
			if !ok {
				return fmt.Errorf("%w: for compare-and-swap operation on %s %s", ErrInvalidTTL, key, ttl)
			}

			return c.client.CompareAndSwap(&memcache.Item{
				Key:        key,
				Value:      value,
				Expiration: ttlSeconds,
			})
		}
	})
}

func (c *MemcachedClient) FlushAll(ctx context.Context) error {
	var err error
	start := time.Now()
	c.metrics.operations.WithLabelValues(opFlush).Inc()

	select {
	case <-ctx.Done():
		err = ctx.Err()
	default:
		err = c.client.FlushAll()
	}
	if err != nil {
		c.trackError(
			opFlush, err,
			"msg", "failed to flush all cache",
		)
	} else {
		c.metrics.duration.WithLabelValues(opFlush).Observe(time.Since(start).Seconds())
	}

	return err
}

func (c *MemcachedClient) storeOperation(ctx context.Context, key string, value []byte, ttl time.Duration, operation string, f func(ctx context.Context, key string, value []byte, ttl time.Duration) error) error {
	if !c.itemSizeIsAcceptable(key, value, operation) {
		return nil
	}

	start := time.Now()
	c.metrics.operations.WithLabelValues(operation).Inc()

	err := f(ctx, key, value, ttl)
	if err != nil {
		c.trackError(
			operation, err,
			"msg", "failed to store item to cache",
			"operation", operation,
			"key", key,
			"sizeBytes", len(value),
		)
	}

	c.metrics.dataSize.WithLabelValues(operation).Observe(float64(len(value)))
	c.metrics.duration.WithLabelValues(operation).Observe(time.Since(start).Seconds())
	return err
}

// wait submits an async task and blocks until it completes. This can be used during
// tests to ensure that async "sets" have completed before attempting to read them.
func (c *MemcachedClient) wait() error {
	var wg sync.WaitGroup

	wg.Add(1)
	err := c.queue.submit(func() {
		wg.Done()
	})
	if err != nil {
		return err
	}

	wg.Wait()
	return nil
}

func (c *MemcachedClient) trackError(op string, err error, msg ...interface{}) {
	severity := level.DebugValue()

	var connErr *memcache.ConnectTimeoutError
	var netErr net.Error
	switch {
	case errors.As(err, &connErr):
		c.metrics.failures.WithLabelValues(op, reasonConnectTimeout).Inc()
	case errors.As(err, &netErr):
		if netErr.Timeout() {
			c.metrics.failures.WithLabelValues(op, reasonTimeout).Inc()
		} else {
			c.metrics.failures.WithLabelValues(op, reasonNetworkError).Inc()
		}
	case errors.Is(err, ErrNotStored):
		c.metrics.failures.WithLabelValues(op, reasonNotStored).Inc()
	case errors.Is(err, ErrInvalidTTL):
		c.metrics.failures.WithLabelValues(op, reasonInvalidTTL).Inc()
	case errors.Is(err, memcache.ErrMalformedKey):
		c.metrics.failures.WithLabelValues(op, reasonMalformedKey).Inc()
	case errors.Is(err, memcache.ErrServerError):
		c.metrics.failures.WithLabelValues(op, reasonServerError).Inc()
	case errors.Is(err, context.Canceled):
		c.metrics.failures.WithLabelValues(op, reasonCanceled).Inc()
	default:
		c.metrics.failures.WithLabelValues(op, reasonOther).Inc()
		severity = level.WarnValue() // Log unexpected kinds of errors with higher severity so they're easier to diagnose.
	}

	logger := log.WithPrefix(c.logger, level.Key(), severity)
	logger = log.WithSuffix(logger, "err", err)
	logger.Log(msg...)
}

func (c *MemcachedClient) getMultiBatched(ctx context.Context, keys []string, opts ...memcache.Option) ([]map[string]*memcache.Item, error) {
	// Do not batch if the input keys are less than the max batch size.
	if (c.config.MaxGetMultiBatchSize <= 0) || (len(keys) <= c.config.MaxGetMultiBatchSize) {
		// Even if we're not splitting the input into batches, make sure that our single request
		// still counts against the concurrency limit.
		if c.config.MaxGetMultiConcurrency > 0 {
			if err := c.getMultiGate.Start(ctx); err != nil {
				return nil, errors.Wrapf(err, "failed to wait for turn. Instance: %s", c.name)
			}

			defer c.getMultiGate.Done()
		}

		items, err := c.getMultiSingle(ctx, keys, opts...)
		if err != nil {
			return nil, err
		}

		return []map[string]*memcache.Item{items}, nil
	}

	// Calculate the number of expected results.
	batchSize := c.config.MaxGetMultiBatchSize
	numResults := len(keys) / batchSize
	if len(keys)%batchSize != 0 {
		numResults++
	}

	// If max concurrency is disabled, use a nil gate for the doWithBatch method which will
	// not apply any limit to the number goroutines started to make batch requests in that case.
	var getMultiGate gate.Gate
	if c.config.MaxGetMultiConcurrency > 0 {
		getMultiGate = c.getMultiGate
	}

	// Sort keys based on which memcached server they will be sharded to. Sorting keys that
	// are on the same server together before splitting into batches reduces the number of
	// connections required and increases the number of "gets" per connection.
	sortedKeys := c.sortKeysByServer(keys)

	// Allocate a channel to store results for each batch request. The max concurrency will be
	// enforced by doWithBatch.
	results := make(chan *memcachedGetMultiResult, numResults)
	defer close(results)

	// Ignore the error here since it can only be returned by our provided function which
	// always returns nil. NOTE also we are using a background context here for the doWithBatch
	// method. This is to ensure that it runs the expected number of batches _even if_ our
	// context (`ctx`) is canceled since we expect a certain number of batches to be read
	// from `results` below. The wrapped `getMultiSingle` method will still check our context
	// and short-circuit if it has been canceled.
	_ = doWithBatch(context.Background(), len(keys), c.config.MaxGetMultiBatchSize, getMultiGate, func(startIndex, endIndex int) error {
		batchKeys := sortedKeys[startIndex:endIndex]

		res := &memcachedGetMultiResult{}
		res.items, res.err = c.getMultiSingle(ctx, batchKeys, opts...)

		results <- res
		return nil
	})

	// Wait for all batch results. In case of error, we keep
	// track of the last error occurred.
	items := make([]map[string]*memcache.Item, 0, numResults)
	var lastErr error

	for i := 0; i < numResults; i++ {
		result := <-results
		if result.err != nil {
			lastErr = result.err
			continue
		}

		items = append(items, result.items)
	}

	return items, lastErr
}

func (c *MemcachedClient) getMultiSingle(ctx context.Context, keys []string, opts ...memcache.Option) (items map[string]*memcache.Item, err error) {
	start := time.Now()
	c.metrics.operations.WithLabelValues(opGetMulti).Inc()

	items, err = c.client.GetMulti(ctx, keys, opts...)

	if err != nil {
		c.trackError(
			opGetMulti, err,
			"msg", "failed to get multiple items from memcached",
		)
	} else {
		var total int
		for _, it := range items {
			total += len(it.Value)
		}
		c.metrics.dataSize.WithLabelValues(opGetMulti).Observe(float64(total))
		c.metrics.duration.WithLabelValues(opGetMulti).Observe(time.Since(start).Seconds())
	}

	return items, err
}

// doWithBatch do func with batch and gate. batchSize==0 means one batch. gate==nil means no gate.
func doWithBatch(ctx context.Context, totalSize int, batchSize int, ga gate.Gate, f func(startIndex, endIndex int) error) error {
	if totalSize == 0 {
		return nil
	}
	if batchSize <= 0 {
		return f(0, totalSize)
	}
	g, ctx := errgroup.WithContext(ctx)
	for i := 0; i < totalSize; i += batchSize {
		j := i + batchSize
		if j > totalSize {
			j = totalSize
		}
		if ga != nil {
			if err := ga.Start(ctx); err != nil {
				return nil
			}
		}
		startIndex, endIndex := i, j
		g.Go(func() error {
			if ga != nil {
				defer ga.Done()
			}
			return f(startIndex, endIndex)
		})
	}
	return g.Wait()
}

// sortKeysByServer sorts cache keys within a slice based on which server they are
// sharded to using a memcache.ServerSelector instance. The keys are ordered so keys
// on the same server are next to each other. Any errors encountered determining which
// server a key should be on will result in returning keys unsorted (in the same order
// they were supplied in). Note that output is not guaranteed to be any particular order
// *except* that keys sharded to the same server will be together. The order of keys
// returned may change from call to call.
func (c *MemcachedClient) sortKeysByServer(keys []string) []string {
	bucketed := make(map[string][]string)

	for _, key := range keys {
		addr, err := c.selector.PickServer(key)
		// If we couldn't determine the correct server, return keys in existing order
		if err != nil {
			return keys
		}

		addrString := addr.String()
		bucketed[addrString] = append(bucketed[addrString], key)
	}

	// Record the number of servers we've bucketed keys into. Note that this is only
	// done in the "batching" code path because sorting keys by server doesn't help
	// when there's only a single batch being used. This is fine since we expect the
	// number of servers touched to be higher for requests with more keys.
	c.metrics.serversTouched.WithLabelValues(opGetMulti).Observe(float64(len(bucketed)))

	out := make([]string, 0, len(keys))
	for srv := range bucketed {
		out = append(out, bucketed[srv]...)
	}

	return out
}

func (c *MemcachedClient) resolveAddrsLoop() {
	ticker := time.NewTicker(dnsProviderUpdateInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			err := c.resolveAddrs()
			if err != nil {
				level.Warn(c.logger).Log("msg", "failed update memcached servers list", "err", err)
			}
		case <-c.stop:
			return
		}
	}
}

func (c *MemcachedClient) resolveAddrs() error {
	// Resolve configured addresses with a reasonable timeout.
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// If some dns resolution fails, log the error.
	if err := c.addressProvider.Resolve(ctx, c.config.Addresses); err != nil {
		level.Error(c.logger).Log("msg", "failed to resolve addresses for memcached", "addresses", strings.Join(c.config.Addresses, ","), "err", err)
	}
	// Return error in case no server address is resolved.
	servers := c.addressProvider.Addresses()
	if len(servers) == 0 {
		return fmt.Errorf("no memcached server addresses were resolved")
	}

	return c.selector.SetServers(servers...)
}
