feat: implement kafka headers support

Author: janiodev

src/Take.Elephant.Kafka/IKafkaHeaderProvider.cs

@@ -0,0 +1,15 @@
+using System.Collections.Generic;
+using Confluent.Kafka;
+
+namespace Take.Elephant.Kafka
+{
+    /// <summary>
+    /// Provides Kafka message headers to be injected into every produced message.
+    /// Implementations are resolved at produce-time, so headers can be dynamic
+    /// (e.g. trace IDs) or static (e.g. deployment metadata read once at startup).
+    /// </summary>
+    public interface IKafkaHeaderProvider
+    {
+        IEnumerable<IHeader> GetHeaders();
+    }
+}

src/Take.Elephant.Kafka/IKafkaReceiverQueue.cs

@@ -0,0 +1,26 @@
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Take.Elephant.Kafka
+{
+    /// <summary>
+    /// Optional Kafka receiver contract that allows reading message headers together with the payload.
+    /// </summary>
+    /// <typeparam name="T">Payload type.</typeparam>
+    public interface IKafkaReceiverQueue<T> : IReceiverQueue<T>, IBlockingReceiverQueue<T>
+    {
+        /// <summary>
+        /// Dequeues a payload with its Kafka headers, if available.
+        /// </summary>
+        Task<KafkaConsumedMessage<T>> DequeueWithHeadersOrDefaultAsync(
+            CancellationToken cancellationToken = default
+        );
+
+        /// <summary>
+        /// Dequeues a payload with its Kafka headers, awaiting for a message if needed.
+        /// </summary>
+        Task<KafkaConsumedMessage<T>> DequeueWithHeadersAsync(
+            CancellationToken cancellationToken
+        );
+    }
+}

src/Take.Elephant.Kafka/KafkaConsumedMessage.cs

@@ -0,0 +1,49 @@
+using System.Collections.Generic;
+using System.Collections.ObjectModel;
+using System.Text;
+
+namespace Take.Elephant.Kafka
+{
+    internal static class KafkaConsumedMessageDefaults
+    {
+        internal static readonly IReadOnlyDictionary<string, byte[]> EmptyHeaders =
+            new ReadOnlyDictionary<string, byte[]>(new Dictionary<string, byte[]>());
+    }
+
+    /// <summary>
+    /// Represents a consumed Kafka payload with its message headers.
+    /// </summary>
+    /// <typeparam name="T">Payload type.</typeparam>
+    public sealed class KafkaConsumedMessage<T>(T item, IReadOnlyDictionary<string, byte[]> headers)
+    {
+
+        private readonly IReadOnlyDictionary<string, byte[]> _headers =
+             headers is null
+                 ? KafkaConsumedMessageDefaults.EmptyHeaders
+                 : new ReadOnlyDictionary<string, byte[]>(new Dictionary<string, byte[]>(headers));
+
+        public T Item { get; } = item;
+
+        public IReadOnlyDictionary<string, byte[]> Headers => _headers;
+
+        public bool TryGetHeader(string key, out byte[] value)
+        {
+            if (!string.IsNullOrWhiteSpace(key)) return Headers.TryGetValue(key, out value);
+            value = null;
+            return false;
+
+        }
+
+        public bool TryGetHeaderAsUtf8String(string key, out string value)
+        {
+            value = null;
+            if (!TryGetHeader(key, out var bytes) || bytes == null)
+            {
+                return false;
+            }
+
+            value = Encoding.UTF8.GetString(bytes);
+            return true;
+        }
+    }
+}

src/Take.Elephant.Kafka/KafkaEventStreamPublisher.cs

