Performance change to improve throughput

Reduce thread contention and await states
This change currently breaks the maximumMessageBuffer

Author: Jroland

src/kafka-net/Producer.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -23,17 +24,22 @@ public class Producer : IMetadataQueries
 
         private readonly CancellationTokenSource _stopToken = new CancellationTokenSource();
         private readonly int _maximumAsyncRequests;
-        private readonly int _maximumMessageBuffer;
         private readonly AsyncCollection<TopicMessage> _asyncCollection;
         private readonly SemaphoreSlim _semaphoreMaximumAsync;
-        private readonly SemaphoreSlim _boundedCapacitySemaphore;
         private readonly IMetadataQueries _metadataQueries;
         private readonly Task _postTask;
 
+        private int _inFlightMessageCount = 0;
+
         /// <summary>
         /// Get the number of messages sitting in the buffer waiting to be sent. 
         /// </summary>
-        public int BufferCount { get { return _maximumMessageBuffer - _boundedCapacitySemaphore.CurrentCount; } }
+        public int BufferCount { get { return _asyncCollection.Count; } }
+
+        /// <summary>
+        /// Get the number of messages staged for Async upload.
+        /// </summary>
+        public int InFlightMessageCount { get { return _inFlightMessageCount; } }
 
         /// <summary>
         /// Get the number of active async threads sending messages.
