github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/ruler/base/lifecycle_test.go (about)

     1  package base
     2  
     3  import (
     4  	"context"
     5  	"sort"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/go-kit/log"
    10  	"github.com/grafana/dskit/kv"
    11  	"github.com/grafana/dskit/kv/consul"
    12  	"github.com/grafana/dskit/ring"
    13  	"github.com/grafana/dskit/services"
    14  	"github.com/stretchr/testify/assert"
    15  	"github.com/stretchr/testify/require"
    16  
    17  	"github.com/grafana/loki/pkg/storage"
    18  	"github.com/grafana/loki/pkg/util/test"
    19  )
    20  
    21  // TestRulerShutdown tests shutting down ruler unregisters correctly
    22  func TestRulerShutdown(t *testing.T) {
    23  	ctx := context.Background()
    24  
    25  	config := defaultRulerConfig(t, newMockRuleStore(mockRules))
    26  
    27  	m := storage.NewClientMetrics()
    28  	defer m.Unregister()
    29  	r := buildRuler(t, config, nil, m, nil)
    30  
    31  	r.cfg.EnableSharding = true
    32  	ringStore, closer := consul.NewInMemoryClient(ring.GetCodec(), log.NewNopLogger(), nil)
    33  	t.Cleanup(func() { assert.NoError(t, closer.Close()) })
    34  
    35  	err := enableSharding(r, ringStore)
    36  	require.NoError(t, err)
    37  
    38  	require.NoError(t, services.StartAndAwaitRunning(ctx, r))
    39  	defer services.StopAndAwaitTerminated(ctx, r) //nolint:errcheck
    40  
    41  	// Wait until the tokens are registered in the ring
    42  	test.Poll(t, 100*time.Millisecond, config.Ring.NumTokens, func() interface{} {
    43  		return numTokens(ringStore, "localhost", ringKey)
    44  	})
    45  
    46  	require.Equal(t, ring.ACTIVE, r.lifecycler.GetState())
    47  
    48  	require.NoError(t, services.StopAndAwaitTerminated(context.Background(), r))
    49  
    50  	// Wait until the tokens are unregistered from the ring
    51  	test.Poll(t, 100*time.Millisecond, 0, func() interface{} {
    52  		return numTokens(ringStore, "localhost", ringKey)
    53  	})
    54  }
    55  
    56  func TestRuler_RingLifecyclerShouldAutoForgetUnhealthyInstances(t *testing.T) {
    57  	const unhealthyInstanceID = "unhealthy-id"
    58  	const heartbeatTimeout = time.Minute
    59  
    60  	ctx := context.Background()
    61  	config := defaultRulerConfig(t, newMockRuleStore(mockRules))
    62  	m := storage.NewClientMetrics()
    63  	defer m.Unregister()
    64  	r := buildRuler(t, config, nil, m, nil)
    65  	r.cfg.EnableSharding = true
    66  	r.cfg.Ring.HeartbeatPeriod = 100 * time.Millisecond
    67  	r.cfg.Ring.HeartbeatTimeout = heartbeatTimeout
    68  
    69  	ringStore, closer := consul.NewInMemoryClient(ring.GetCodec(), log.NewNopLogger(), nil)
    70  	t.Cleanup(func() { assert.NoError(t, closer.Close()) })
    71  
    72  	err := enableSharding(r, ringStore)
    73  	require.NoError(t, err)
    74  
    75  	require.NoError(t, services.StartAndAwaitRunning(ctx, r))
    76  	defer services.StopAndAwaitTerminated(ctx, r) //nolint:errcheck
    77  
    78  	// Add an unhealthy instance to the ring.
    79  	require.NoError(t, ringStore.CAS(ctx, ringKey, func(in interface{}) (interface{}, bool, error) {
    80  		ringDesc := ring.GetOrCreateRingDesc(in)
    81  
    82  		instance := ringDesc.AddIngester(unhealthyInstanceID, "1.1.1.1", "", generateSortedTokens(config.Ring.NumTokens), ring.ACTIVE, time.Now())
    83  		instance.Timestamp = time.Now().Add(-(ringAutoForgetUnhealthyPeriods + 1) * heartbeatTimeout).Unix()
    84  		ringDesc.Ingesters[unhealthyInstanceID] = instance
    85  
    86  		return ringDesc, true, nil
    87  	}))
    88  
    89  	// Ensure the unhealthy instance is removed from the ring.
    90  	test.Poll(t, time.Second*5, false, func() interface{} {
    91  		d, err := ringStore.Get(ctx, ringKey)
    92  		if err != nil {
    93  			return err
    94  		}
    95  
    96  		_, ok := ring.GetOrCreateRingDesc(d).Ingesters[unhealthyInstanceID]
    97  		return ok
    98  	})
    99  }
   100  
   101  func generateSortedTokens(numTokens int) ring.Tokens {
   102  	tokens := ring.GenerateTokens(numTokens, nil)
   103  
   104  	// Ensure generated tokens are sorted.
   105  	sort.Slice(tokens, func(i, j int) bool {
   106  		return tokens[i] < tokens[j]
   107  	})
   108  
   109  	return ring.Tokens(tokens)
   110  }
   111  
   112  // numTokens determines the number of tokens owned by the specified
   113  // address
   114  func numTokens(c kv.Client, name, ringKey string) int {
   115  	ringDesc, err := c.Get(context.Background(), ringKey)
   116  
   117  	// The ringDesc may be null if the lifecycler hasn't stored the ring
   118  	// to the KVStore yet.
   119  	if ringDesc == nil || err != nil {
   120  		return 0
   121  	}
   122  	rd := ringDesc.(*ring.Desc)
   123  	return len(rd.Ingesters[name].Tokens)
   124  }