traces/vendor/github.com/grafana/dskit/ring/replication_set.go

package ring

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"sync"
	"time"

	kitlog "github.com/go-kit/log"
	"github.com/go-kit/log/level"

	"github.com/grafana/dskit/cancellation"
	"github.com/grafana/dskit/spanlogger"
)

// Used to communicate, via a context, the number of replicas that are available for a DoUntilQuorum call.
type availableReplicasContextKey struct{}

// ReplicationSet describes the instances to talk to for a given key, and how
// many errors to tolerate.
type ReplicationSet struct {
	Instances []InstanceDesc

	// Maximum number of tolerated failing instances. Max errors and max unavailable zones are
	// mutually exclusive.
	MaxErrors int

	// Maximum number of different zones in which instances can fail. Max unavailable zones and
	// max errors are mutually exclusive.
	MaxUnavailableZones int

	ZoneAwarenessEnabled bool
}

// Do function f in parallel for all replicas in the set, erroring if we exceed
// MaxErrors and returning early otherwise.
// Return a slice of all results from f, or nil if an error occurred.
func (r ReplicationSet) Do(ctx context.Context, delay time.Duration, f func(context.Context, *InstanceDesc) (interface{}, error)) ([]interface{}, error) {
	// Initialise the result tracker, which is use to keep track of successes and failures.
	var tracker replicationSetResultTracker
	if r.MaxUnavailableZones > 0 {
		tracker = newZoneAwareResultTracker(r.Instances, r.MaxUnavailableZones, nil, kitlog.NewNopLogger())
	} else {
		tracker = newDefaultResultTracker(r.Instances, r.MaxErrors, kitlog.NewNopLogger())
	}

	var (
		ch         = make(chan instanceResult[any], len(r.Instances))
		forceStart = make(chan struct{}, r.MaxErrors)
	)
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	// Spawn a goroutine for each instance.
	for i := range r.Instances {
		go func(i int, ing *InstanceDesc) {
			// Wait to send extra requests. Works only when zone-awareness is disabled.
			if delay > 0 && r.MaxUnavailableZones == 0 && i >= len(r.Instances)-r.MaxErrors {
				after := time.NewTimer(delay)
				defer after.Stop()
				select {
				case <-ctx.Done():
					return
				case <-forceStart:
				case <-after.C:
				}
			}
			result, err := f(ctx, ing)
			ch <- instanceResult[any]{
				result:   result,
				err:      err,
				instance: ing,
			}
		}(i, &r.Instances[i])
	}

	results := make([]interface{}, 0, len(r.Instances))

	for !tracker.succeeded() {
		select {
		case res := <-ch:
			tracker.done(res.instance, res.err)
			if res.err != nil {
				if tracker.failed() {
					return nil, res.err
				}

				// force one of the delayed requests to start
				if delay > 0 && r.MaxUnavailableZones == 0 {
					forceStart <- struct{}{}
				}
			} else {
				results = append(results, res.result)
			}

		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}

	return results, nil
}

type DoUntilQuorumConfig struct {
	// MinimizeRequests enables request minimization.
	// See docs for DoUntilQuorum for more information.
	MinimizeRequests bool

	// HedgingDelay configures the delay used before initiating hedged requests.
	// Hedging is only enabled if HedgingDelay is non-zero and MinimizeRequests is true.
	// See docs for DoUntilQuorum for more information.
	HedgingDelay time.Duration

	// Logger to emit log lines and span events to during the call.
	// Can be nil, in which case no log lines or span events are emitted.
	Logger *spanlogger.SpanLogger

	// IsTerminalError allows DoUntilQuorum to detect terminal errors generated by requests.
	//
	// If IsTerminalError is non-nil and a request returns an error for which IsTerminalError returns true,
	// DoUntilQuorum will immediately cancel any inflight requests and return the error.
	//
	// This is useful to cancel DoUntilQuorum when an unrecoverable error occurs and it does not
	// make sense to attempt requests to other instances. For example, if a client-side limit on the
	// total response size across all instances is reached, making further requests to other
	// instances would not be worthwhile.
	IsTerminalError func(error) bool

	// ZoneSorter orders the provided zones in preference order, for use when MinimizeRequests is true
	// and DoUntilQuorum is operating in zone-aware mode. If not set, zones will be used in a
	// randomly-selected order.
	//
	// Earlier zones will be used first.
	// The function can modify the provided slice of zones in place.
	// All provided zones must be returned exactly once.
	//
	// This can be used to prioritise zones that are more likely to succeed, or are expected to complete
	// faster, for example.
	ZoneSorter ZoneSorter

	// IncludeReplicaCount indicates that we should unclude in a DoUntilQuorum call, as part of the context, a count of the
	// number of available replicas. This can be propagated to a replica and used to make informed decisions about load
	// shedding.
	IncludeReplicaCount bool
}

func (c DoUntilQuorumConfig) Validate() error {
	if c.HedgingDelay < 0 {
		return errors.New("invalid DoUntilQuorumConfig: HedgingDelay must be non-negative")
	}

	return nil
}

// DoUntilQuorum runs function f in parallel for all replicas in r.
//
// # Result selection
//
// If r.MaxUnavailableZones is greater than zero, or r.ZoneAwarenessEnabled is true, DoUntilQuorum operates in zone-aware mode:
//   - DoUntilQuorum returns an error if calls to f for instances in more than r.MaxUnavailableZones zones return errors
//   - Otherwise, DoUntilQuorum returns all results from all replicas in the first zones for which f succeeds
//     for every instance in that zone (eg. if there are 3 zones and r.MaxUnavailableZones is 1, DoUntilQuorum will
//     return the results from all instances in 2 zones, even if all calls to f succeed).
//
// Otherwise, DoUntilQuorum operates in non-zone-aware mode:
//   - DoUntilQuorum returns an error if more than r.MaxErrors calls to f return errors
//   - Otherwise, DoUntilQuorum returns all results from the first len(r.Instances) - r.MaxErrors instances
//     (eg. if there are 6 replicas and r.MaxErrors is 2, DoUntilQuorum will return the results from the first 4
//     successful calls to f, even if all 6 calls to f succeed).
//
// # Request minimization
//
// cfg.MinimizeRequests enables or disables request minimization.
//
// Regardless of the value of cfg.MinimizeRequests, if one of the termination conditions above is satisfied or ctx is
// cancelled before f is called for an instance, f may not be called for that instance at all.
//
// ## When disabled
//
// If request minimization is disabled, DoUntilQuorum will call f for each instance in r. The value of cfg.HedgingDelay
// is ignored.
//
// ## When enabled
//
// If request minimization is enabled, DoUntilQuorum will initially call f for the minimum number of instances needed to
// reach the termination conditions above, and later call f for further instances if required. For example, if
// r.MaxUnavailableZones is 1 and there are three zones, DoUntilQuorum will initially only call f for instances in two
// zones, and only call f for instances in the remaining zone if a request in the initial two zones fails.
//
// If cfg.ZoneSorter is non-nil and DoUntilQuorum is operating in zone-aware mode, DoUntilQuorum will initiate requests
// to zones in the order returned by the sorter.
//
// If cfg.ZoneSorter is nil, or DoUntilQuorum is operating in non-zone-aware mode, DoUntilQuorum will randomly select
// available zones / instances such that calling DoUntilQuorum multiple times with the same ReplicationSet should evenly
// distribute requests across all zones / instances.
//
// If cfg.HedgingDelay is non-zero, DoUntilQuorum will call f for an additional zone's instances (if zone-aware) / an
// additional instance (if not zone-aware) every cfg.HedgingDelay until one of the termination conditions above is
// reached. For example, if r.MaxUnavailableZones is 2, cfg.HedgingDelay is 4 seconds and there are fives zones,
// DoUntilQuorum will initially only call f for instances in three zones, and unless one of the termination conditions
// is reached earlier, will then call f for instances in a fourth zone approximately 4 seconds later, and then call f
// for instances in the final zone approximately 4 seconds after that (ie. roughly 8 seconds since the call to
// DoUntilQuorum began).
//
// # Cleanup
//
// Any results from successful calls to f that are not returned by DoUntilQuorum will be passed to cleanupFunc,
// including when DoUntilQuorum returns an error or only returns a subset of successful results. cleanupFunc may
// be called both before and after DoUntilQuorum returns.
//
// A call to f is considered successful if it returns a nil error.
//
// # Contexts
//
// The context.Context passed to an invocation of f may be cancelled at any time if the result of that invocation of
// f will not be used.
//
// DoUntilQuorum cancels the context.Context passed to each invocation of f before DoUntilQuorum returns.
func DoUntilQuorum[T any](ctx context.Context, r ReplicationSet, cfg DoUntilQuorumConfig, f func(context.Context, *InstanceDesc) (T, error), cleanupFunc func(T)) ([]T, error) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	wrappedF := func(ctx context.Context, desc *InstanceDesc, _ context.CancelCauseFunc) (T, error) {
		return f(ctx, desc)
	}

	return DoUntilQuorumWithoutSuccessfulContextCancellation(ctx, r, cfg, wrappedF, cleanupFunc)
}

