Publisher confirms as PublishMiddleware

src/RabbitMQ.Next.Abstractions/Channels/IChannel.cs

@@ -21,7 +21,7 @@ Task SendAsync<TRequest>(TRequest request, CancellationToken cancellation = defa
         where TRequest : struct, IOutgoingMethod
         where TResponse : struct, IIncomingMethod;
 
-    Task PublishAsync<TState>(TState contentBuilderState, string exchange, string routingKey, IMessageProperties properties, Action<TState, IBufferWriter<byte>> payload, PublishFlags flags = PublishFlags.None, CancellationToken cancellation = default);
+    Task<ulong> PublishAsync<TState>(TState contentBuilderState, string exchange, string routingKey, IMessageProperties properties, Action<TState, IBufferWriter<byte>> payload, PublishFlags flags = PublishFlags.None, CancellationToken cancellation = default);
 
     Task<TMethod> WaitAsync<TMethod>(CancellationToken cancellation = default)
         where TMethod : struct, IIncomingMethod;

src/RabbitMQ.Next.Consumer.Abstractions/IConsumerMiddleware.cs

@@ -0,0 +1,9 @@
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace RabbitMQ.Next.Consumer;
+
+public interface IConsumerMiddleware
+{
+    ValueTask InvokeAsync(IDeliveredMessage message, CancellationToken cancellation);
+}

src/RabbitMQ.Next.Publisher.Abstractions/IPublishMiddleware.cs

@@ -1,9 +1,12 @@
 using System.Threading;
 using System.Threading.Tasks;
+using RabbitMQ.Next.Channels;
 
 namespace RabbitMQ.Next.Publisher;
 
 public interface IPublishMiddleware
 {
-    ValueTask InvokeAsync<TContent>(TContent content, IMessageBuilder message, CancellationToken cancellation);
+    ValueTask InitAsync(IChannel channel, CancellationToken cancellation);
+
+    ValueTask<ulong> InvokeAsync<TContent>(TContent content, IMessageBuilder message, CancellationToken cancellation);
 }

src/RabbitMQ.Next.Publisher.Abstractions/IPublisher.cs

@@ -8,5 +8,5 @@ public interface IPublisher : IAsyncDisposable
 {
     Task PublishAsync<TContent>(TContent content, Action<IMessageBuilder> propertiesBuilder = null, CancellationToken cancellation = default);
 
-    Task PublishAsync<TState, TContent>(TState state, TContent content, Action<TState, IMessageBuilder> propertiesBuilder = null, CancellationToken cancellation = default);
+    Task PublishAsync<TState, TContent>(TState state, TContent content, Action<TState, IMessageBuilder> propertiesBuilder, CancellationToken cancellation = default);
 }

src/RabbitMQ.Next.Publisher.Abstractions/IPublisherBuilder.cs

@@ -8,6 +8,4 @@ public interface IPublisherBuilder: ISerializationBuilder<IPublisherBuilder>
     IPublisherBuilder UsePublishMiddleware(Func<IPublishMiddleware, IPublishMiddleware> middlewareFactory);
 
     IPublisherBuilder UseReturnMiddleware(Func<IReturnMiddleware,IReturnMiddleware> middlewareFactory);
-
-    IPublisherBuilder NoConfirms();
 }

src/RabbitMQ.Next.Publisher.Abstractions/PublisherBuilderExtensions.cs

@@ -4,9 +4,6 @@ namespace RabbitMQ.Next.Publisher;
 
 public static class PublisherBuilderExtensions
 {
-    public static IPublisherBuilder UsePublishMiddleware(this IPublisherBuilder builder, Action<object, IMessageBuilder> middleware)
-        => builder.UsePublishMiddleware(p => new DelegatePublishMiddleware(p, middleware));
-
     public static IPublisherBuilder UseReturnMiddleware(this IPublisherBuilder builder, Action<IReturnedMessage> middleware)
         => builder.UseReturnMiddleware(p => new DelegateReturnMiddleware(p, middleware));
 }

src/RabbitMQ.Next.Publisher.Attributes/AttributePublishMiddleware.cs

@@ -4,6 +4,7 @@
 using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
+using RabbitMQ.Next.Channels;
 
 namespace RabbitMQ.Next.Publisher.Attributes;
 
