github.com/grafana/pyroscope@v1.18.0/pkg/distributor/instance_count.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  
     3  package distributor
     4  
     5  import (
     6  	"time"
     7  
     8  	"github.com/grafana/dskit/ring"
     9  	"go.uber.org/atomic"
    10  )
    11  
    12  // healthyInstanceDelegate counts the number of healthy instances that are part of the ring
    13  // and stores the count to the provided atomic integer. Used here to count the number of
    14  // distributors in the ring to determine how to enforce rate limiting.
    15  type healthyInstanceDelegate struct {
    16  	count            *atomic.Uint32
    17  	heartbeatTimeout time.Duration
    18  	next             ring.BasicLifecyclerDelegate
    19  }
    20  
    21  func newHealthyInstanceDelegate(count *atomic.Uint32, heartbeatTimeout time.Duration, next ring.BasicLifecyclerDelegate) *healthyInstanceDelegate {
    22  	return &healthyInstanceDelegate{count: count, heartbeatTimeout: heartbeatTimeout, next: next}
    23  }
    24  
    25  // OnRingInstanceRegister implements the ring.BasicLifecyclerDelegate interface
    26  func (d *healthyInstanceDelegate) OnRingInstanceRegister(lifecycler *ring.BasicLifecycler, ringDesc ring.Desc, instanceExists bool, instanceID string, instanceDesc ring.InstanceDesc) (ring.InstanceState, ring.Tokens) {
    27  	return d.next.OnRingInstanceRegister(lifecycler, ringDesc, instanceExists, instanceID, instanceDesc)
    28  }
    29  
    30  // OnRingInstanceTokens implements the ring.BasicLifecyclerDelegate interface
    31  func (d *healthyInstanceDelegate) OnRingInstanceTokens(lifecycler *ring.BasicLifecycler, tokens ring.Tokens) {
    32  	d.next.OnRingInstanceTokens(lifecycler, tokens)
    33  }
    34  
    35  // OnRingInstanceStopping implements the ring.BasicLifecyclerDelegate interface
    36  func (d *healthyInstanceDelegate) OnRingInstanceStopping(lifecycler *ring.BasicLifecycler) {
    37  	d.next.OnRingInstanceStopping(lifecycler)
    38  }
    39  
    40  // OnRingInstanceHeartbeat implements the ring.BasicLifecyclerDelegate interface
    41  func (d *healthyInstanceDelegate) OnRingInstanceHeartbeat(lifecycler *ring.BasicLifecycler, ringDesc *ring.Desc, instanceDesc *ring.InstanceDesc) {
    42  	activeMembers := uint32(0)
    43  	now := time.Now()
    44  
    45  	for _, instance := range ringDesc.Ingesters {
    46  		if ring.ACTIVE == instance.State && instance.IsHeartbeatHealthy(d.heartbeatTimeout, now) {
    47  			activeMembers++
    48  		}
    49  	}
    50  
    51  	d.count.Store(activeMembers)
    52  	d.next.OnRingInstanceHeartbeat(lifecycler, ringDesc, instanceDesc)
    53  }