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

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  
     3  package distributor
     4  
     5  import (
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/grafana/dskit/ring"
    10  	"github.com/stretchr/testify/assert"
    11  	"go.uber.org/atomic"
    12  )
    13  
    14  type nopDelegate struct{}
    15  
    16  func (n nopDelegate) OnRingInstanceRegister(lifecycler *ring.BasicLifecycler, ringDesc ring.Desc, instanceExists bool, instanceID string, instanceDesc ring.InstanceDesc) (ring.InstanceState, ring.Tokens) {
    17  	return instanceDesc.State, instanceDesc.GetTokens()
    18  }
    19  
    20  func (n nopDelegate) OnRingInstanceTokens(lifecycler *ring.BasicLifecycler, tokens ring.Tokens) {
    21  }
    22  
    23  func (n nopDelegate) OnRingInstanceStopping(lifecycler *ring.BasicLifecycler) {
    24  }
    25  
    26  func (n nopDelegate) OnRingInstanceHeartbeat(lifecycler *ring.BasicLifecycler, ringDesc *ring.Desc, instanceDesc *ring.InstanceDesc) {
    27  }
    28  
    29  func TestHealthyInstanceDelegate_OnRingInstanceHeartbeat(t *testing.T) {
    30  	// addInstance registers a new instance with the given ring and sets its last heartbeat timestamp
    31  	addInstance := func(desc *ring.Desc, id string, state ring.InstanceState, timestamp int64) {
    32  		instance := desc.AddIngester(id, "127.0.0.1", "", []uint32{1}, state, time.Now(), false, time.Now())
    33  		instance.Timestamp = timestamp
    34  		desc.Ingesters[id] = instance
    35  	}
    36  
    37  	tests := map[string]struct {
    38  		ringSetup        func(desc *ring.Desc)
    39  		heartbeatTimeout time.Duration
    40  		expectedCount    uint32
    41  	}{
    42  		"all instances healthy and active": {
    43  			ringSetup: func(desc *ring.Desc) {
    44  				now := time.Now()
    45  				addInstance(desc, "distributor-1", ring.ACTIVE, now.Unix())
    46  				addInstance(desc, "distributor-2", ring.ACTIVE, now.Unix())
    47  				addInstance(desc, "distributor-3", ring.ACTIVE, now.Unix())
    48  			},
    49  			heartbeatTimeout: time.Minute,
    50  			expectedCount:    3,
    51  		},
    52  
    53  		"all instances healthy not all instances active": {
    54  			ringSetup: func(desc *ring.Desc) {
    55  				now := time.Now()
    56  				addInstance(desc, "distributor-1", ring.ACTIVE, now.Unix())
    57  				addInstance(desc, "distributor-2", ring.LEAVING, now.Unix())
    58  				addInstance(desc, "distributor-3", ring.ACTIVE, now.Unix())
    59  			},
    60  			heartbeatTimeout: time.Minute,
    61  			expectedCount:    2,
    62  		},
    63  
    64  		"some instances healthy all instances active": {
    65  			ringSetup: func(desc *ring.Desc) {
    66  				now := time.Now()
    67  				addInstance(desc, "distributor-1", ring.ACTIVE, now.Unix())
    68  				addInstance(desc, "distributor-2", ring.ACTIVE, now.Unix())
    69  				addInstance(desc, "distributor-3", ring.ACTIVE, now.Add(-5*time.Minute).Unix())
    70  			},
    71  			heartbeatTimeout: time.Minute,
    72  			expectedCount:    2,
    73  		},
    74  
    75  		"some instances healthy but timeout disabled all instances active": {
    76  			ringSetup: func(desc *ring.Desc) {
    77  				now := time.Now()
    78  				addInstance(desc, "distributor-1", ring.ACTIVE, now.Unix())
    79  				addInstance(desc, "distributor-2", ring.ACTIVE, now.Unix())
    80  				addInstance(desc, "distributor-3", ring.ACTIVE, now.Add(-5*time.Minute).Unix())
    81  			},
    82  			heartbeatTimeout: 0,
    83  			expectedCount:    3,
    84  		},
    85  	}
    86  
    87  	for testName, testData := range tests {
    88  		t.Run(testName, func(t *testing.T) {
    89  			count := atomic.NewUint32(0)
    90  			ringDesc := ring.NewDesc()
    91  
    92  			testData.ringSetup(ringDesc)
    93  			instance := ringDesc.Ingesters["distributor-1"]
    94  
    95  			delegate := newHealthyInstanceDelegate(count, testData.heartbeatTimeout, &nopDelegate{})
    96  			delegate.OnRingInstanceHeartbeat(&ring.BasicLifecycler{}, ringDesc, &instance)
    97  
    98  			assert.Equal(t, testData.expectedCount, count.Load())
    99  		})
   100  	}
   101  }