@@ -1,4 +1,4 @@
-using System;
+using System;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
@@ -9,6 +9,7 @@ namespace Take.Elephant.Kafka
     public class KafkaEventStreamPublisher<TKey, TEvent> : IEventStreamPublisher<TKey, TEvent>, IDisposable
     {
         private readonly IProducer<TKey, TEvent> _producer;
+        private readonly IKafkaHeaderProvider _headerProvider;
 
         public KafkaEventStreamPublisher(string bootstrapServers, string topic, ISerializer<TEvent> serializer)
             : this(new ProducerConfig() { BootstrapServers = bootstrapServers }, topic, serializer)
@@ -21,14 +22,29 @@ public KafkaEventStreamPublisher(string bootstrapServers, string topic, Confluen
         }
 
         public KafkaEventStreamPublisher(
-            ProducerConfig producerConfig,            
+            ProducerConfig producerConfig,
             string topic,
             ISerializer<TEvent> serializer)
             : this(
                   new ProducerBuilder<TKey, TEvent>(producerConfig)
                         .SetValueSerializer(new EventSerializer(serializer))
                         .Build(),
-                  topic)
+                  topic,
+                  null)
+        {
+        }
+
+        public KafkaEventStreamPublisher(
+            ProducerConfig producerConfig,
+            string topic,
+            ISerializer<TEvent> serializer,
+            IKafkaHeaderProvider headerProvider)
+            : this(
+                  new ProducerBuilder<TKey, TEvent>(producerConfig)
+                        .SetValueSerializer(new EventSerializer(serializer))
+                        .Build(),
+                  topic,
+                  headerProvider)
         {
         }
 
@@ -40,7 +56,22 @@ public KafkaEventStreamPublisher(
                   new ProducerBuilder<TKey, TEvent>(producerConfig)
                         .SetValueSerializer(kafkaSerializer)
                         .Build(),
-                  topic)
+                  topic,
+                  null)
+        {
+        }
+
+        public KafkaEventStreamPublisher(
+            ProducerConfig producerConfig,
+            string topic,
+            Confluent.Kafka.ISerializer<TEvent> kafkaSerializer,
+            IKafkaHeaderProvider headerProvider)
+            : this(
+                  new ProducerBuilder<TKey, TEvent>(producerConfig)
+                        .SetValueSerializer(kafkaSerializer)
+                        .Build(),
+                  topic,
+                  headerProvider)
         {
         }
 
@@ -53,22 +84,64 @@ public KafkaEventStreamPublisher(
                 throw new ArgumentException("Value cannot be null or whitespace.", nameof(topic));
             }
 
+            if (producer == null)
+             {
+                 throw new ArgumentNullException(nameof(producer));
+             }
+
+            _producer = producer;
+            Topic = topic;
+            _headerProvider = null;
+        }
+
+        public KafkaEventStreamPublisher(
+            IProducer<TKey, TEvent> producer,
+            string topic,
+            IKafkaHeaderProvider headerProvider)
+        {
+            if (string.IsNullOrWhiteSpace(topic))
+            {
+                throw new ArgumentException("Value cannot be null or whitespace.", nameof(topic));
+            }
+
+            if (producer == null)
+             {
+                 throw new ArgumentNullException(nameof(producer));
+             }
+
             _producer = producer;
             Topic = topic;
+            _headerProvider = headerProvider;
         }
 
         public string Topic { get; }
 
         public async Task PublishAsync(TKey key, TEvent item, CancellationToken cancellationToken)
         {
-            await _producer.ProduceAsync(
-                Topic,
-                new Message<TKey, TEvent>
+            var message = new Message<TKey, TEvent>
+            {
+                Key = key,
+                Value = item
+            };
+
+            if (_headerProvider != null)
+            {
+                message.Headers = [];
+                var headers = _headerProvider.GetHeaders();
+                 if (headers != null)
                 {
-                    Key = key,
-                    Value = item
-                },
-                cancellationToken);
+                    foreach (var header in headers)
+                     {
+                         if (header == null || header.Key == null)
+                         {
+                             continue;
+                         }
+                         message.Headers.Add(header.Key, header.GetValueBytes());
+                     }
+                }
+            }
+
+            await _producer.ProduceAsync(Topic, message, cancellationToken);
         }
 
         protected virtual void Dispose(bool disposing)
@@ -85,18 +158,12 @@ public void Dispose()
             GC.SuppressFinalize(this);
         }
 
-        public class EventSerializer : Confluent.Kafka.ISerializer<TEvent>
+        public class EventSerializer(ISerializer<TEvent> serializer) : Confluent.Kafka.ISerializer<TEvent>
         {
-            private readonly ISerializer<TEvent> _serializer;
-
-            public EventSerializer(ISerializer<TEvent> serializer)
-            {
-                _serializer = serializer;
-            }
             public byte[] Serialize(TEvent data, SerializationContext context)
             {
-                return Encoding.UTF8.GetBytes(_serializer.Serialize(data));
+                return Encoding.UTF8.GetBytes(serializer.Serialize(data));
             }
         }
     }
-}
+}

src/Take.Elephant.Kafka/KafkaHeadersConverter.cs

