fix(consumer): prevent offset loss during rebalance under high load

src/KafkaFlow/Consumers/ConsumerManager.cs

@@ -109,10 +109,17 @@ private void OnPartitionRevoked(IEnumerable<Confluent.Kafka.TopicPartitionOffset
             this.GetConsumerLogInfo(topicPartitions?.Select(x => x.TopicPartition).ToArray()
                                     ?? Array.Empty<Confluent.Kafka.TopicPartition>()));
 
+        // 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<Confluent.Kafka.TopicPartitionOffset
             .StartAsync(assignedPartitions, workersCount)
             .GetAwaiter()
             .GetResult();
+
+        this.Feeder.Start();
     }
 
     private void OnPartitionAssigned(IReadOnlyCollection<Confluent.Kafka.TopicPartition> partitions)
@@ -131,6 +140,12 @@ private void OnPartitionAssigned(IReadOnlyCollection<Confluent.Kafka.TopicPartit
             "Partitions assigned",
             this.GetConsumerLogInfo(partitions));
 
+        // Cancel any currently-running feeder before rebuilding the worker pool.
+        // On initial startup the feeder was started in StartAsync; cancelling it here
+        // and restarting after the pool is ready closes the window where a message from
+        // a no-longer-assigned partition could slip through without offset tracking.
+        this.Feeder.Cancel();
+
         if (_stopTheWorldStrategy is false)
         {
             this.WorkerPool.StopAsync().GetAwaiter().GetResult();
@@ -143,6 +158,8 @@ private void OnPartitionAssigned(IReadOnlyCollection<Confluent.Kafka.TopicPartit
             .StartAsync(assignedPartitions, workersCount)
             .GetAwaiter()
             .GetResult();
+
+        this.Feeder.Start();
     }
 
     private object GetConsumerLogInfo(IEnumerable<Confluent.Kafka.TopicPartition> partitions) => new

src/KafkaFlow/Consumers/ConsumerWorkerPool.cs

@@ -166,7 +166,13 @@ public async Task EnqueueAsync(ConsumeResult<byte[], byte[]> 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);
     }

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);
 

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();
 }

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)
     {

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() =>

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);
     }
 }

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<Confluent.Kafka.TopicPartition>(6).ToList();
+        var revokedPartitions = currentPartitions.Take(3).ToArray();
+        var leftPartitions = currentPartitions.Except(revokedPartitions).ToArray();
+
+        var callOrder = new List<string>();
+
+        _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<int>()))
+            .Callback(() => callOrder.Add("WorkerPool.StartAsync"))
+            .Returns(Task.CompletedTask);
+
+        _logHandlerMock
+            .Setup(x => x.Warning(It.IsAny<string>(), It.IsAny<object>()));
+
+        _consumerMock.SetupGet(x => x.Assignment).Returns(leftPartitions);
+
+        // Act
+        _onPartitionRevokedHandler(
+            _dependencyResolver.Object,
+            Mock.Of<Confluent.Kafka.IConsumer<byte[], byte[]>>(),
+            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<List<Confluent.Kafka.TopicPartition>>();
+        var newAssignedPartitions = _fixture.Create<List<Confluent.Kafka.TopicPartition>>();
+        var allPartitions = currentPartitions.Concat(newAssignedPartitions).ToArray();
+
+        var callOrder = new List<string>();
+
+        _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<int>()))
+            .Callback(() => callOrder.Add("WorkerPool.StartAsync"))
+            .Returns(Task.CompletedTask);
+
+        _logHandlerMock
+            .Setup(x => x.Info(It.IsAny<string>(), It.IsAny<object>()));
+
+        _consumerMock.SetupGet(x => x.Assignment).Returns(allPartitions);
+
+        // Act
+        _onPartitionAssignedHandler(
+            _dependencyResolver.Object,
+            Mock.Of<Confluent.Kafka.IConsumer<byte[], byte[]>>(),
+            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]);
+    }
 }

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<List<Confluent.Kafka.TopicPartitionOffset>>();
+
+        var callOrder = new List<string>();
+
+        _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<string>(), It.IsAny<object>()));
+
+        // 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<List<Confluent.Kafka.TopicPartition>>();
+
+        var callOrder = new List<string>();
+
+        _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<int>()))
+            .Callback(() => callOrder.Add("WorkerPool.StartAsync"))
+            .Returns(Task.CompletedTask);
+
+        _logHandlerMock
+            .Setup(x => x.Info(It.IsAny<string>(), It.IsAny<object>()));
+
+        _consumerMock.SetupGet(x => x.Assignment).Returns(partitions.ToArray());
+
+        // Act
+        _onPartitionAssignedHandler(_dependencyResolver.Object, Mock.Of<Confluent.Kafka.IConsumer<byte[], byte[]>>(), 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]);
+    }
 }

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<TopicPartitionOffset>()),
+            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<IConsumerContext>();
-        var tpo = new TopicPartitionOffset(_topicPartition.Topic, _topicPartition.Partition, offset);
+        var tpo = new TopicPartitionOffset(topic, partition, offset);
 
         mock
             .SetupGet(x => x.Offset)