@@ -42,7 +43,9 @@ private static IReadOnlyList<MessageAttributeBase> AsMessageAttributes(object[]
         return typed;
     }
 
-    public ValueTask InvokeAsync<TContent>(TContent content, IMessageBuilder message, CancellationToken cancellation)
+    public ValueTask InitAsync(IChannel channel, CancellationToken cancellation) => default;
+
+    public ValueTask<ulong> InvokeAsync<TContent>(TContent content, IMessageBuilder message, CancellationToken cancellation)
     {
         var type = typeof(TContent);
 

src/RabbitMQ.Next.Publisher/ConfirmMessageHandler.cs

@@ -2,9 +2,11 @@
 using System.Collections.Concurrent;
 using System.Linq;
 using System.Runtime.CompilerServices;
+using System.Threading;
 using System.Threading.Tasks;
 using RabbitMQ.Next.Channels;
 using RabbitMQ.Next.Messaging;
+using RabbitMQ.Next.Tasks;
 using RabbitMQ.Next.Transport.Methods.Basic;
 
 namespace RabbitMQ.Next.Publisher;
@@ -25,15 +27,15 @@ static ConfirmMessageHandler()
         NegativeCompletedTcs.SetResult(false);
     }
 
-    public Task<bool> WaitForConfirmAsync(ulong deliveryTag)
+    public Task<bool> WaitForConfirmAsync(ulong deliveryTag, CancellationToken cancellation)
     {
         var tcs = this.pendingConfirms.GetOrAdd(deliveryTag, _ => new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously));
         if (tcs.Task.IsCompleted)
         {
             this.pendingConfirms.TryRemove(deliveryTag, out _);
         }
 
-        return tcs.Task;
+        return tcs.Task.WithCancellation(cancellation);
     }
 
     public void Handle(AckMethod method, IPayload payload)

src/RabbitMQ.Next.Publisher/ConfirmMessageMiddleware.cs

@@ -0,0 +1,41 @@
+using System.Threading;
+using System.Threading.Tasks;
+using RabbitMQ.Next.Channels;
+using RabbitMQ.Next.Transport.Methods.Basic;
+using RabbitMQ.Next.Transport.Methods.Confirm;
+
+namespace RabbitMQ.Next.Publisher;
+
+internal class ConfirmMessageMiddleware : IPublishMiddleware
+{
+    private readonly IPublishMiddleware next;
+    private ConfirmMessageHandler confirms;
+
+    public ConfirmMessageMiddleware(IPublishMiddleware next)
+    {
+        this.next = next;
+    }
+
+    public ValueTask InitAsync(IChannel channel, CancellationToken cancellation)
+    {
+        this.confirms = new ConfirmMessageHandler();
+        channel.WithMessageHandler<AckMethod>(this.confirms);
+        channel.WithMessageHandler<NackMethod>(this.confirms);
+
+        return new ValueTask(channel.SendAsync<SelectMethod, SelectOkMethod>(new SelectMethod(), cancellation));
+    }
+
+    public async ValueTask<ulong> InvokeAsync<TContent>(TContent content, IMessageBuilder message, CancellationToken cancellation)
+    {
+        var deliveryTag = await this.next.InvokeAsync(content, message, cancellation);
+
+        var confirmed = await this.confirms.WaitForConfirmAsync(deliveryTag, cancellation);
+        if (!confirmed)
+        {
+            // todo: provide some useful info here
+            throw new DeliveryFailedException();
+        }
+
+        return deliveryTag;
+    }
+}

src/RabbitMQ.Next.Publisher/ConnectionExtensions.cs