// DoUntilQuorumWithoutSuccessfulContextCancellation behaves the same as DoUntilQuorum, except it does not cancel
// the context.Context passed to invocations of f whose results are returned.
//
// For example, this is useful in situations where DoUntilQuorumWithoutSuccessfulContextCancellation is used
// to establish a set of streams that will be used after DoUntilQuorumWithoutSuccessfulContextCancellation returns.
//
// It is the caller's responsibility to ensure that either of the following are eventually true:
//   - ctx is cancelled, or
//   - the corresponding context.CancelFunc is called for all invocations of f whose results are returned by
//     DoUntilQuorumWithoutSuccessfulContextCancellation
//
// Failing to do this may result in a memory leak.
func DoUntilQuorumWithoutSuccessfulContextCancellation[T any](ctx context.Context, r ReplicationSet, cfg DoUntilQuorumConfig, f func(context.Context, *InstanceDesc, context.CancelCauseFunc) (T, error), cleanupFunc func(T)) ([]T, error) {
	if err := cfg.Validate(); err != nil {
		return nil, err
	}

	if r.ZoneAwarenessEnabled && r.MaxErrors > 0 {
		return nil, fmt.Errorf("invalid ReplicationSet: MaxErrors is non-zero (is %v) and ZoneAwarenessEnabled is true", r.MaxErrors)
	}

	var logger kitlog.Logger = cfg.Logger
	if cfg.Logger == nil {
		logger = kitlog.NewNopLogger()
	}

	if cfg.IncludeReplicaCount {
		ctx = ContextWithAvailableReplicas(ctx, len(r.Instances))
	}

	resultsChan := make(chan instanceResult[T], len(r.Instances))
	resultsRemaining := len(r.Instances)

	defer func() {
		go func() {
			for resultsRemaining > 0 {
				result := <-resultsChan
				resultsRemaining--

				if result.err == nil {
					cleanupFunc(result.result)
				}
			}
		}()
	}()

	var resultTracker replicationSetResultTracker
	var contextTracker replicationSetContextTracker
	if r.MaxUnavailableZones > 0 || r.ZoneAwarenessEnabled {
		resultTracker = newZoneAwareResultTracker(r.Instances, r.MaxUnavailableZones, cfg.ZoneSorter, logger)
		contextTracker = newZoneAwareContextTracker(ctx, r.Instances)
	} else {
		resultTracker = newDefaultResultTracker(r.Instances, r.MaxErrors, logger)
		contextTracker = newDefaultContextTracker(ctx, r.Instances)
	}

	if cfg.MinimizeRequests {
		resultTracker.startMinimumRequests()
	} else {
		resultTracker.startAllRequests()
	}

	for i := range r.Instances {
		instance := &r.Instances[i]
		ctx, ctxCancel := contextTracker.contextFor(instance)

		go func(desc *InstanceDesc) {
			if err := resultTracker.awaitStart(ctx, desc); err != nil {
				// Post to resultsChan so that the deferred cleanup handler above eventually terminates.
				resultsChan <- instanceResult[T]{
					err:      err,
					instance: desc,
				}

				return
			}

			result, err := f(ctx, desc, ctxCancel)
			resultsChan <- instanceResult[T]{
				result:   result,
				err:      err,
				instance: desc,
			}
		}(instance)
	}

	resultsMap := make(map[*InstanceDesc]T, len(r.Instances))
	cleanupResultsAlreadyReceived := func() {
		for _, result := range resultsMap {
			cleanupFunc(result)
		}
	}

	terminate := func(err error, cause string) ([]T, error) {
		if cfg.Logger != nil && !errors.Is(err, context.Canceled) { // Cancellation is not an error.
			cfg.Logger.SetError()
		}

		contextTracker.cancelAllContexts(cancellation.NewError(errors.New(cause)))
		cleanupResultsAlreadyReceived()
		return nil, err
	}

	var hedgingTrigger <-chan time.Time

	if cfg.HedgingDelay > 0 {
		ticker := time.NewTicker(cfg.HedgingDelay)
		defer ticker.Stop()
		hedgingTrigger = ticker.C
	}

	for !resultTracker.succeeded() {
		select {
		case <-ctx.Done():
			err := context.Cause(ctx)
			level.Debug(logger).Log("msg", "parent context done, returning", "err", err)

			// No need to cancel individual instance contexts, as they inherit the cancellation from ctx.
			cleanupResultsAlreadyReceived()

			return nil, err
		case <-hedgingTrigger:
			resultTracker.startAdditionalRequests()
		case result := <-resultsChan:
			resultsRemaining--

			if result.err != nil && cfg.IsTerminalError != nil && cfg.IsTerminalError(result.err) {
				level.Warn(logger).Log("msg", "cancelling all outstanding requests because a terminal error occurred", "err", result.err)
				// We must return before calling resultTracker.done() below, otherwise done() might start further requests if request minimisation is enabled.
				return terminate(result.err, "a terminal error occurred")
			}

			resultTracker.done(result.instance, result.err)

			if result.err == nil {
				resultsMap[result.instance] = result.result
			} else {
				contextTracker.cancelContextFor(result.instance, cancellation.NewErrorf("this instance returned an error: %w", result.err))

				if resultTracker.failed() {
					level.Error(logger).Log("msg", "cancelling all outstanding requests because quorum cannot be reached")
					return terminate(result.err, "quorum cannot be reached")
				}
			}
		}
	}

	level.Debug(logger).Log("msg", "quorum reached")

	results := make([]T, 0, len(r.Instances))

	for i := range r.Instances {
		instance := &r.Instances[i]
		result, haveResult := resultsMap[instance]

		if haveResult {
			if resultTracker.shouldIncludeResultFrom(instance) {
				results = append(results, result)
			} else {
				contextTracker.cancelContextFor(instance, cancellation.NewErrorf("quorum reached, result not required from this instance"))
				cleanupFunc(result)
			}
		} else {
			// Nothing to clean up (yet) - this will be handled by deferred call above.
			contextTracker.cancelContextFor(instance, cancellation.NewErrorf("quorum reached, result not required from this instance"))
		}
	}

	return results, nil
}

