github.com/grafana/pyroscope@v1.18.0/pkg/util/servicediscovery/ring.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  
     3  package servicediscovery
     4  
     5  import (
     6  	"context"
     7  	"sort"
     8  	"time"
     9  
    10  	"github.com/grafana/dskit/ring"
    11  	"github.com/grafana/dskit/services"
    12  	"github.com/pkg/errors"
    13  )
    14  
    15  var (
    16  	// Ring operation used to get healthy active instances in the ring.
    17  	activeRingOp = ring.NewOp([]ring.InstanceState{ring.ACTIVE}, nil)
    18  )
    19  
    20  type ringServiceDiscovery struct {
    21  	services.Service
    22  
    23  	ringClient         *ring.Ring
    24  	ringCheckPeriod    time.Duration
    25  	maxUsedInstances   int
    26  	subservicesWatcher *services.FailureWatcher
    27  	receiver           Notifications
    28  
    29  	// Keep track of the instances that have been discovered and notified so far.
    30  	notifiedByAddress map[string]Instance
    31  }
    32  
    33  func NewRing(ringClient *ring.Ring, ringCheckPeriod time.Duration, maxUsedInstances int, receiver Notifications) services.Service {
    34  	r := &ringServiceDiscovery{
    35  		ringClient:         ringClient,
    36  		ringCheckPeriod:    ringCheckPeriod,
    37  		maxUsedInstances:   maxUsedInstances,
    38  		subservicesWatcher: services.NewFailureWatcher(),
    39  		notifiedByAddress:  make(map[string]Instance),
    40  		receiver:           receiver,
    41  	}
    42  
    43  	r.Service = services.NewBasicService(r.starting, r.running, r.stopping)
    44  	return r
    45  }
    46  
    47  func (r *ringServiceDiscovery) starting(ctx context.Context) error {
    48  	r.subservicesWatcher.WatchService(r.ringClient)
    49  
    50  	return errors.Wrap(services.StartAndAwaitRunning(ctx, r.ringClient), "failed to start ring client")
    51  }
    52  
    53  func (r *ringServiceDiscovery) stopping(_ error) error {
    54  	return errors.Wrap(services.StopAndAwaitTerminated(context.Background(), r.ringClient), "failed to stop ring client")
    55  }
    56  
    57  func (r *ringServiceDiscovery) running(ctx context.Context) error {
    58  	ringTicker := time.NewTicker(r.ringCheckPeriod)
    59  	defer ringTicker.Stop()
    60  
    61  	// Notifies the initial state.
    62  	all, _ := r.ringClient.GetAllHealthy(activeRingOp) // nolint:errcheck
    63  	r.notifyChanges(all)
    64  
    65  	for {
    66  		select {
    67  		case <-ringTicker.C:
    68  			all, _ := r.ringClient.GetAllHealthy(activeRingOp) // nolint:errcheck
    69  			r.notifyChanges(all)
    70  		case <-ctx.Done():
    71  			return nil
    72  		case err := <-r.subservicesWatcher.Chan():
    73  			return errors.Wrap(err, "a subservice of ring-based service discovery has failed")
    74  		}
    75  	}
    76  }
    77  
    78  // notifyChanges is not concurrency safe. The input all and inUse ring.ReplicationSet may be the same object.
    79  func (r *ringServiceDiscovery) notifyChanges(all ring.ReplicationSet) {
    80  	// Build a map with the discovered instances.
    81  	instancesByAddress := make(map[string]Instance, len(all.Instances))
    82  	for _, instance := range selectInUseInstances(all.Instances, r.maxUsedInstances) {
    83  		instancesByAddress[instance.Addr] = Instance{Address: instance.Addr, InUse: true}
    84  	}
    85  	for _, instance := range all.Instances {
    86  		if _, ok := instancesByAddress[instance.Addr]; !ok {
    87  			instancesByAddress[instance.Addr] = Instance{Address: instance.Addr, InUse: false}
    88  		}
    89  	}
    90  
    91  	// Notify new instances.
    92  	for addr, instance := range instancesByAddress {
    93  		if _, ok := r.notifiedByAddress[addr]; !ok {
    94  			r.receiver.InstanceAdded(instance)
    95  		}
    96  	}
    97  
    98  	// Notify changed instances.
    99  	for addr, instance := range instancesByAddress {
   100  		if n, ok := r.notifiedByAddress[addr]; ok && !n.Equal(instance) {
   101  			r.receiver.InstanceChanged(instance)
   102  		}
   103  	}
   104  
   105  	// Notify removed instances.
   106  	for addr, instance := range r.notifiedByAddress {
   107  		if _, ok := instancesByAddress[addr]; !ok {
   108  			r.receiver.InstanceRemoved(instance)
   109  		}
   110  	}
   111  
   112  	r.notifiedByAddress = instancesByAddress
   113  }
   114  
   115  func selectInUseInstances(instances []ring.InstanceDesc, maxInstances int) []ring.InstanceDesc {
   116  	if maxInstances <= 0 || len(instances) <= maxInstances {
   117  		return instances
   118  	}
   119  
   120  	// Select the first N instances (sorted by address) to be used.
   121  	sort.Sort(ring.ByAddr(instances))
   122  	return instances[:maxInstances]
   123  }