Function TestMinimizeSpreadTokenGenerator

pkg/ring/token_generator_test.go:78–145

(t *testing.T)

}

func TestMinimizeSpreadTokenGenerator(t *testing.T) {
	rindDesc := NewDesc()
	zones := []string{"zone1", "zone2", "zone3"}

	mTokenGenerator := newMockedTokenGenerator(len(zones))
	minimizeTokenGenerator := &MinimizeSpreadTokenGenerator{
		innerGenerator: mTokenGenerator,
	}

	dups := map[uint32]bool{}

	// First time we should generate tokens using the inner generator
	generateTokensForIngesters(t, rindDesc, "initial", zones, minimizeTokenGenerator, dups)
	require.Equal(t, mTokenGenerator.called, len(zones))

	// Should Generate tokens based on the ring state
	for i := range 50 {
		generateTokensForIngesters(t, rindDesc, fmt.Sprintf("minimize-%v", i), zones, minimizeTokenGenerator, dups)
		assertDistancePerIngester(t, rindDesc, 0.01)
	}
	require.Equal(t, mTokenGenerator.called, len(zones))

	// Test generating tokens for an ingester that already have tokens in the ring
	rg := NewRandomTokenGenerator()
	rTokens := rg.GenerateTokens(rindDesc, "partial", zones[0], 256, true)
	rindDesc.AddIngester("partial", "partial", zones[0], rTokens, ACTIVE, time.Now())
	nTokens := minimizeTokenGenerator.GenerateTokens(rindDesc, "partial", zones[0], 256, true)
	rindDesc.AddIngester("partial", "partial", zones[0], append(rTokens, nTokens...), ACTIVE, time.Now())
	assertDistancePerIngester(t, rindDesc, 0.01)

	mTokenGenerator.called = 0
	// Should fallback to random generator when more than 1 ingester does not have tokens and force flag is set
	rindDesc.AddIngester("pendingIngester-1", "pendingIngester-1", zones[0], []uint32{}, PENDING, time.Now())
	rindDesc.AddIngester("pendingIngester-2", "pendingIngester-2", zones[0], []uint32{}, PENDING, time.Now().Add(-10*time.Minute))
	tokens := minimizeTokenGenerator.GenerateTokens(rindDesc, "pendingIngester-1", zones[0], 512, true)
	require.Len(t, tokens, 512)
	require.Equal(t, mTokenGenerator.called, 1)

	// Should generate for the ingester with with the smaller registered
	tokens = minimizeTokenGenerator.GenerateTokens(rindDesc, "pendingIngester-2", zones[0], 512, false)
	require.Len(t, tokens, 512)
	require.Equal(t, mTokenGenerator.called, 1)

	// Should generate tokens on other AZs
	rindDesc.AddIngester("pendingIngester-1-az-2", "pendingIngester-1-az-2", zones[0], []uint32{}, PENDING, time.Now())
	tokens = minimizeTokenGenerator.GenerateTokens(rindDesc, "pendingIngester-1-az-2", zones[1], 512, false)
	require.Len(t, tokens, 512)
	require.Equal(t, mTokenGenerator.called, 1)
	// Lets remove this pending ingester to no interfere with the next tests
	rindDesc.RemoveIngester("pendingIngester-1-az-2")

	// Should generate tokens only for the ingesters with the smaller registered time when multiples
	// ingesters does not have tokens
	tokens = minimizeTokenGenerator.GenerateTokens(rindDesc, "pendingIngester-1", zones[0], 512, false)
	require.Len(t, tokens, 0)
	tokens = minimizeTokenGenerator.GenerateTokens(rindDesc, "pendingIngester-2", zones[0], 512, false)
	require.Len(t, tokens, 512)