Producer target block

Author: kmcclellan

Confluent.Kafka.Dataflow/BlockOptions.cs

@@ -1,5 +1,6 @@
 namespace Confluent.Kafka.Dataflow
 {
+    using System;
     using System.Threading.Tasks.Dataflow;
 
     /// <summary>
@@ -19,4 +20,15 @@ public class ConsumeBlockOptions : DataflowBlockOptions
     public class OffsetBlockOptions : ExecutionDataflowBlockOptions
     {
     }
+
+    /// <summary>
+    /// Dataflow block options for producing to Kafka.
+    /// </summary>
+    public class ProduceBlockOptions : ExecutionDataflowBlockOptions
+    {
+        /// <summary>
+        /// Gets or sets the handler for produced message offsets.
+        /// </summary>
+        public Action<TopicPartitionOffset>? OffsetHandler { get; set; }
+    }
 }

Confluent.Kafka.Dataflow/ClientExtensions.cs

@@ -50,5 +50,28 @@ public static ITargetBlock<TopicPartitionOffset> AsOffsetBlock<TKey, TValue>(
                 x => consumer.StoreOffset(new TopicPartitionOffset(x.TopicPartition, x.Offset + 1)),
                 options ?? new());
         }
+
+        /// <summary>
+        /// Represents a producer as a target block for Kafka messages.
+        /// </summary>
+        /// <remarks>
+        /// A producer can be represented as multiple targets.
+        /// </remarks>
+        /// <typeparam name="TKey">The producer key type.</typeparam>
+        /// <typeparam name="TValue">The producer value type.</typeparam>
+        /// <param name="producer">The producer.</param>
+        /// <param name="topicPartition">The topic/partition receiving the messages. Use <see cref="Partition.Any"/> for automatic partitioning.</param>
+        /// <param name="options">Block options for producing.</param>
+        /// <returns>The producer target block.</returns>
+        public static ITargetBlock<Message<TKey, TValue>> AsTargetBlock<TKey, TValue>(
+            this IProducer<TKey, TValue> producer,
+            TopicPartition topicPartition,
+            ProduceBlockOptions? options = null)
+        {
+            return new ProduceBlock<TKey, TValue>(
+                producer ?? throw new ArgumentNullException(nameof(producer)),
+                topicPartition ?? throw new ArgumentNullException(nameof(topicPartition)),
+                options ?? new ProduceBlockOptions());
+        }
     }
 }

Confluent.Kafka.Dataflow/Internal/ProduceBlock.cs

@@ -0,0 +1,75 @@
+namespace Confluent.Kafka.Dataflow.Internal
+{
+    using System;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using System.Threading.Tasks.Dataflow;
+
+    class ProduceBlock<TKey, TValue> : ITargetBlock<Message<TKey, TValue>>
+    {
+        readonly ITargetBlock<Message<TKey, TValue>> target;
+        readonly Action<DeliveryReport<TKey, TValue>> deliveryHandler;
+
+        readonly IProducer<TKey, TValue> producer;
+        readonly TopicPartition topicPartition;
+        readonly ProduceBlockOptions options;
+
+        public ProduceBlock(IProducer<TKey, TValue> producer, TopicPartition topicPartition, ProduceBlockOptions options)
+        {
+            this.producer = producer;
+            this.topicPartition = topicPartition;
+            this.options = options;
+
+            this.deliveryHandler = this.OnDelivery;
+            this.target = new ActionBlock<Message<TKey, TValue>>(this.Produce, options);
+
+            var flushTask = this.target.Completion.ContinueWith(
+                (task, obj) =>
+                {
+                    var block = (ProduceBlock<TKey, TValue>)obj!;
+                    block.producer.Flush(block.options.CancellationToken);
+                },
+                this,
+                this.options.CancellationToken,
+                TaskContinuationOptions.OnlyOnRanToCompletion,
+                this.options.TaskScheduler);
+
+            this.Completion = Task.Factory.ContinueWhenAll(
+                new[] { this.target.Completion, flushTask },
+                t => t[0].IsFaulted ? t[0] : t[1],
+                CancellationToken.None,
+                TaskContinuationOptions.None,
+                this.options.TaskScheduler).Unwrap();
+        }
+
+        public Task Completion { get; }
+
+        public DataflowMessageStatus OfferMessage(
+            DataflowMessageHeader messageHeader,
+            Message<TKey, TValue> messageValue,
+            ISourceBlock<Message<TKey, TValue>>? source,
+            bool consumeToAccept)
+        {
+            return this.target.OfferMessage(messageHeader, messageValue, source, consumeToAccept);
+        }
+
+        void Produce(Message<TKey, TValue> message) =>
+            this.producer.Produce(this.topicPartition, message, this.deliveryHandler);
+
+        void OnDelivery(DeliveryReport<TKey, TValue> report)
+        {
+            if (report.Error.IsError)
+            {
+                this.target.Fault(new ProduceException<TKey, TValue>(report.Error, report));
+            }
+            else
+            {
+                this.options.OffsetHandler?.Invoke(report.TopicPartitionOffset);
+            }
+        }
+
+        public void Complete() => this.target.Complete();
+
+        public void Fault(Exception exception) => this.target.Fault(exception);
+    }
+}

README.md

@@ -110,3 +110,38 @@ var processor = new TransformBlock<Tuple<TopicPartitionOffset, Message<string, s
 var offsetTarget = consumer.AsOffsetBlock();
 processor.LinkTo(offsetTarget, linkOptions);
 ```
+
+### Producing using `ITargetBlock<T>`
+
+Use `IProducer<TKey, TValue>.AsTargetBlock(...)` to direct a message pipeline into a destination Kafka topic:
+
+```c#
+using System.Threading.Tasks.Dataflow;
+using Confluent.Kafka;
+using Confluent.Kafka.Dataflow;
+
+using var producer = new ProducerBuilder<string, string>(
+    new ProducerConfig
+    {
+        BootstrapServers = "localhost:9092",
+    }).Build();
+
+var target = producer.AsTargetBlock(new TopicPartition("my-topic", Partition.Any));
+
+var generator = new TransformBlock<int, Message<string, string>>(
+    i => new Message<string, string>
+    {
+        Key = i.ToString(),
+        Value = $"Value #{i}"
+    });
+
+generator.LinkTo(target, new DataflowLinkOptions { PropagateCompletion = true });
+
+for (var i = 0; i < 10; i++)
+{
+    generator.Post(i);
+}
+
+generator.Complete();
+await target.Completion;
+```