// DoMultiUntilQuorumWithoutSuccessfulContextCancellation behaves similar to DoUntilQuorumWithoutSuccessfulContextCancellation
// with the following exceptions:
//
//   - This function calls DoUntilQuorumWithoutSuccessfulContextCancellation for each input ReplicationSet and requires
//     DoUntilQuorumWithoutSuccessfulContextCancellation to successfully run for each of them. Execution breaks on the
//     first error returned by DoUntilQuorumWithoutSuccessfulContextCancellation on any ReplicationSet.
//
//   - This function requires that the callback function f always call context.CancelCauseFunc once done. Failing to
//     cancel the context will leak resources.
func DoMultiUntilQuorumWithoutSuccessfulContextCancellation[T any](ctx context.Context, sets []ReplicationSet, cfg DoUntilQuorumConfig, f func(context.Context, *InstanceDesc, context.CancelCauseFunc) (T, error), cleanupFunc func(T)) ([]T, error) {
	if len(sets) == 0 {
		return nil, errors.New("no replication sets")
	}
	if len(sets) == 1 {
		return DoUntilQuorumWithoutSuccessfulContextCancellation[T](ctx, sets[0], cfg, f, cleanupFunc)
	}

	results, _, err := doMultiUntilQuorumWithoutSuccessfulContextCancellation[T](ctx, sets, cfg, f, cleanupFunc)
	return results, err
}

