fix: solve Offset Manager and Worker racing condition when stopping

Farfetch · Dec 6, 2023 · 453996e · 453996e
1 parent 40f6bac
commit 453996e
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 18 deletions.
diff --git a/src/KafkaFlow/Consumers/ConsumerWorker.cs b/src/KafkaFlow/Consumers/ConsumerWorker.cs
@@ -59,11 +59,9 @@ public ConsumerWorker(
 
         public IEvent<IMessageContext> WorkerProcessingEnded => _workerProcessingEnded;
 
-        public ValueTask EnqueueAsync(
-            IMessageContext context,
-            CancellationToken stopCancellationToken)
+        public ValueTask EnqueueAsync(IMessageContext context)
         {
-            return _messagesBuffer.Writer.WriteAsync(context, stopCancellationToken);
+            return _messagesBuffer.Writer.WriteAsync(context, CancellationToken.None);
         }
 
         public Task StartAsync()

diff --git a/src/KafkaFlow/Consumers/ConsumerWorkerPool.cs b/src/KafkaFlow/Consumers/ConsumerWorkerPool.cs
@@ -127,7 +127,7 @@ public async Task StopAsync()
 
         public async Task EnqueueAsync(ConsumeResult<byte[], byte[]> message, CancellationToken stopCancellationToken)
         {
-            await _startedTaskSource.Task.ConfigureAwait(false);
+            await _startedTaskSource.Task;
 
             var worker = (IConsumerWorker)await _distributionStrategy
                 .GetWorkerAsync(
@@ -136,8 +136,7 @@ public async Task EnqueueAsync(ConsumeResult<byte[], byte[]> message, Cancellati
                         message.Topic,
                         message.Partition.Value,
                         message.Message.Key,
-                        stopCancellationToken))
-                .ConfigureAwait(false);
+                        stopCancellationToken));
 
             if (worker is null)
             {
@@ -146,11 +145,9 @@ public async Task EnqueueAsync(ConsumeResult<byte[], byte[]> message, Cancellati
 
             var context = this.CreateMessageContext(message, worker);
 
-            await worker
-                .EnqueueAsync(context, stopCancellationToken)
-                .ConfigureAwait(false);
-
             _offsetManager.Enqueue(context.ConsumerContext);
+
+            await worker.EnqueueAsync(context);
         }
 
         private MessageContext CreateMessageContext(ConsumeResult<byte[], byte[]> message, IConsumerWorker worker)

diff --git a/src/KafkaFlow/Consumers/IConsumerWorker.cs b/src/KafkaFlow/Consumers/IConsumerWorker.cs
@@ -10,7 +10,7 @@ internal interface IConsumerWorker : IWorker, IDisposable
 
         IDependencyResolver WorkerDependencyResolver { get; }
 
-        ValueTask EnqueueAsync(IMessageContext context, CancellationToken stopCancellationToken);
+        ValueTask EnqueueAsync(IMessageContext context);
 
         Task StartAsync();
 

diff --git a/src/KafkaFlow/Consumers/PartitionOffsets.cs b/src/KafkaFlow/Consumers/PartitionOffsets.cs
@@ -42,16 +42,24 @@ public bool TryDequeue(IConsumerContext context)
                 {
                     this.DequeuedContext = _receivedContexts.First.Value;
                     _receivedContexts.RemoveFirst();
-                }
-                while (_receivedContexts.Count > 0 && _processedContexts.Remove(_receivedContexts.First.Value.Offset));
+                } while (_receivedContexts.Count > 0 && _processedContexts.Remove(_receivedContexts.First.Value.Offset));
             }
 
             return true;
         }
 
-        public Task WaitContextsCompletionAsync() => Task.WhenAll(
-            _receivedContexts
-                .Select(x => x.Completion)
-                .ToList());
+        public Task WaitContextsCompletionAsync()
+        {
+            List<Task> tasks;
+
+            lock (_receivedContexts)
+            {
+                tasks = _receivedContexts
+                    .Select(x => (Task)x.Completion)
+                    .ToList();
+            }
+
+            return Task.WhenAll(tasks);
+        }
     }
 }