@@ -5,9 +5,6 @@ namespace RabbitMQ.Next.Publisher;
 
 public static class ConnectionExtensions
 {
-    private static readonly Func<IPublishMiddleware, IPublishMiddleware> DefaultPublishPipelineFactory = (p) => p;
-    private static readonly Func<IReturnMiddleware, IReturnMiddleware> DefaultReturnPipelineFactory = (p) => null;
-
     public static IPublisher Publisher(this IConnection connection, string exchange, Action<IPublisherBuilder> builder = null)
     {
         if (string.IsNullOrWhiteSpace(exchange))
@@ -27,44 +24,13 @@ public static IPublisher Publisher(this IConnection connection, string exchange,
             new MessageBuilderPoolPolicy(),
             10);
 
-        var publishPipelineFactory = DefaultPublishPipelineFactory;
-        if (publisherBuilder.PublishMiddlewares.Count > 0)
-        {
-            publishPipelineFactory = (p) =>
-            {
-                var last = p;
-                for (var i = 0; i < publisherBuilder.PublishMiddlewares.Count; i++)
-                {
-                    last = publisherBuilder.PublishMiddlewares[i]?.Invoke(last);
-                }
-
-                return last;
-            };
-        }
-
-        var returnPipelineFactory = DefaultReturnPipelineFactory;
-        if (publisherBuilder.ReturnMiddlewares.Count > 0)
-        {
-            returnPipelineFactory = (p) =>
-            {
-                var last = p;
-                for (var i = 0; i < publisherBuilder.ReturnMiddlewares.Count; i++)
-                {
-                    last = publisherBuilder.ReturnMiddlewares[i]?.Invoke(last);
-                }
-
-                return last;
-            };
-        }
-
         var publisher = new Publisher(
             connection,
             messagePropsPool,
             publisherBuilder.Serializer,
             exchange,
-            publisherBuilder.PublisherConfirms,
-            publishPipelineFactory,
-            returnPipelineFactory);
+            publisherBuilder.PublishMiddlewares,
+            publisherBuilder.ReturnMiddlewares);
 
         return publisher;
     }

src/RabbitMQ.Next.Publisher/InternalMessagePublisher.cs

@@ -3,47 +3,37 @@
 using RabbitMQ.Next.Channels;
 using RabbitMQ.Next.Messaging;
 using RabbitMQ.Next.Serialization;
+using RabbitMQ.Next.Tasks;
 
 namespace RabbitMQ.Next.Publisher;
 
 internal sealed class InternalMessagePublisher : IPublishMiddleware
 {
-    private readonly IChannel channel;
     private readonly ISerializer serializer;
-    private readonly ConfirmMessageHandler confirms;
     private readonly string exchange;
-
-    private long lastDeliveryTag;
-
-    public InternalMessagePublisher(IChannel channel, string exchange, ISerializer serializer, ConfirmMessageHandler confirms)
+    private IChannel publisherChannel;
+
+    public InternalMessagePublisher(string exchange, ISerializer serializer)
     {
-        this.channel = channel;
         this.exchange = exchange;
         this.serializer = serializer;
-        this.confirms = confirms;
     }
 
-    public async ValueTask InvokeAsync<TContent>(TContent content, IMessageBuilder message, CancellationToken cancellation = default)
+    public ValueTask InitAsync(IChannel channel, CancellationToken cancellation)
+    {
+        this.publisherChannel = channel;
+        return default;
+    }
+
+    public ValueTask<ulong> InvokeAsync<TContent>(TContent content, IMessageBuilder message, CancellationToken cancellation = default)
     {
         var flags = this.ComposePublishFlags(message);
 
-        await this.channel.PublishAsync(
+        return this.publisherChannel.PublishAsync(
             (content, this.serializer, message),
             this.exchange, message.RoutingKey, message,
             (st, buffer) => st.serializer.Serialize(st.message, st.content, buffer),
-            flags, cancellation);
-
-        var deliveryTag = Interlocked.Increment(ref this.lastDeliveryTag);
-
-        if (this.confirms != null)
-        {
-            var confirmed = await this.confirms.WaitForConfirmAsync((ulong)deliveryTag);
-            if (!confirmed)
-            {
-                // todo: provide some useful info here
-                throw new DeliveryFailedException();
-            }
-        }
+            flags, cancellation).AsValueTask();
     }
 
     private PublishFlags ComposePublishFlags(IMessageBuilder message)

src/RabbitMQ.Next.Publisher/MessageBuilder.cs

@@ -91,8 +91,6 @@ public IMessageBuilder SetContentEncoding(string value)
         return this;
     }
 
-    public IMessageBuilder SetSetHeader(string key, object value) => throw new NotImplementedException();
-
     public IMessageBuilder SetHeader(string key, object value)
     {
         if (string.IsNullOrWhiteSpace(key))