@@ -0,0 +1,35 @@
+using System.Collections.Generic;
+using Confluent.Kafka;
+
+namespace Take.Elephant.Kafka
+{
+    internal static class KafkaHeadersConverter
+    {
+        internal static KafkaConsumedMessage<T> BuildConsumedMessage<T>(T item, Headers headers)
+        {
+            return new KafkaConsumedMessage<T>(item, ToDictionary(headers));
+        }
+
+        internal static Dictionary<string, byte[]> ToDictionary(Headers headers)
+        {
+            if (headers == null || headers.Count == 0)
+            {
+                return null;
+            }
+
+            var result = new Dictionary<string, byte[]>(headers.Count);
+            foreach (var header in headers)
+            {
+                if (header?.Key == null)
+                {
+                    continue;
+                }
+
+                var valueBytes = header.GetValueBytes();
+                result[header.Key] = valueBytes != null ? (byte[])valueBytes.Clone() : null;
+            }
+
+            return result;
+        }
+    }
+}

src/Take.Elephant.Kafka/KafkaPartitionQueue.cs

@@ -1,13 +1,11 @@
 using System;
-using System.Collections.Generic;
-using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using Confluent.Kafka;
 
 namespace Take.Elephant.Kafka
 {
-    public class KafkaPartitionQueue<T> : IReceiverQueue<T>, IPartitionSenderQueue<T>, ICloseable, IDisposable
+    public class KafkaPartitionQueue<T> : IKafkaReceiverQueue<T>, IPartitionSenderQueue<T>, ICloseable, IDisposable
     {
         private readonly KafkaPartitionSenderQueue<T> _senderQueue;
         private readonly KafkaReceiverQueue<T> _receiverQueue;
@@ -16,12 +14,13 @@ public KafkaPartitionQueue(
             ProducerConfig producerConfig,
             ConsumerConfig consumerConfig,
             string topic,
-            Take.Elephant.ISerializer<T> serializer,
+            ISerializer<T> serializer,
             Confluent.Kafka.ISerializer<string> kafkaSerializer = null,
-            IDeserializer<string> kafkaDeserializer = null)
+            IDeserializer<string> kafkaDeserializer = null,
+            IKafkaHeaderProvider headerProvider = null)
         {
-            _senderQueue = new KafkaPartitionSenderQueue<T>(producerConfig, topic, serializer, kafkaSerializer);
-            _receiverQueue = new KafkaReceiverQueue<T>(consumerConfig, topic, serializer, kafkaDeserializer);
+            _senderQueue = new(producerConfig, topic, serializer, kafkaSerializer, headerProvider);
+            _receiverQueue = new(consumerConfig, topic, serializer, kafkaDeserializer);
         }
 
         public Task EnqueueAsync(T item, string key, CancellationToken cancellationToken = default)
@@ -34,11 +33,25 @@ public Task<T> DequeueOrDefaultAsync(CancellationToken cancellationToken = defau
             return _receiverQueue.DequeueOrDefaultAsync(cancellationToken);
         }
 
-        public Task<T> DequeueAsync(CancellationToken cancellationToken = default)
+        public Task<KafkaConsumedMessage<T>> DequeueWithHeadersOrDefaultAsync(
+            CancellationToken cancellationToken = default
+        )
+        {
+            return _receiverQueue.DequeueWithHeadersOrDefaultAsync(cancellationToken);
+        }
+
+        public Task<T> DequeueAsync(CancellationToken cancellationToken)
         {
             return _receiverQueue.DequeueAsync(cancellationToken);
         }
 
+        public Task<KafkaConsumedMessage<T>> DequeueWithHeadersAsync(
+            CancellationToken cancellationToken
+        )
+        {
+            return _receiverQueue.DequeueWithHeadersAsync(cancellationToken);
+        }
+
         public Task CloseAsync(CancellationToken cancellationToken) => _receiverQueue.CloseAsync(cancellationToken);
 
         public void Dispose()
@@ -47,13 +60,11 @@ public void Dispose()
             GC.SuppressFinalize(this);
         }
 
-        protected void Dispose(bool disposing)
+        protected virtual void Dispose(bool disposing)
         {
-            if (disposing)
-            {
-                _senderQueue?.Dispose();
-                _receiverQueue?.Dispose();
-            }
+            if (!disposing) return;
+            _senderQueue?.Dispose();
+            _receiverQueue?.Dispose();
         }
     }
 }