Switch to RoundRobbinBalancer

pkg/usage/service.go

@@ -37,8 +37,18 @@ func NewService(kafkaCfg kafka.Config, consumerGroup string, logger log.Logger,
 	client, err := client.NewReaderClient(kafkaCfg, kprom, logger,
 		kgo.ConsumerGroup(consumerGroup),
 		kgo.ConsumeTopics(topic),
+		kgo.Balancers(kgo.RoundRobinBalancer()),
 		kgo.ConsumeResetOffset(kgo.NewOffset().AfterMilli(time.Now().Add(-windowSize).UnixMilli())),
 		kgo.DisableAutoCommit(),
+		kgo.OnPartitionsAssigned(func(_ context.Context, _ *kgo.Client, partitions map[string][]int32) {
+			level.Info(logger).Log("msg", "assigned partitions", "partitions", partitions)
+		}),
+		kgo.OnPartitionsRevoked(func(_ context.Context, _ *kgo.Client, partitions map[string][]int32) {
+			level.Info(logger).Log("msg", "revoked partitions", "partitions", partitions)
+		}),
+		kgo.OnPartitionsLost(func(_ context.Context, _ *kgo.Client, partitions map[string][]int32) {
+			level.Info(logger).Log("msg", "lost partitions", "partitions", partitions)
+		}),
 	)
 	if err != nil {
 		return nil, err

pkg/usage/stats.go

@@ -55,6 +55,11 @@ func (u *usageStats) addEntry(partition int32, tenantID string, streamHash uint6
 		u.stats[partition] = pStats
 	}
 
+	// If the partition has already processed this offset, skip it
+	if pStats.offset != 0 && pStats.offset >= offset {
+		return
+	}
+
 	pStats.offset = offset
 	pStats.timestamp = timestamp