Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Rhythm] Block builder test updates #4510

Merged
merged 5 commits into from
Jan 7, 2025
Merged

[Rhythm] Block builder test updates #4510

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
34e5ac3
Make blockbuilder tests closer to real kafka and less implementation …
mdisibio Jan 3, 2025
28e0d11
Verify last committed offset in each test
mdisibio Jan 3, 2025
f1a6c1e
hide test function
mdisibio Jan 3, 2025
0a6eb3d
lint
mdisibio Jan 3, 2025
b51df12
lint
mdisibio Jan 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 53 additions & 30 deletions modules/blockbuilder/blockbuilder_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,20 +30,23 @@ import (
"go.uber.org/atomic"
)

const testTopic = "test-topic"
const (
testTopic = "test-topic"
testConsumerGroup = "test-consumer-group"
testPartition = int32(0)
)

// When the partition starts with no existing commit,
// the block-builder looks back to consume all available records from the start and ensures they are committed and flushed into a block.
func TestBlockbuilder_lookbackOnNoCommit(t *testing.T) {
ctx, cancel := context.WithCancelCause(context.Background())
t.Cleanup(func() { cancel(errors.New("test done")) })

k, address := testkafka.CreateClusterWithoutCustomConsumerGroupsSupport(t, 1, "test-topic")
t.Cleanup(k.Close)
k, address := testkafka.CreateCluster(t, 1, testTopic)

kafkaCommits := atomic.NewInt32(0)
k.ControlKey(kmsg.OffsetCommit.Int16(), func(kmsg.Request) (kmsg.Response, error, bool) {
kafkaCommits.Add(1)
k.ControlKey(kmsg.OffsetCommit, func(kmsg.Request) (kmsg.Response, error, bool) {
kafkaCommits.Inc()
return nil, nil, false
})

Expand All @@ -57,7 +60,7 @@ func TestBlockbuilder_lookbackOnNoCommit(t *testing.T) {
})

client := newKafkaClient(t, cfg.IngestStorageConfig.Kafka)
sendReq(t, ctx, client)
producedRecords := sendReq(t, ctx, client)

// Wait for record to be consumed and committed.
require.Eventually(t, func() bool {
Expand All @@ -66,8 +69,11 @@ func TestBlockbuilder_lookbackOnNoCommit(t *testing.T) {

// Wait for the block to be flushed.
require.Eventually(t, func() bool {
return len(store.BlockMetas(util.FakeTenantID)) == 1
return len(store.BlockMetas(util.FakeTenantID)) == 1 && countFlushedTraces(store) == 1
}, time.Minute, time.Second)

// Check committed offset
requireLastCommitEquals(t, ctx, client, producedRecords[len(producedRecords)-1].Offset+1)
}

// Starting with a pre-existing commit,
Expand All @@ -77,12 +83,11 @@ func TestBlockbuilder_startWithCommit(t *testing.T) {
ctx, cancel := context.WithCancelCause(context.Background())
t.Cleanup(func() { cancel(errors.New("test done")) })

k, address := testkafka.CreateClusterWithoutCustomConsumerGroupsSupport(t, 1, "test-topic")
t.Cleanup(k.Close)
k, address := testkafka.CreateCluster(t, 1, testTopic)

kafkaCommits := atomic.NewInt32(0)
k.ControlKey(kmsg.OffsetCommit.Int16(), func(kmsg.Request) (kmsg.Response, error, bool) {
kafkaCommits.Add(1)
k.ControlKey(kmsg.OffsetCommit, func(kmsg.Request) (kmsg.Response, error, bool) {
kafkaCommits.Inc()
return nil, nil, false
})

Expand Down Expand Up @@ -121,26 +126,28 @@ func TestBlockbuilder_startWithCommit(t *testing.T) {
require.Eventually(t, func() bool {
return countFlushedTraces(store) == len(producedRecords)-commitedAt
}, time.Minute, time.Second)

// Check committed offset
requireLastCommitEquals(t, ctx, client, producedRecords[len(producedRecords)-1].Offset+1)
}

// In case a block flush initially fails, the system retries until it succeeds.
func TestBlockbuilder_flushingFails(t *testing.T) {
ctx, cancel := context.WithCancelCause(context.Background())
t.Cleanup(func() { cancel(errors.New("test done")) })

k, address := testkafka.CreateClusterWithoutCustomConsumerGroupsSupport(t, 1, "test-topic")
t.Cleanup(k.Close)
k, address := testkafka.CreateCluster(t, 1, "test-topic")

kafkaCommits := atomic.NewInt32(0)
k.ControlKey(kmsg.OffsetCommit.Int16(), func(kmsg.Request) (kmsg.Response, error, bool) {
kafkaCommits.Add(1)
k.ControlKey(kmsg.OffsetCommit, func(kmsg.Request) (kmsg.Response, error, bool) {
kafkaCommits.Inc()
return nil, nil, false
})

storageWrites := atomic.NewInt32(0)
store := newStoreWrapper(newStore(ctx, t), func(ctx context.Context, block tempodb.WriteableBlock, store storage.Store) error {
// Fail the first block write
if storageWrites.Add(1) == 1 {
if storageWrites.Inc() == 1 {
return errors.New("failed to write block")
}
return store.WriteBlock(ctx, block)
Expand All @@ -149,7 +156,7 @@ func TestBlockbuilder_flushingFails(t *testing.T) {
logger := test.NewTestingLogger(t)

client := newKafkaClient(t, cfg.IngestStorageConfig.Kafka)
sendTracesFor(t, ctx, client, time.Second, 100*time.Millisecond) // Send for 1 second, <1 consumption cycles
producedRecords := sendTracesFor(t, ctx, client, time.Second, 100*time.Millisecond) // Send for 1 second, <1 consumption cycles

b := New(cfg, logger, newPartitionRingReader(), &mockOverrides{}, store)
require.NoError(t, services.StartAndAwaitRunning(ctx, b))
Expand All @@ -164,6 +171,9 @@ func TestBlockbuilder_flushingFails(t *testing.T) {
require.Eventually(t, func() bool {
return len(store.BlockMetas(util.FakeTenantID)) >= 1
}, time.Minute, time.Second)

// Check committed offset
requireLastCommitEquals(t, ctx, client, producedRecords[len(producedRecords)-1].Offset+1)
}

// Receiving records with older timestamps the block-builder processes them in the current cycle,
Expand All @@ -172,13 +182,11 @@ func TestBlockbuilder_receivesOldRecords(t *testing.T) {
ctx, cancel := context.WithCancelCause(context.Background())
t.Cleanup(func() { cancel(errors.New("test done")) })

k, address := testkafka.CreateClusterWithoutCustomConsumerGroupsSupport(t, 1, "test-topic")
t.Cleanup(k.Close)
k, address := testkafka.CreateCluster(t, 1, "test-topic")

kafkaCommits := atomic.NewInt32(0)
k.ControlKey(kmsg.OffsetCommit.Int16(), func(kmsg.Request) (kmsg.Response, error, bool) {
k.KeepControl()
kafkaCommits.Add(1)
k.ControlKey(kmsg.OffsetCommit, func(kmsg.Request) (kmsg.Response, error, bool) {
kafkaCommits.Inc()
return nil, nil, false
})

Expand Down Expand Up @@ -221,6 +229,9 @@ func TestBlockbuilder_receivesOldRecords(t *testing.T) {
require.Eventually(t, func() bool {
return len(store.BlockMetas(util.FakeTenantID)) == 2
}, time.Minute, time.Second)

// Check committed offset
requireLastCommitEquals(t, ctx, client, producedRecords[len(producedRecords)-1].Offset+1)
}

// FIXME - Test is unstable and will fail if records cross two consumption cycles,
Expand All @@ -238,15 +249,13 @@ func TestBlockbuilder_committingFails(t *testing.T) {
ctx, cancel := context.WithCancelCause(context.Background())
t.Cleanup(func() { cancel(errors.New("test done")) })

k, address := testkafka.CreateClusterWithoutCustomConsumerGroupsSupport(t, 1, "test-topic")
t.Cleanup(k.Close)
k, address := testkafka.CreateCluster(t, 1, "test-topic")

kafkaCommits := atomic.NewInt32(0)
k.ControlKey(kmsg.OffsetCommit.Int16(), func(req kmsg.Request) (kmsg.Response, error, bool) {
k.KeepControl()
defer kafkaCommits.Add(1)
k.ControlKey(kmsg.OffsetCommit, func(req kmsg.Request) (kmsg.Response, error, bool) {
kafkaCommits.Inc()

if kafkaCommits.Load() == 0 { // First commit fails
if kafkaCommits.Load() == 1 { // First commit fails
res := kmsg.NewOffsetCommitResponse()
res.Version = req.GetVersion()
res.Topics = []kmsg.OffsetCommitResponseTopic{
Expand All @@ -271,7 +280,7 @@ func TestBlockbuilder_committingFails(t *testing.T) {
logger := test.NewTestingLogger(t)

client := newKafkaClient(t, cfg.IngestStorageConfig.Kafka)
sendTracesFor(t, ctx, client, time.Second, 100*time.Millisecond) // Send for 1 second, <1 consumption cycles
producedRecords := sendTracesFor(t, ctx, client, time.Second, 100*time.Millisecond) // Send for 1 second, <1 consumption cycles

b := New(cfg, logger, newPartitionRingReader(), &mockOverrides{}, store)
require.NoError(t, services.StartAndAwaitRunning(ctx, b))
Expand All @@ -288,6 +297,9 @@ func TestBlockbuilder_committingFails(t *testing.T) {
require.Eventually(t, func() bool {
return len(store.BlockMetas(util.FakeTenantID)) == 1 // Only one block should have been written
}, time.Minute, time.Second)

// Check committed offset
requireLastCommitEquals(t, ctx, client, producedRecords[len(producedRecords)-1].Offset+1)
}

func TestCycleEndAtStartup(t *testing.T) {
Expand Down Expand Up @@ -375,7 +387,7 @@ func blockbuilderConfig(t *testing.T, address string) Config {
flagext.DefaultValues(&cfg.IngestStorageConfig.Kafka)
cfg.IngestStorageConfig.Kafka.Address = address
cfg.IngestStorageConfig.Kafka.Topic = testTopic
cfg.IngestStorageConfig.Kafka.ConsumerGroup = "test-consumer-group"
cfg.IngestStorageConfig.Kafka.ConsumerGroup = testConsumerGroup

cfg.AssignedPartitions = map[string][]int32{cfg.InstanceID: {0}}
cfg.LookbackOnNoCommit = 15 * time.Second
Expand Down Expand Up @@ -491,6 +503,7 @@ func countFlushedTraces(store storage.Store) int {
return count
}

// nolint: revive
func sendReq(t *testing.T, ctx context.Context, client *kgo.Client) []*kgo.Record {
traceID := generateTraceID(t)

Expand All @@ -504,6 +517,7 @@ func sendReq(t *testing.T, ctx context.Context, client *kgo.Client) []*kgo.Recor
return records
}

// nolint: revive
func sendTracesFor(t *testing.T, ctx context.Context, client *kgo.Client, dur, interval time.Duration) []*kgo.Record {
ticker := time.NewTicker(interval)
defer ticker.Stop()
Expand Down Expand Up @@ -532,3 +546,12 @@ func generateTraceID(t *testing.T) []byte {
require.NoError(t, err)
return traceID
}

// nolint: revive
func requireLastCommitEquals(t testing.TB, ctx context.Context, client *kgo.Client, expectedOffset int64) {
offsets, err := kadm.NewClient(client).FetchOffsetsForTopics(ctx, testConsumerGroup, testTopic)
require.NoError(t, err)
offset, ok := offsets.Lookup(testTopic, testPartition)
require.True(t, ok)
require.Equal(t, expectedOffset, offset.At)
}
Loading