v2/feature/[Add support for Keys in Kafka Stream] #127:

src/Cortex.Streams.Kafka/Cortex.Streams.Kafka.csproj

@@ -4,8 +4,8 @@
 		<TargetFrameworks>net8.0</TargetFrameworks>
 		<Nullable>enable</Nullable>
 
-		<AssemblyVersion>1.0.1</AssemblyVersion>
-		<FileVersion>1.0.1</FileVersion>
+		<AssemblyVersion>2.0.0</AssemblyVersion>
+		<FileVersion>2.0.0</FileVersion>
 		<Product>Buildersoft Cortex Framework</Product>
 		<Company>Buildersoft</Company>
 		<Authors>Buildersoft,EnesHoxha</Authors>
@@ -16,7 +16,7 @@
 		<RepositoryUrl>https://github.com/buildersoftio/cortex</RepositoryUrl>
 		<PackageTags>cortex vortex eda streaming distributed streams states kafka pulsar rocksdb</PackageTags>
 
-		<Version>1.0.1</Version>
+		<Version>2.0.0</Version>
 		<PackageLicenseFile>license.md</PackageLicenseFile>
 		<PackageIcon>cortex.png</PackageIcon>
 		<PackageId>Cortex.Streams.Kafka</PackageId>
@@ -52,9 +52,9 @@
 
 
 	<ItemGroup>
-		<PackageReference Include="Confluent.Kafka" Version="2.9.0" />
-		<PackageReference Include="Google.Protobuf" Version="3.30.2" />
-		<PackageReference Include="protobuf-net" Version="3.2.46" />
+		<PackageReference Include="Confluent.Kafka" Version="2.11.0" />
+		<PackageReference Include="Google.Protobuf" Version="3.31.1" />
+		<PackageReference Include="protobuf-net" Version="3.2.55" />
 	</ItemGroup>
 
 	<ItemGroup>

src/Cortex.Streams.Kafka/KafkaKeyValueSinkOperator.cs

@@ -0,0 +1,65 @@
+using Confluent.Kafka;
+using Cortex.Streams.Kafka.Serializers;
+using Cortex.Streams.Operators;
+using System;
+using System.Collections.Generic;
+
+namespace Cortex.Streams.Kafka
+{
+    /// <summary>
+    /// Kafka sink that accepts KeyValuePair<TKey, TValue> so message keys are produced.
+    /// </summary>
+    public sealed class KafkaSinkOperator<TKey, TValue> : ISinkOperator<KeyValuePair<TKey, TValue>>
+    {
+        private readonly string _bootstrapServers;
+        private readonly string _topic;
+        private readonly IProducer<TKey, TValue> _producer;
+
+        public KafkaSinkOperator(
+                   string bootstrapServers,
+                   string topic,
+                   ProducerConfig config = null,
+                   ISerializer<TKey> keySerializer = null,
+                   ISerializer<TValue> valueSerializer = null)
+        {
+            _bootstrapServers = bootstrapServers ?? throw new ArgumentNullException(nameof(bootstrapServers));
+            _topic = topic ?? throw new ArgumentNullException(nameof(topic));
+
+            var producerConfig = config ?? new ProducerConfig
+            {
+                BootstrapServers = _bootstrapServers
+            };
+
+            keySerializer ??= new DefaultJsonSerializer<TKey>();
+            valueSerializer ??= new DefaultJsonSerializer<TValue>();
+
+            _producer = new ProducerBuilder<TKey, TValue>(producerConfig)
+                .SetKeySerializer(keySerializer)
+                .SetValueSerializer(valueSerializer)
+                .Build();
+        }
+
+        public void Process(KeyValuePair<TKey, TValue> input)
+        {
+            var msg = new Message<TKey, TValue> { Key = input.Key, Value = input.Value };
+            _producer.Produce(_topic, msg, deliveryReport =>
+            {
+                if (deliveryReport.Error.IsError)
+                {
+                    Console.WriteLine($"Delivery Error: {deliveryReport.Error.Reason}");
+                }
+            });
+        }
+
+        public void Start()
+        {
+            // no-op
+        }
+
+        public void Stop()
+        {
+            _producer.Flush(TimeSpan.FromSeconds(10));
+            _producer.Dispose();
+        }
+    }
+}

src/Cortex.Streams.Kafka/KafkaKeyValueSourceOperator.cs