// See DoMultiUntilQuorumWithoutSuccessfulContextCancellation().
//
// The returned context.Context is the internal context used by workers and it's used for testing purposes.
func doMultiUntilQuorumWithoutSuccessfulContextCancellation[T any](ctx context.Context, sets []ReplicationSet, cfg DoUntilQuorumConfig, f func(context.Context, *InstanceDesc, context.CancelCauseFunc) (T, error), cleanupFunc func(T)) ([]T, context.Context, error) {
	var (
		returnResultsMx = sync.Mutex{}
		returnResults   = make([]T, 0, len(sets)*len(sets[0].Instances)) // Assume all replication sets have the same number of instances.

		returnErrOnce sync.Once
		returnErr     error // The first error occurred.

		workersGroup                 = sync.WaitGroup{}
		workersCtx, cancelWorkersCtx = context.WithCancelCause(ctx)

		inflightTracker = newInflightInstanceTracker(sets)
	)

	cancelWorkersCtxIfSafe := func() {
		if inflightTracker.allInstancesCompleted() {
			cancelWorkersCtx(errors.New("all requests completed"))
		}
	}

	// Start a worker for each set. A worker is responsible to call DoUntilQuorumWithoutSuccessfulContextCancellation()
	// for the given replication set and handle the result.
	workersGroup.Add(len(sets))

	for idx, set := range sets {
		go func(idx int, set ReplicationSet) {
			defer workersGroup.Done()

			wrappedFn := func(ctx context.Context, instance *InstanceDesc, cancelCtx context.CancelCauseFunc) (T, error) {
				// The callback function has been called, so we need to track it.
				inflightTracker.addInstance(idx, instance)

				// Inject custom logic in the context.CancelCauseFunc.
				return f(ctx, instance, func(cause error) {
					// Call the original one.
					cancelCtx(cause)

					// The callback has done, so we can remove it from tracker and then check if it's safe
					// to cancel the workers context.
					inflightTracker.removeInstance(idx, instance)
					cancelWorkersCtxIfSafe()
				})
			}

			setResults, setErr := DoUntilQuorumWithoutSuccessfulContextCancellation[T](workersCtx, set, cfg, wrappedFn, cleanupFunc)

			if setErr != nil {
				returnErrOnce.Do(func() {
					returnErr = setErr

					// Interrupt the execution of all workers.
					cancelWorkersCtx(setErr)
				})

				return
			}

			// Keep track of the results.
			returnResultsMx.Lock()
			returnResults = append(returnResults, setResults...)
			returnResultsMx.Unlock()
		}(idx, set)
	}

	// Wait until all goroutines have terminated.
	workersGroup.Wait()

	// All workers completed, so it's guaranteed returnResults and returnErr won't be accessed by workers anymore,
	// and it's safe to read them with no locking.
	if returnErr != nil {
		return nil, workersCtx, returnErr
	}

	// No error occurred. It means workers context hasn't been canceled yet, and we don't expect more callbacks
	// to get tracked, so we can check if the cancelling condition has already been reached and eventually do it.
	inflightTracker.allInstancesAdded()
	cancelWorkersCtxIfSafe()

	return returnResults, workersCtx, nil
}