@@ -78,11 +84,9 @@ public Producer(IBrokerRouter brokerRouter, int maximumAsyncRequests = MaximumAs
         {
             BrokerRouter = brokerRouter;
             _maximumAsyncRequests = maximumAsyncRequests;
-            _maximumMessageBuffer = maximumMessageBuffer;
             _metadataQueries = new MetadataQueries(BrokerRouter);
             _asyncCollection = new AsyncCollection<TopicMessage>();
             _semaphoreMaximumAsync = new SemaphoreSlim(maximumAsyncRequests, maximumAsyncRequests);
-            _boundedCapacitySemaphore = new SemaphoreSlim(maximumMessageBuffer, maximumMessageBuffer);
 
             BatchSize = DefaultBatchSize;
             BatchDelayTime = TimeSpan.FromMilliseconds(DefaultBatchDelayMS);
@@ -104,10 +108,11 @@ public Producer(IBrokerRouter brokerRouter, int maximumAsyncRequests = MaximumAs
         /// <param name="timeout">Interal kafka timeout to wait for the requested level of ack to occur before returning. Defaults to 1000ms.</param>
         /// <param name="codec">The codec to apply to the message collection.  Defaults to none.</param>
         /// <returns>List of ProduceResponses from each partition sent to or empty list if acks = 0.</returns>
-        public async Task<List<ProduceResponse>> SendMessageAsync(string topic, IEnumerable<Message> messages, Int16 acks = 1,
+        public Task<List<ProduceResponse>> SendMessageAsync(string topic, IEnumerable<Message> messages, Int16 acks = 1,
             TimeSpan? timeout = null, MessageCodec codec = MessageCodec.CodecNone)
         {
-            if (_stopToken.IsCancellationRequested) throw new ObjectDisposedException("Cannot send new documents as producer is disposing.");
+            if (_stopToken.IsCancellationRequested)
+                throw new ObjectDisposedException("Cannot send new documents as producer is disposing.");
             if (timeout == null) timeout = TimeSpan.FromMilliseconds(DefaultAckTimeoutMS);
 
             var batch = messages.Select(message => new TopicMessage
@@ -119,20 +124,17 @@ public async Task<List<ProduceResponse>> SendMessageAsync(string topic, IEnumera
                 Message = message
             }).ToList();
 
-            foreach (var item in batch)
-            {
-                item.Tcs.Task.ContinueWith(t => _boundedCapacitySemaphore.Release(), TaskContinuationOptions.ExecuteSynchronously);
-                _boundedCapacitySemaphore.Wait(_stopToken.Token);
-                _asyncCollection.Add(item);
-            }
+            _asyncCollection.AddRange(batch);
 
-            var results = new List<ProduceResponse>();
-            foreach (var topicMessage in batch)
-            {
-                results.Add(await topicMessage.Tcs.Task.ConfigureAwait(false));
-            }
+            return Task.WhenAll(batch.Select(x => x.Tcs.Task))
+                        .ContinueWith(t =>
+                        {
+                            t.ThrowOnFault();
 
-            return results.Distinct().ToList();
+                            return batch.Select(topicMessage => topicMessage.Tcs.Task.Result)
+                                .Distinct()
+                                .ToList();
+                        });
         }
 
         /// <summary>
@@ -181,6 +183,8 @@ private async Task BatchSendAsync()
                 {
                     try
                     {
+                        await _asyncCollection.OnHasDataAvailable(_stopToken.Token).ConfigureAwait(false);
+
                         batch = await _asyncCollection.TakeAsync(BatchSize, BatchDelayTime, _stopToken.Token).ConfigureAwait(false);
                     }
                     catch (OperationCanceledException ex)
@@ -229,16 +233,18 @@ private async Task BatchSendAsync()
             }
         }
 
-        private async Task ProduceAndSendBatchAsync(List<TopicMessage> batchs, CancellationToken cancellationToken)
+        private async Task ProduceAndSendBatchAsync(List<TopicMessage> messages, CancellationToken cancellationToken)
         {
+            Interlocked.Add(ref _inFlightMessageCount, messages.Count);
+
             //we must send a different produce request for each ack level and timeout combination.
-            foreach (var ackLevelBatch in batchs.GroupBy(batch => new { batch.Acks, batch.Timeout }))
+            foreach (var ackLevelBatch in messages.GroupBy(batch => new { batch.Acks, batch.Timeout }))
             {
                 var messageByRouter = ackLevelBatch.Select(batch => new
-                                            {
-                                                TopicMessage = batch,
-                                                Route = BrokerRouter.SelectBrokerRoute(batch.Topic, batch.Message.Key),
-                                            })
+                {
+                    TopicMessage = batch,
+                    Route = BrokerRouter.SelectBrokerRoute(batch.Topic, batch.Message.Key),
+                })
                                          .GroupBy(x => new { x.Route, x.TopicMessage.Topic, x.TopicMessage.Codec });
 
                 var sendTasks = new List<BrokerRouteSendBatch>();
@@ -269,8 +275,7 @@ private async Task ProduceAndSendBatchAsync(List<TopicMessage> batchs, Cancellat
                     };
 
                     //ensure the async is released as soon as each task is completed
-                    brokerSendTask.Task.ContinueWith(t => { _semaphoreMaximumAsync.Release(); }, cancellationToken, 
-                        TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
+                    brokerSendTask.Task.ContinueWith(t => { _semaphoreMaximumAsync.Release(); }, cancellationToken);
 
                     sendTasks.Add(brokerSendTask);
                 }
@@ -293,11 +298,17 @@ private async Task ProduceAndSendBatchAsync(List<TopicMessage> batchs, Cancellat
                     {
                         foreach (var topicMessageBatch in ackLevelBatch)
                         {
-                            topicMessageBatch.Tcs.TrySetException(new KafkaApplicationException("An exception occured while executing a send operation against {0}.  Exception:{1}",
-                                failedTask.Route, failedTask.Task.Exception));
+                            topicMessageBatch.Tcs.TrySetException(
+                                new KafkaApplicationException(
+                                    "An exception occured while executing a send operation against {0}.  Exception:{1}",
+                                    failedTask.Route, failedTask.Task.Exception));
                         }
                     }
                 }
+                finally
+                {
+                    Interlocked.Add(ref _inFlightMessageCount, messages.Count * -1);
+                }
             }
         }
 
@@ -312,7 +323,7 @@ public void Dispose()
             using (_metadataQueries)
             {
             }
-        } 
+        }
         #endregion
     }