Add a UTF-8 statistics handler

Author: verdie-g

src/Confluent.Kafka/Consumer.cs

@@ -17,6 +17,7 @@
 // Refer to LICENSE for more information.
 
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.Linq;
 using System.Runtime.InteropServices;
@@ -39,6 +40,9 @@ internal class Config
             internal IEnumerable<KeyValuePair<string, string>> config;
             internal Action<Error> errorHandler;
             internal Action<LogMessage> logHandler;
+#if NET6_0_OR_GREATER
+            internal ReadOnlySpanAction<byte, object> statisticsUtf8Handler;
+#endif
             internal Action<string> statisticsHandler;
             internal Action<CommittedOffsets> offsetsCommittedHandler;
             internal Action<string> oAuthBearerTokenRefreshHandler;
@@ -104,6 +108,9 @@ private void ErrorCallback(IntPtr rk, ErrorCode err, string reason, IntPtr opaqu
             }
         }
 
+#if NET6_0_OR_GREATER
+        private ReadOnlySpanAction<byte, object> statisticsUtf8Handler;
+#endif
         private Action<string> statisticsHandler;
         private Librdkafka.StatsDelegate statisticsCallbackDelegate;
         private int StatisticsCallback(IntPtr rk, IntPtr json, UIntPtr json_len, IntPtr opaque)