type instanceResult[T any] struct {
	result   T
	err      error
	instance *InstanceDesc
}

// Includes returns whether the replication set includes the replica with the provided addr.
func (r ReplicationSet) Includes(addr string) bool {
	for _, instance := range r.Instances {
		if instance.GetAddr() == addr {
			return true
		}
	}

	return false
}

// GetIDs returns the IDs of all instances within the replication set. Returned slice
// order is not guaranteed.
func (r ReplicationSet) GetIDs() []string {
	ids := make([]string, 0, len(r.Instances))
	for _, desc := range r.Instances {
		ids = append(ids, desc.Id)
	}
	return ids
}

// GetAddresses returns the addresses of all instances within the replication set. Returned slice
// order is not guaranteed.
func (r ReplicationSet) GetAddresses() []string {
	addrs := make([]string, 0, len(r.Instances))
	for _, desc := range r.Instances {
		addrs = append(addrs, desc.Addr)
	}
	return addrs
}

// GetAddressesWithout returns the addresses of all instances within the replication set while
// excluding the specified address. Returned slice order is not guaranteed.
func (r ReplicationSet) GetAddressesWithout(exclude string) []string {
	addrs := make([]string, 0, len(r.Instances))
	for _, desc := range r.Instances {
		if desc.Addr != exclude {
			addrs = append(addrs, desc.Addr)
		}
	}
	return addrs
}

