From 2d73cc6690022f53781912e715f9e7fb0cae15ca Mon Sep 17 00:00:00 2001 From: nimaniko Date: Wed, 24 Jun 2026 21:12:54 +0330 Subject: [PATCH] fix(consumer): prevent offset loss during rebalance under high load --- src/KafkaFlow/Consumers/ConsumerManager.cs | 17 +++ src/KafkaFlow/Consumers/ConsumerWorkerPool.cs | 8 +- src/KafkaFlow/Consumers/IOffsetManager.cs | 4 +- src/KafkaFlow/Consumers/IWorkerPoolFeeder.cs | 5 + src/KafkaFlow/Consumers/NullOffsetManager.cs | 5 +- src/KafkaFlow/Consumers/OffsetManager.cs | 9 +- src/KafkaFlow/Consumers/WorkerPoolFeeder.cs | 6 +- .../ConsumerManagerCooperativeStickyTests.cs | 106 ++++++++++++++++++ .../Consumer/ConsumerManagerTests.cs | 79 +++++++++++++ .../KafkaFlow.UnitTests/OffsetManagerTests.cs | 31 ++++- 10 files changed, 258 insertions(+), 12 deletions(-) diff --git a/src/KafkaFlow/Consumers/ConsumerManager.cs b/src/KafkaFlow/Consumers/ConsumerManager.cs index f668c51f7..82d58f4b3 100644 --- a/src/KafkaFlow/Consumers/ConsumerManager.cs +++ b/src/KafkaFlow/Consumers/ConsumerManager.cs @@ -109,10 +109,17 @@ private void OnPartitionRevoked(IEnumerable x.TopicPartition).ToArray() ?? Array.Empty())); + // Cancel the feeder before draining workers. We cannot await the feeder task here + // because this callback runs synchronously on the same thread as consumer.Consume(), + // which is the thread the feeder is blocked on — awaiting would deadlock. + // Cancelling the token is safe: the feeder exits after Consume() returns, and + // EnqueueAsync discards any message whose partition is no longer tracked. + this.Feeder.Cancel(); this.WorkerPool.StopAsync().GetAwaiter().GetResult(); if (_stopTheWorldStrategy) { + // For stop-the-world strategies the feeder is restarted in OnPartitionAssigned. return; } @@ -123,6 +130,8 @@ private void OnPartitionRevoked(IEnumerable partitions) @@ -131,6 +140,12 @@ private void OnPartitionAssigned(IReadOnlyCollection partitions) => new diff --git a/src/KafkaFlow/Consumers/ConsumerWorkerPool.cs b/src/KafkaFlow/Consumers/ConsumerWorkerPool.cs index 17fe80a09..212bf109f 100644 --- a/src/KafkaFlow/Consumers/ConsumerWorkerPool.cs +++ b/src/KafkaFlow/Consumers/ConsumerWorkerPool.cs @@ -166,7 +166,13 @@ public async Task EnqueueAsync(ConsumeResult message, Cancellati var context = this.CreateMessageContext(message, worker); - _offsetManager.Enqueue(context.ConsumerContext); + if (!_offsetManager.Enqueue(context.ConsumerContext)) + { + // Partition was revoked between when the message was polled and when it + // reached this point (feeder race during rebalance). Discard it here — + // Kafka will replay it on the new partition owner after the rebalance. + return; + } await worker.EnqueueAsync(context).ConfigureAwait(false); } diff --git a/src/KafkaFlow/Consumers/IOffsetManager.cs b/src/KafkaFlow/Consumers/IOffsetManager.cs index 48b09aa9e..27d0ccf11 100644 --- a/src/KafkaFlow/Consumers/IOffsetManager.cs +++ b/src/KafkaFlow/Consumers/IOffsetManager.cs @@ -4,7 +4,9 @@ namespace KafkaFlow.Consumers; internal interface IOffsetManager { - void Enqueue(IConsumerContext context); + // Returns false when the context's partition is not tracked (e.g. was revoked). + // The caller must skip further processing if false is returned. + bool Enqueue(IConsumerContext context); void MarkAsProcessed(IConsumerContext offset); diff --git a/src/KafkaFlow/Consumers/IWorkerPoolFeeder.cs b/src/KafkaFlow/Consumers/IWorkerPoolFeeder.cs index 09fb23fcb..0887e636d 100644 --- a/src/KafkaFlow/Consumers/IWorkerPoolFeeder.cs +++ b/src/KafkaFlow/Consumers/IWorkerPoolFeeder.cs @@ -6,5 +6,10 @@ internal interface IWorkerPoolFeeder { void Start(); + // Cancels the feeder token without awaiting the feeder task. + // Safe to call from Confluent.Kafka rebalance callbacks (which run on the polling + // thread), where awaiting the feeder task would deadlock. + void Cancel(); + Task StopAsync(); } diff --git a/src/KafkaFlow/Consumers/NullOffsetManager.cs b/src/KafkaFlow/Consumers/NullOffsetManager.cs index 19cd0026c..fd6320077 100644 --- a/src/KafkaFlow/Consumers/NullOffsetManager.cs +++ b/src/KafkaFlow/Consumers/NullOffsetManager.cs @@ -4,10 +4,7 @@ namespace KafkaFlow.Consumers; internal class NullOffsetManager : IOffsetManager { - public void Enqueue(IConsumerContext context) - { - // Do nothing - } + public bool Enqueue(IConsumerContext context) => true; public void MarkAsProcessed(IConsumerContext offset) { diff --git a/src/KafkaFlow/Consumers/OffsetManager.cs b/src/KafkaFlow/Consumers/OffsetManager.cs index 515b1b854..d8036b88d 100644 --- a/src/KafkaFlow/Consumers/OffsetManager.cs +++ b/src/KafkaFlow/Consumers/OffsetManager.cs @@ -35,14 +35,17 @@ public void MarkAsProcessed(IConsumerContext context) } } - public void Enqueue(IConsumerContext context) + public bool Enqueue(IConsumerContext context) { - if (_partitionsOffsets.TryGetValue( + if (!_partitionsOffsets.TryGetValue( (context.Topic, context.Partition), out var offsets)) { - offsets.Enqueue(context); + return false; } + + offsets.Enqueue(context); + return true; } public Task WaitContextsCompletionAsync() => diff --git a/src/KafkaFlow/Consumers/WorkerPoolFeeder.cs b/src/KafkaFlow/Consumers/WorkerPoolFeeder.cs index 2b02d8652..d2038b161 100644 --- a/src/KafkaFlow/Consumers/WorkerPoolFeeder.cs +++ b/src/KafkaFlow/Consumers/WorkerPoolFeeder.cs @@ -64,14 +64,18 @@ await _workerPool CancellationToken.None); } - public async Task StopAsync() + public void Cancel() { if (_stopTokenSource is { IsCancellationRequested: false }) { _stopTokenSource.Cancel(); _stopTokenSource.Dispose(); } + } + public async Task StopAsync() + { + this.Cancel(); await (_feederTask ?? Task.CompletedTask).ConfigureAwait(false); } } diff --git a/tests/KafkaFlow.UnitTests/Consumer/ConsumerManagerCooperativeStickyTests.cs b/tests/KafkaFlow.UnitTests/Consumer/ConsumerManagerCooperativeStickyTests.cs index 7de6090d0..96297f0d7 100644 --- a/tests/KafkaFlow.UnitTests/Consumer/ConsumerManagerCooperativeStickyTests.cs +++ b/tests/KafkaFlow.UnitTests/Consumer/ConsumerManagerCooperativeStickyTests.cs @@ -190,4 +190,110 @@ public void OnPartitionsRevoked_StopWorkerPool() _workerPoolMock.VerifyNoOtherCalls(); _logHandlerMock.VerifyAll(); } + + [TestMethod] + public void OnPartitionsRevoked_ShouldCancelFeederBeforeWorkerPool_ToPreventOffsetLossUnderHighLoad() + { + // Arrange + // The rebalance callback runs on the same thread as consumer.Consume(), which is the + // thread the feeder is blocked on. Calling Feeder.StopAsync() (which awaits the feeder + // task) from here would deadlock. Cancel() just signals the token; the feeder exits + // after Consume() returns, and EnqueueAsync discards messages from revoked partitions. + var currentPartitions = _fixture.CreateMany(6).ToList(); + var revokedPartitions = currentPartitions.Take(3).ToArray(); + var leftPartitions = currentPartitions.Except(revokedPartitions).ToArray(); + + var callOrder = new List(); + + _feederMock + .Setup(x => x.Cancel()) + .Callback(() => callOrder.Add("Feeder.Cancel")); + + _feederMock + .Setup(x => x.Start()) + .Callback(() => callOrder.Add("Feeder.Start")); + + _workerPoolMock + .Setup(x => x.StopAsync()) + .Callback(() => callOrder.Add("WorkerPool.StopAsync")) + .Returns(Task.CompletedTask); + + _workerPoolMock + .Setup(x => x.StartAsync(leftPartitions, It.IsAny())) + .Callback(() => callOrder.Add("WorkerPool.StartAsync")) + .Returns(Task.CompletedTask); + + _logHandlerMock + .Setup(x => x.Warning(It.IsAny(), It.IsAny())); + + _consumerMock.SetupGet(x => x.Assignment).Returns(leftPartitions); + + // Act + _onPartitionRevokedHandler( + _dependencyResolver.Object, + Mock.Of>(), + revokedPartitions.Select(x => new Confluent.Kafka.TopicPartitionOffset(x, 123)).ToList()); + + // Assert + _feederMock.Verify(x => x.Cancel(), Times.Once); + _feederMock.Verify(x => x.Start(), Times.Once); + + Assert.AreEqual(4, callOrder.Count); + Assert.AreEqual("Feeder.Cancel", callOrder[0]); + Assert.AreEqual("WorkerPool.StopAsync", callOrder[1]); + Assert.AreEqual("WorkerPool.StartAsync", callOrder[2]); + Assert.AreEqual("Feeder.Start", callOrder[3]); + } + + [TestMethod] + public void OnPartitionsAssigned_ShouldCancelAndRestartFeeder_ToPreventOffsetLossUnderHighLoad() + { + // Arrange + // Same race as OnPartitionsRevoked: feeder must be cancelled and restarted around + // every pool rebuild so no message slips through without offset tracking. + var currentPartitions = _fixture.Create>(); + var newAssignedPartitions = _fixture.Create>(); + var allPartitions = currentPartitions.Concat(newAssignedPartitions).ToArray(); + + var callOrder = new List(); + + _feederMock + .Setup(x => x.Cancel()) + .Callback(() => callOrder.Add("Feeder.Cancel")); + + _feederMock + .Setup(x => x.Start()) + .Callback(() => callOrder.Add("Feeder.Start")); + + _workerPoolMock + .Setup(x => x.StopAsync()) + .Callback(() => callOrder.Add("WorkerPool.StopAsync")) + .Returns(Task.CompletedTask); + + _workerPoolMock + .Setup(x => x.StartAsync(allPartitions, It.IsAny())) + .Callback(() => callOrder.Add("WorkerPool.StartAsync")) + .Returns(Task.CompletedTask); + + _logHandlerMock + .Setup(x => x.Info(It.IsAny(), It.IsAny())); + + _consumerMock.SetupGet(x => x.Assignment).Returns(allPartitions); + + // Act + _onPartitionAssignedHandler( + _dependencyResolver.Object, + Mock.Of>(), + newAssignedPartitions); + + // Assert + _feederMock.Verify(x => x.Cancel(), Times.Once); + _feederMock.Verify(x => x.Start(), Times.Once); + + Assert.AreEqual(4, callOrder.Count); + Assert.AreEqual("Feeder.Cancel", callOrder[0]); + Assert.AreEqual("WorkerPool.StopAsync", callOrder[1]); + Assert.AreEqual("WorkerPool.StartAsync", callOrder[2]); + Assert.AreEqual("Feeder.Start", callOrder[3]); + } } diff --git a/tests/KafkaFlow.UnitTests/Consumer/ConsumerManagerTests.cs b/tests/KafkaFlow.UnitTests/Consumer/ConsumerManagerTests.cs index 3eeccdf11..7d13ed518 100644 --- a/tests/KafkaFlow.UnitTests/Consumer/ConsumerManagerTests.cs +++ b/tests/KafkaFlow.UnitTests/Consumer/ConsumerManagerTests.cs @@ -175,4 +175,83 @@ public void OnPartitionsRevoked_StopWorkerPool() _workerPoolMock.VerifyNoOtherCalls(); _logHandlerMock.VerifyAll(); } + + [TestMethod] + public void OnPartitionsRevoked_ShouldCancelFeederBeforeWorkerPool_ToPreventOffsetLossUnderHighLoad() + { + // Arrange + // Cancel() is used instead of StopAsync() because the callback runs on the same + // thread as consumer.Consume() (the feeder thread). Awaiting the feeder task from + // there would deadlock. Cancel() just signals the cancellation token; the feeder + // exits naturally after Consume() returns. + var partitions = _fixture.Create>(); + + var callOrder = new List(); + + _feederMock + .Setup(x => x.Cancel()) + .Callback(() => callOrder.Add("Feeder.Cancel")); + + _workerPoolMock + .Setup(x => x.StopAsync()) + .Callback(() => callOrder.Add("WorkerPool.StopAsync")) + .Returns(Task.CompletedTask); + + _logHandlerMock + .Setup(x => x.Warning(It.IsAny(), It.IsAny())); + + // Act + _onPartitionRevokedHandler(_dependencyResolver.Object, null, partitions); + + // Assert - feeder cancelled before pool drain + _feederMock.Verify(x => x.Cancel(), Times.Once); + + Assert.AreEqual(2, callOrder.Count); + Assert.AreEqual("Feeder.Cancel", callOrder[0]); + Assert.AreEqual("WorkerPool.StopAsync", callOrder[1]); + } + + [TestMethod] + public void OnPartitionsAssigned_ShouldCancelAndRestartFeeder_ToCloseRebalanceRaceWindow() + { + // Arrange + // OnPartitionAssigned must cancel any running feeder (including the one started at + // startup) before the pool is rebuilt, then restart it after. This closes the window + // where a stale message from a no-longer-assigned partition could be enqueued without + // offset tracking. For stop-the-world strategy the pool is NOT stopped here (it was + // already stopped in OnPartitionsRevoked). + var partitions = _fixture.Create>(); + + var callOrder = new List(); + + _feederMock + .Setup(x => x.Cancel()) + .Callback(() => callOrder.Add("Feeder.Cancel")); + + _feederMock + .Setup(x => x.Start()) + .Callback(() => callOrder.Add("Feeder.Start")); + + _workerPoolMock + .Setup(x => x.StartAsync(partitions, It.IsAny())) + .Callback(() => callOrder.Add("WorkerPool.StartAsync")) + .Returns(Task.CompletedTask); + + _logHandlerMock + .Setup(x => x.Info(It.IsAny(), It.IsAny())); + + _consumerMock.SetupGet(x => x.Assignment).Returns(partitions.ToArray()); + + // Act + _onPartitionAssignedHandler(_dependencyResolver.Object, Mock.Of>(), partitions); + + // Assert - order: cancel → start pool → start feeder + _feederMock.Verify(x => x.Cancel(), Times.Once); + _feederMock.Verify(x => x.Start(), Times.Once); + + Assert.AreEqual(3, callOrder.Count); + Assert.AreEqual("Feeder.Cancel", callOrder[0]); + Assert.AreEqual("WorkerPool.StartAsync", callOrder[1]); + Assert.AreEqual("Feeder.Start", callOrder[2]); + } } diff --git a/tests/KafkaFlow.UnitTests/OffsetManagerTests.cs b/tests/KafkaFlow.UnitTests/OffsetManagerTests.cs index 830b65a66..dee17edb8 100644 --- a/tests/KafkaFlow.UnitTests/OffsetManagerTests.cs +++ b/tests/KafkaFlow.UnitTests/OffsetManagerTests.cs @@ -60,10 +60,37 @@ public void MarkAsProcessed_WithGaps_ShouldStoreOffsetJustOnce() Times.Once); } - private IConsumerContext MockConsumerContext(int offset) + [TestMethod] + public void Enqueue_ForRevokedPartition_SilentlyDropsMessage_CausingOffsetLoss() + { + // This test proves the mechanism behind Bug #1 in ConsumerManager: + // When the feeder is not stopped during rebalance, it can consume a message + // from a partition that was just revoked. That message gets enqueued into a + // new OffsetManager (created for the new partition set). Since the revoked + // partition is not in the new OffsetManager, Enqueue and MarkAsProcessed + // both silently no-op. The message is processed by a worker but its offset + // is NEVER committed. After the rebalance the broker replays that offset. + var revokedPartition = new Confluent.Kafka.TopicPartition("topic-A", new Partition(99)); + var context = MockConsumerContext(100, revokedPartition.Topic, revokedPartition.Partition.Value); + + // Act - message from revoked partition arrives in new OffsetManager + _target.Enqueue(context); + _target.MarkAsProcessed(context); + + // Assert - offset is silently dropped, committer is never called + _committerMock.Verify( + c => c.MarkAsProcessed(It.IsAny()), + Times.Never, + "Offset from a partition unknown to this OffsetManager is silently discarded"); + } + + private IConsumerContext MockConsumerContext(int offset) => + MockConsumerContext(offset, _topicPartition.Topic, _topicPartition.Partition.Value); + + private IConsumerContext MockConsumerContext(int offset, string topic, int partition) { var mock = new Mock(); - var tpo = new TopicPartitionOffset(_topicPartition.Topic, _topicPartition.Partition, offset); + var tpo = new TopicPartitionOffset(topic, partition, offset); mock .SetupGet(x => x.Offset)