@@ -0,0 +1,97 @@
+using Confluent.Kafka;
+using Cortex.Streams.Kafka.Deserializers;
+using Cortex.Streams.Operators;
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Cortex.Streams.Kafka
+{
+    /// <summary>
+    /// Kafka source that emits KeyValuePair<TKey, TValue> so the pipeline can use message keys.
+    /// </summary>
+    public sealed class KafkaSourceOperator<TKey, TValue> : ISourceOperator<KeyValuePair<TKey, TValue>>
+    {
+        private readonly string _bootstrapServers;
+        private readonly string _topic;
+        private readonly IConsumer<TKey, TValue> _consumer;
+        private CancellationTokenSource _cts;
+        private Task _consumeTask;
+
+
+        public KafkaSourceOperator(string bootstrapServers,
+            string topic,
+            ConsumerConfig config = null,
+            IDeserializer<TKey> keyDeserializer = null,
+            IDeserializer<TValue> valueDeserializer = null)
+        {
+            _bootstrapServers = bootstrapServers ?? throw new ArgumentNullException(nameof(bootstrapServers));
+            _topic = topic ?? throw new ArgumentNullException(nameof(topic));
+
+            var consumerConfig = config ?? new ConsumerConfig
+            {
+                BootstrapServers = _bootstrapServers,
+                GroupId = Guid.NewGuid().ToString(),
+                AutoOffsetReset = AutoOffsetReset.Earliest,
+                EnableAutoCommit = true,
+            };
+
+            keyDeserializer ??= new DefaultJsonDeserializer<TKey>();
+            valueDeserializer ??= new DefaultJsonDeserializer<TValue>();
+
+            _consumer = new ConsumerBuilder<TKey, TValue>(consumerConfig)
+                .SetKeyDeserializer(keyDeserializer)
+                .SetValueDeserializer(valueDeserializer)
+                .Build();
+        }
+
+
+        public void Start(Action<KeyValuePair<TKey, TValue>> emit)
+        {
+            if (emit == null) throw new ArgumentNullException(nameof(emit));
+
+            _cts = new CancellationTokenSource();
+            _consumer.Subscribe(_topic);
+
+            _consumeTask = Task.Run(() =>
+            {
+                try
+                {
+                    while (!_cts.Token.IsCancellationRequested)
+                    {
+                        var result = _consumer.Consume(_cts.Token);
+                        emit(new KeyValuePair<TKey, TValue>(result.Message.Key, result.Message.Value));
+                    }
+                }
+                catch (OperationCanceledException)
+                {
+                    // shutting down - consume loop canceled
+                }
+                finally
+                {
+                    _consumer.Close();
+                }
+            }, _cts.Token);
+        }
+
+        public void Stop()
+        {
+            if (_cts == null)
+                return;
+
+            _cts.Cancel();
+            try
+            {
+                _consumeTask?.Wait();
+            }
+            catch
+            {
+                /* swallow aggregate canceled */
+            }
+
+            _consumer.Dispose();
+            _cts.Dispose();
+        }
+    }
+}

src/Cortex.Streams.Kafka/KafkaSinkOperator.cs

@@ -5,7 +5,7 @@
 
 namespace Cortex.Streams.Kafka
 {
-    public class KafkaSinkOperator<TInput> : ISinkOperator<TInput>
+    public sealed class KafkaSinkOperator<TInput> : ISinkOperator<TInput>
     {
         private readonly string _bootstrapServers;
         private readonly string _topic;

src/Cortex.Streams.Kafka/KafkaSourceOperator.cs

@@ -7,15 +7,18 @@
 
 namespace Cortex.Streams.Kafka
 {
-    public class KafkaSourceOperator<TOutput> : ISourceOperator<TOutput>
+    public sealed class KafkaSourceOperator<TOutput> : ISourceOperator<TOutput>
     {
         private readonly string _bootstrapServers;
         private readonly string _topic;
         private readonly IConsumer<Ignore, TOutput> _consumer;
         private CancellationTokenSource _cts;
         private Task _consumeTask;
 
-        public KafkaSourceOperator(string bootstrapServers, string topic, ConsumerConfig config = null, IDeserializer<TOutput> deserializer = null)
+        public KafkaSourceOperator(string bootstrapServers,
+            string topic,
+            ConsumerConfig config = null,
+            IDeserializer<TOutput> deserializer = null)
         {
             _bootstrapServers = bootstrapServers;
             _topic = topic;
@@ -53,7 +56,7 @@ public void Start(Action<TOutput> emit)
                 }
                 catch (OperationCanceledException)
                 {
-                    // Consume loop canceled
+                    // shutting down - consume loop canceled
                 }
                 finally
                 {
@@ -64,8 +67,21 @@ public void Start(Action<TOutput> emit)
 
         public void Stop()
         {
+            if (_cts == null)
+                return;
+
             _cts.Cancel();
-            _consumeTask.Wait();
+            try
+            {
+                _consumeTask?.Wait();
+            }
+            catch
+            {
+                /* swallow aggregate canceled */
+            }
+
+            _consumer.Dispose();
+            _cts.Dispose();
         }
     }
 }