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

     1  package testhelper
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/grafana/dskit/ring"
     7  )
     8  
     9  type MockRing struct {
    10  	ingesters         []ring.InstanceDesc
    11  	replicationFactor uint32
    12  }
    13  
    14  func NewMockRing(ingesters []ring.InstanceDesc, replicationFactor uint32) MockRing {
    15  	return MockRing{
    16  		ingesters:         ingesters,
    17  		replicationFactor: replicationFactor,
    18  	}
    19  }
    20  
    21  func (r MockRing) Get(key uint32, op ring.Operation, buf []ring.InstanceDesc, _ []string, _ []string) (ring.ReplicationSet, error) {
    22  	result := ring.ReplicationSet{
    23  		MaxErrors: 1,
    24  		Instances: buf[:0],
    25  	}
    26  
    27  	for i := uint32(0); i < r.replicationFactor; i++ {
    28  		n := (key + i) % uint32(len(r.ingesters))
    29  		result.Instances = append(result.Instances, r.ingesters[n])
    30  	}
    31  	return result, nil
    32  }
    33  
    34  func (r MockRing) GetWithOptions(key uint32, op ring.Operation, opts ...ring.Option) (ring.ReplicationSet, error) {
    35  	return r.Get(key, op, nil, nil, nil)
    36  }
    37  
    38  func (r MockRing) GetAllHealthy(op ring.Operation) (ring.ReplicationSet, error) {
    39  	return r.GetReplicationSetForOperation(op)
    40  }
    41  
    42  func (r MockRing) GetReplicationSetForOperation(op ring.Operation) (ring.ReplicationSet, error) {
    43  	return ring.ReplicationSet{
    44  		Instances: r.ingesters,
    45  		MaxErrors: 1,
    46  	}, nil
    47  }
    48  
    49  func (r MockRing) GetSubringForOperationStates(op ring.Operation) ring.ReadRing {
    50  	return r
    51  }
    52  
    53  func (r MockRing) ReplicationFactor() int {
    54  	return int(r.replicationFactor)
    55  }
    56  
    57  func (r MockRing) InstancesCount() int {
    58  	return len(r.ingesters)
    59  }
    60  
    61  func (r MockRing) InstancesWithTokensCount() int {
    62  	return len(r.ingesters)
    63  }
    64  
    65  func (r MockRing) Subring(key uint32, n int) ring.ReadRing {
    66  	return r
    67  }
    68  
    69  func (r MockRing) HasInstance(instanceID string) bool {
    70  	for _, ing := range r.ingesters {
    71  		if ing.Addr != instanceID {
    72  			return true
    73  		}
    74  	}
    75  	return false
    76  }
    77  
    78  func (r MockRing) ShuffleShard(identifier string, size int) ring.ReadRing {
    79  	// Nothing to do if the shard size is not smaller then the actual ring.
    80  	if size <= 0 || r.InstancesCount() <= size {
    81  		return r
    82  	}
    83  
    84  	if rf := int(r.replicationFactor); size < rf {
    85  		size = rf
    86  	}
    87  
    88  	return &MockRing{
    89  		ingesters:         r.ingesters[:size],
    90  		replicationFactor: r.replicationFactor,
    91  	}
    92  }
    93  
    94  func (r MockRing) ShuffleShardWithLookback(identifier string, size int, lookbackPeriod time.Duration, now time.Time) ring.ReadRing {
    95  	return r
    96  }
    97  
    98  func (r MockRing) CleanupShuffleShardCache(identifier string) {}
    99  
   100  func (r MockRing) GetInstanceState(instanceID string) (ring.InstanceState, error) {
   101  	return 0, nil
   102  }
   103  
   104  func (r MockRing) GetTokenRangesForInstance(instanceID string) (ring.TokenRanges, error) {
   105  	return nil, nil
   106  }
   107  
   108  func (r MockRing) InstancesInZoneCount(zone string) int {
   109  	return len(r.ingesters)
   110  }
   111  
   112  func (r MockRing) InstancesWithTokensInZoneCount(zone string) int {
   113  	return len(r.ingesters)
   114  }
   115  
   116  func (r MockRing) WritableInstancesWithTokensCount() int {
   117  	return len(r.ingesters)
   118  }
   119  
   120  func (r MockRing) WritableInstancesWithTokensInZoneCount(zone string) int {
   121  	return len(r.ingesters)
   122  }
   123  
   124  func (r MockRing) ZonesCount() int {
   125  	return 1
   126  }
   127  
   128  func (r *MockRing) SetInstances(instances []ring.InstanceDesc) {
   129  	r.ingesters = instances
   130  }