@@ -112,7 +119,20 @@ private int StatisticsCallback(IntPtr rk, IntPtr json, UIntPtr json_len, IntPtr
             try
             {
                 // Ensure registered handlers are never called as a side-effect of Dispose/Finalize (prevents deadlocks in common scenarios).
-                statisticsHandler?.Invoke(Util.Marshal.PtrToStringUTF8(json, json_len));
+#if NET6_0_OR_GREATER
+                if (statisticsUtf8Handler != null)
+                {
+                    unsafe
+                    {
+                        statisticsUtf8Handler.Invoke(new ReadOnlySpan<byte>(json.ToPointer(), (int)json_len), null);
+                    }
+                }
+#endif
+
+                if (statisticsHandler != null)
+                {
+                    statisticsHandler.Invoke(Util.Marshal.PtrToStringUTF8(json, json_len));
+                }
             }
             catch (Exception e)
             {
@@ -645,6 +665,9 @@ internal Consumer(ConsumerBuilder<TKey, TValue> builder)
         {
             var baseConfig = builder.ConstructBaseConfig(this);
 
+#if NET6_0_OR_GREATER
+            this.statisticsUtf8Handler = baseConfig.statisticsUtf8Handler;
+#endif
             this.statisticsHandler = baseConfig.statisticsHandler;
             this.logHandler = baseConfig.logHandler;
             this.errorHandler = baseConfig.errorHandler;
@@ -727,7 +750,11 @@ internal Consumer(ConsumerBuilder<TKey, TValue> builder)
             {
                 Librdkafka.conf_set_log_cb(configPtr, logCallbackDelegate);
             }
+#if NET6_0_OR_GREATER
+            if (statisticsUtf8Handler != null || statisticsHandler != null)
+#else
             if (statisticsHandler != null)
+#endif
             {
                 Librdkafka.conf_set_stats_cb(configPtr, statisticsCallbackDelegate);
             }

src/Confluent.Kafka/ConsumerBuilder.cs

@@ -15,6 +15,7 @@
 // Refer to LICENSE for more information.
 
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.Linq;
 
@@ -41,8 +42,18 @@ public class ConsumerBuilder<TKey, TValue>
         /// </summary>
         internal protected Action<IConsumer<TKey, TValue>, LogMessage> LogHandler { get; set; }
 
+#if NET6_0_OR_GREATER
         /// <summary>
-        ///     The configured statistics handler.
+        ///     The configured statistics handler. Unlike <see cref="StatisticsHandler"/>, this handler gives access
+        ///     to the raw UTF-8 which is more performance friendly to UTF-8 parsers such as System.Text.Json. Also,
+        ///     if <see cref="StatisticsHandler"/> is not set, the JSON string allocation is completely avoided.
+        /// </summary>
+        internal protected ReadOnlySpanAction<byte, IConsumer<TKey, TValue>> StatisticsUtf8Handler { get; set; }
+#endif
+
+        /// <summary>
+        ///     The configured statistics handler. With .NET 6+, prefer the <see cref="StatisticsUtf8Handler"/> which
+        ///     doesn't allocate the JSON string and gives access to the raw UTF-8 bytes.
         /// </summary>
         internal protected Action<IConsumer<TKey, TValue>, string> StatisticsHandler { get; set; }
 
@@ -98,6 +109,11 @@ internal Consumer<TKey,TValue>.Config ConstructBaseConfig(Consumer<TKey, TValue>
                 logHandler = this.LogHandler == null
                     ? default(Action<LogMessage>)
                     : logMessage => this.LogHandler(consumer, logMessage),
+#if NET6_0_OR_GREATER
+                statisticsUtf8Handler = this.StatisticsUtf8Handler == null
+                    ? default(ReadOnlySpanAction<byte, object>)
+                    : (stats, _) => this.StatisticsUtf8Handler(stats, consumer),
+#endif
                 statisticsHandler = this.StatisticsHandler == null
                     ? default(Action<string>)
                     : stats => this.StatisticsHandler(consumer, stats),
@@ -136,6 +152,24 @@ public ConsumerBuilder(IEnumerable<KeyValuePair<string, string>> config)
             this.Config = config;
         }
 
+#if NET6_0_OR_GREATER
+        /// <summary>
+        ///     Set the handler to call on statistics events. Unlike <see cref="SetStatisticsHandler"/>, this handler
+        ///     gives access to the raw UTF-8 which is more performance friendly to UTF-8 parsers such as
+        ///     System.Text.Json. Also, it doesn't allocate a potentially large string for the JSON.
+        /// </summary>
+        public ConsumerBuilder<TKey, TValue> SetStatisticsUtf8Handler(
+            ReadOnlySpanAction<byte, IConsumer<TKey, TValue>> statisticsHandler)
+        {
+            if (this.StatisticsUtf8Handler != null)
+            {
+                throw new InvalidOperationException("Statistics handler may not be specified more than once.");
+            }
+            this.StatisticsUtf8Handler = statisticsHandler;
+            return this;
+        }
+#endif
+
         /// <summary>
         ///     Set the handler to call on statistics events. Statistics 
         ///     are provided as a JSON formatted string as defined here:

test/Confluent.Kafka.IntegrationTests/Tests/Transactions_Statistics.cs

@@ -15,6 +15,7 @@
 // Refer to LICENSE for more information.
 
 using System;
+using System.Text.Json;
 using System.Threading;
 using Newtonsoft.Json.Linq;
 using Xunit;
@@ -90,5 +91,69 @@ public void Transactions_Statistics(string bootstrapServers)
             Assert.Equal(0, Library.HandleCount);
             LogToFile("end   Transactions_Statistics");
         }
+        
+        [Theory, MemberData(nameof(KafkaParameters))]
+        public void Transactions_Statistics_Utf8(string bootstrapServers)
+        {
+            LogToFile("start Transactions_Statistics");
+
+            var groupName = Guid.NewGuid().ToString();
+
+            var cConfig = new ConsumerConfig
+            {
+                IsolationLevel = IsolationLevel.ReadCommitted,
+                BootstrapServers = bootstrapServers,
+                GroupId = groupName,
+                EnableAutoCommit = false,
+                StatisticsIntervalMs = 1000
+            };
+
+            int ls_offset = -1;
+            int hi_offset = -1;
+            bool done = false;
+
+            using (var topic = new TemporaryTopic(bootstrapServers, 1))
+            using (var producer = new TestProducerBuilder<string, string>(new ProducerConfig { BootstrapServers = bootstrapServers, TransactionalId = Guid.NewGuid().ToString(), LingerMs = 0 }).Build())
+            using (var consumer = new TestConsumerBuilder<string, string>(cConfig)
+                .SetStatisticsUtf8Handler((jsonBytes, _) =>
+                {
+                    var stats = JsonSerializer.Deserialize<JsonElement>(jsonBytes);
+                    var firstPartitionStats = stats.GetProperty("topics").GetProperty(topic.Name)
+                        .GetProperty("partitions").GetProperty("0");
+                    ls_offset = firstPartitionStats.GetProperty("ls_offset").GetInt32();
+                    hi_offset = firstPartitionStats.GetProperty("hi_offset").GetInt32();
+                    if (hi_offset > 4) { done = true; }
+                })
+                .Build())
+            {
+                consumer.Assign(new TopicPartitionOffset(topic.Name, 0, 0));
+
+                producer.InitTransactions(TimeSpan.FromSeconds(30));
+                producer.BeginTransaction();
+                producer.ProduceAsync(topic.Name, new Message<string, string> { Key = "test", Value = "message_a" }).Wait();
+                producer.CommitTransaction(TimeSpan.FromSeconds(30));
+
+                producer.BeginTransaction();
+                producer.ProduceAsync(topic.Name, new Message<string, string> { Key = "test", Value = "message_b" }).Wait();
+                producer.CommitTransaction(TimeSpan.FromSeconds(30));
+
+                producer.BeginTransaction();
+                producer.ProduceAsync(topic.Name, new Message<string, string> { Key = "test", Value = "message1" }).Wait();
+                producer.ProduceAsync(topic.Name, new Message<string, string> { Key = "test", Value = "message2" }).Wait();
+                producer.ProduceAsync(topic.Name, new Message<string, string> { Key = "test", Value = "message3" }).Wait();
+
+                for (int i=0; i<10; ++i)
+                {
+                    consumer.Consume(TimeSpan.FromMilliseconds(500));
+                    if (done) { break; }
+                }
+
+                Assert.Equal(4, ls_offset);
+                Assert.Equal(7, hi_offset);
+            }
+
+            Assert.Equal(0, Library.HandleCount);
+            LogToFile("end   Transactions_Statistics");
+        }
     }
 }