// ZoneCount returns the number of unique zones represented by instances within this replication set.
func (r *ReplicationSet) ZoneCount() int {
	// Why not use a map here? Using a slice is faster for the small number of zones we expect to typically use.
	zones := []string{}

	for _, i := range r.Instances {
		sawZone := false

		for _, z := range zones {
			if z == i.Zone {
				sawZone = true
				break
			}
		}

		if !sawZone {
			zones = append(zones, i.Zone)
		}
	}

	return len(zones)
}

// HasReplicationSetChanged returns false if two replications sets are the same (with possibly different timestamps),
// true if they differ in any way (number of instances, instance states, tokens, zones, ...).
func HasReplicationSetChanged(before, after ReplicationSet) bool {
	return hasReplicationSetChangedExcluding(before, after, func(i *InstanceDesc) {
		i.Timestamp = 0
	})
}

// HasReplicationSetChangedWithoutState returns false if two replications sets
// are the same (with possibly different timestamps and instance states),
// true if they differ in any other way (number of instances, tokens, zones, ...).
func HasReplicationSetChangedWithoutState(before, after ReplicationSet) bool {
	return hasReplicationSetChangedExcluding(before, after, func(i *InstanceDesc) {
		i.Timestamp = 0
		i.State = PENDING
	})
}

// Has HasReplicationSetChangedWithoutStateOrAddr returns false if two replications sets
// are the same (with possibly different timestamps, instance states, and ip addresses),
// true if they differ in any other way (number of instances, tokens, zones, ...).
func HasReplicationSetChangedWithoutStateOrAddr(before, after ReplicationSet) bool {
	return hasReplicationSetChangedExcluding(before, after, func(i *InstanceDesc) {
		i.Timestamp = 0
		i.State = PENDING
		i.Addr = ""
	})
}

// Do comparison of replicasets, but apply a function first
// to be able to exclude (reset) some values
func hasReplicationSetChangedExcluding(before, after ReplicationSet, exclude func(*InstanceDesc)) bool {
	beforeInstances := before.Instances
	afterInstances := after.Instances

	if len(beforeInstances) != len(afterInstances) {
		return true
	}

	sort.Sort(ByID(beforeInstances))
	sort.Sort(ByID(afterInstances))

	for i := 0; i < len(beforeInstances); i++ {
		b := beforeInstances[i]
		a := afterInstances[i]

		exclude(&a)
		exclude(&b)

		if !b.Equal(a) {
			return true
		}
	}

	return false
}

// ContextWithAvailableReplicas returns a new context that can be used as part of a DoUntilQuorum call, with the number
// of availableReplicas set. This can be later checked via GetAvailableReplicas.
func ContextWithAvailableReplicas(ctx context.Context, availableReplicas int) context.Context {
	return context.WithValue(ctx, availableReplicasContextKey{}, availableReplicas)
}

// GetAvailableReplicas returns the available replicas count from the context, if any, along with an ok flag indicating
// if the availableReplicas was set.
func GetAvailableReplicas(ctx context.Context) (int, bool) {
	result, ok := ctx.Value(availableReplicasContextKey{}).(int)
	return result, ok
}