Write to the writing Pipeline directly

See: https://blog.marcgravell.com/2018/07/pipe-dreams-part-3.html#writes-and-wrongs Fixup API validation Add assertion Fix a missing Advance/1 call Add FUTURE
rabbitmq · Mar 3, 2023 · 1606f7b · 1606f7b
1 parent 9652280
commit 1606f7b
Show file tree

Hide file tree

Showing 45 changed files with 223 additions and 106 deletions.
diff --git a/RabbitMQ.Stream.Client/Client.cs b/RabbitMQ.Stream.Client/Client.cs
@@ -85,7 +85,7 @@ public OutgoingMsg(byte publisherId, ulong publishingId, Message data)
         public Message Data => data;
         public int SizeNeeded => 0;
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
             throw new NotImplementedException();
         }

diff --git a/RabbitMQ.Stream.Client/CloseRequest.cs b/RabbitMQ.Stream.Client/CloseRequest.cs
@@ -2,7 +2,7 @@
 // 2.0, and the Mozilla Public License, version 2.0.
 // Copyright (c) 2007-2023 VMware, Inc.
 
-using System;
+using System.Buffers;
 
 namespace RabbitMQ.Stream.Client
 {
@@ -20,13 +20,15 @@ public CloseRequest(uint correlationId, string reason)
 
         public int SizeNeeded => 10 + WireFormatting.StringSize(reason);
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
+            var span = writer.GetSpan(SizeNeeded);
             var offset = WireFormatting.WriteUInt16(span, Key);
             offset += WireFormatting.WriteUInt16(span.Slice(offset), ((ICommand)this).Version);
             offset += WireFormatting.WriteUInt32(span.Slice(offset), correlationId);
             offset += WireFormatting.WriteUInt16(span.Slice(offset), 1); //ok code
             offset += WireFormatting.WriteString(span.Slice(offset), reason);
+            writer.Advance(offset);
             return offset;
         }
     }

diff --git a/RabbitMQ.Stream.Client/CloseResponse.cs b/RabbitMQ.Stream.Client/CloseResponse.cs
@@ -25,10 +25,11 @@ public CloseResponse(uint correlationId, ResponseCode responseCode)
 
         public ResponseCode ResponseCode => responseCode;
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
             throw new NotImplementedException();
         }
+
         internal static int Read(ReadOnlySequence<byte> frame, out CloseResponse command)
         {
             var offset = WireFormatting.ReadUInt16(frame, out _);

diff --git a/RabbitMQ.Stream.Client/Connection.cs b/RabbitMQ.Stream.Client/Connection.cs
@@ -83,34 +83,94 @@ public static async Task<Connection> Create(EndPoint endpoint, Func<Memory<byte>
             return new Connection(socket, commandCallback, closedCallBack, sslOption);
         }
 
-        public async ValueTask<bool> Write<T>(T command) where T : struct, ICommand
+        public ValueTask<bool> Write<T>(T command) where T : struct, ICommand
         {
-            await WriteCommand(command).ConfigureAwait(false);
-            // we return true to indicate that the command was written
-            // In this PR https://github.com/rabbitmq/rabbitmq-stream-dotnet-client/pull/220
-            // we made all WriteCommand async so await is enough to indicate that the command was written
-            // We decided to keep the return value to avoid a breaking change
-            return true;
+            if (!_writeLock.Wait(0))
+            {
+                // https://blog.marcgravell.com/2018/07/pipe-dreams-part-3.html
+                var writeSlowPath = WriteCommandAsyncSlowPath(command);
+                writeSlowPath.ConfigureAwait(false);
+                return writeSlowPath;
+            }
+            else
+            {
+                var release = true;
+                try
+                {
+                    var payloadSize = WriteCommandPayloadSize(command);
+                    var written = command.Write(writer);
+                    Debug.Assert(payloadSize == written);
+                    var flush = writer.FlushAsync();
+                    flush.ConfigureAwait(false);
+                    if (flush.IsCompletedSuccessfully)
+                    {
+                        // we return true to indicate that the command was written
+                        // In this PR https://github.com/rabbitmq/rabbitmq-stream-dotnet-client/pull/220
+                        // we made all WriteCommand async so await is enough to indicate that the command was written
+                        // We decided to keep the return value to avoid a breaking change
+                        return ValueTask.FromResult(true);
+                    }
+                    else
+                    {
+                        release = false;
+                        return AwaitFlushAndRelease(flush);
+                    }
+                }
+                finally
+                {
+                    if (release)
+                    {
+                        _writeLock.Release();
+                    }
+                }
+            }
         }
 
-        private async Task WriteCommand<T>(T command) where T : struct, ICommand
+        private async ValueTask<bool> WriteCommandAsyncSlowPath<T>(T command) where T : struct, ICommand
         {
             // Only one thread should be able to write to the output pipeline at a time.
             await _writeLock.WaitAsync().ConfigureAwait(false);
             try
             {
-                var size = command.SizeNeeded;
-                var mem = new byte[4 + size]; // + 4 to write the size
-                WireFormatting.WriteUInt32(mem, (uint)size);
-                var written = command.Write(mem.AsSpan()[4..]);
-                await writer.WriteAsync(new ReadOnlyMemory<byte>(mem)).ConfigureAwait(false);
-                Debug.Assert(size == written);
+                var payloadSize = WriteCommandPayloadSize(command);
+                var written = command.Write(writer);
+                Debug.Assert(payloadSize == written);
                 await writer.FlushAsync().ConfigureAwait(false);
             }
             finally
             {
                 _writeLock.Release();
             }
+
+            return true;
+        }
+
+        private async ValueTask<bool> AwaitFlushAndRelease(ValueTask<FlushResult> task)
+        {
+            try
+            {
+                await task.ConfigureAwait(false);
+            }
+            finally
+            {
+                _writeLock.Release();
+            }
+
+            return true;
+        }
+
+        private int WriteCommandPayloadSize<T>(T command) where T : struct, ICommand
+        {
+            /*
+             * TODO FUTURE
+             * This code could be moved into a common base class for all outgoing
+             * commands
+             */
+            var payloadSize = command.SizeNeeded;
+            var mem = writer.GetSpan(WireFormatting.SizeofUInt32);
+            var written = WireFormatting.WriteUInt32(mem, (uint)payloadSize);
+            writer.Advance(written);
+            return payloadSize;
         }
 
         private async Task ProcessIncomingFrames()

diff --git a/RabbitMQ.Stream.Client/ConsumerUpdateQueryResponse.cs b/RabbitMQ.Stream.Client/ConsumerUpdateQueryResponse.cs
@@ -28,7 +28,7 @@ private ConsumerUpdateQueryResponse(uint correlationId, byte subscriptionId, byt
     public bool IsActive => active == 1;
     public int SizeNeeded => throw new NotImplementedException();
 
-    public int Write(Span<byte> span)
+    public int Write(IBufferWriter<byte> writer)
     {
         throw new NotImplementedException();
     }

diff --git a/RabbitMQ.Stream.Client/ConsumerUpdateRequest.cs b/RabbitMQ.Stream.Client/ConsumerUpdateRequest.cs
@@ -2,7 +2,7 @@
 // 2.0, and the Mozilla Public License, version 2.0.
 // Copyright (c) 2007-2023 VMware, Inc.
 
-using System;
+using System.Buffers;
 
 namespace RabbitMQ.Stream.Client;
 
@@ -27,13 +27,15 @@ public int SizeNeeded
         }
     }
 
-    public int Write(Span<byte> span)
+    public int Write(IBufferWriter<byte> writer)
     {
+        var span = writer.GetSpan(SizeNeeded);
         var offset = WireFormatting.WriteUInt16(span, Key);
         offset += WireFormatting.WriteUInt16(span.Slice(offset), ((ICommand)this).Version);
         offset += WireFormatting.WriteUInt32(span.Slice(offset), _correlationId);
         offset += WireFormatting.WriteUInt16(span.Slice(offset), (ushort)ResponseCode.Ok);
         offset += OffsetSpecification.Write(span.Slice(offset));
+        writer.Advance(offset);
         return offset;
     }
 

diff --git a/RabbitMQ.Stream.Client/Create.cs b/RabbitMQ.Stream.Client/Create.cs
@@ -35,8 +35,10 @@ public int SizeNeeded
             }
         }
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
+            var span = writer.GetSpan(SizeNeeded);
+
             var offset = WireFormatting.WriteUInt16(span, Key);
             offset += WireFormatting.WriteUInt16(span.Slice(offset), ((ICommand)this).Version);
             offset += WireFormatting.WriteUInt32(span.Slice(offset), correlationId);
@@ -49,6 +51,7 @@ public int Write(Span<byte> span)
                 offset += WireFormatting.WriteString(span.Slice(offset), value);
             }
 
+            writer.Advance(offset);
             return offset;
         }
     }
@@ -71,10 +74,11 @@ public CreateResponse(uint correlationId, ushort responseCode)
 
         public ResponseCode ResponseCode => (ResponseCode)responseCode;
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
             throw new NotImplementedException();
         }
+
         internal static int Read(ReadOnlySequence<byte> frame, out CreateResponse command)
         {
             var offset = WireFormatting.ReadUInt16(frame, out _);

diff --git a/RabbitMQ.Stream.Client/Credit.cs b/RabbitMQ.Stream.Client/Credit.cs
@@ -2,7 +2,7 @@
 // 2.0, and the Mozilla Public License, version 2.0.
 // Copyright (c) 2007-2023 VMware, Inc.
 
-using System;
+using System.Buffers;
 
 namespace RabbitMQ.Stream.Client
 {
@@ -19,13 +19,15 @@ public CreditRequest(byte subscriptionId, ushort credit)
         }
         public int SizeNeeded => 7;
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
+            var span = writer.GetSpan(SizeNeeded);
             var command = (ICommand)this;
             var offset = WireFormatting.WriteUInt16(span, Key);
             offset += WireFormatting.WriteUInt16(span.Slice(offset), command.Version);
             offset += WireFormatting.WriteByte(span.Slice(offset), subscriptionId);
             offset += WireFormatting.WriteUInt16(span.Slice(offset), credit);
+            writer.Advance(offset);
             return offset;
         }
     }

diff --git a/RabbitMQ.Stream.Client/CreditResponse.cs b/RabbitMQ.Stream.Client/CreditResponse.cs
@@ -24,7 +24,7 @@ private CreditResponse(ResponseCode responseCode, byte subscriptionId)
 
         private ResponseCode ResponseCode { get; }
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
             throw new NotImplementedException();
         }

diff --git a/RabbitMQ.Stream.Client/DeclarePublisherRequest.cs b/RabbitMQ.Stream.Client/DeclarePublisherRequest.cs
@@ -2,7 +2,7 @@
 // 2.0, and the Mozilla Public License, version 2.0.
 // Copyright (c) 2007-2023 VMware, Inc.
 
-using System;
+using System.Buffers;
 
 namespace RabbitMQ.Stream.Client
 {
@@ -24,14 +24,16 @@ public DeclarePublisherRequest(uint correlationId, byte publisherId, string publ
 
         public int SizeNeeded => 8 + 1 + WireFormatting.StringSize(publisherRef) + WireFormatting.StringSize(stream);
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
+            var span = writer.GetSpan(SizeNeeded);
             var offset = WireFormatting.WriteUInt16(span, Key);
             offset += WireFormatting.WriteUInt16(span.Slice(offset), ((ICommand)this).Version);
             offset += WireFormatting.WriteUInt32(span.Slice(offset), correlationId);
             offset += WireFormatting.WriteByte(span.Slice(offset), publisherId);
             offset += WireFormatting.WriteString(span.Slice(offset), publisherRef);
             offset += WireFormatting.WriteString(span.Slice(offset), stream);
+            writer.Advance(offset);
             return offset;
         }
     }

diff --git a/RabbitMQ.Stream.Client/DeclarePublisherResponse.cs b/RabbitMQ.Stream.Client/DeclarePublisherResponse.cs
@@ -25,7 +25,7 @@ public DeclarePublisherResponse(uint correlationId, ResponseCode responseCode)
 
         public ResponseCode ResponseCode => responseCode;
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
             throw new NotImplementedException();
         }

diff --git a/RabbitMQ.Stream.Client/Delete.cs b/RabbitMQ.Stream.Client/Delete.cs
@@ -21,12 +21,14 @@ public DeleteRequest(uint correlationId, string stream)
 
         public int SizeNeeded => 8 + WireFormatting.StringSize(stream);
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
+            var span = writer.GetSpan(SizeNeeded);
             var offset = WireFormatting.WriteUInt16(span, Key);
             offset += WireFormatting.WriteUInt16(span.Slice(offset), ((ICommand)this).Version);
             offset += WireFormatting.WriteUInt32(span.Slice(offset), correlationId);
             offset += WireFormatting.WriteString(span.Slice(offset), stream);
+            writer.Advance(offset);
             return offset;
         }
     }
@@ -49,10 +51,11 @@ public DeleteResponse(uint correlationId, ushort responseCode)
 
         public ResponseCode ResponseCode => (ResponseCode)responseCode;
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
             throw new NotImplementedException();
         }
+
         internal static int Read(ReadOnlySequence<byte> frame, out DeleteResponse command)
         {
             var offset = WireFormatting.ReadUInt16(frame, out _);

diff --git a/RabbitMQ.Stream.Client/DeletePublisherRequest.cs b/RabbitMQ.Stream.Client/DeletePublisherRequest.cs
@@ -2,7 +2,7 @@
 // 2.0, and the Mozilla Public License, version 2.0.
 // Copyright (c) 2007-2023 VMware, Inc.
 
-using System;
+using System.Buffers;
 
 namespace RabbitMQ.Stream.Client
 {
@@ -20,12 +20,14 @@ public DeletePublisherRequest(uint correlationId, byte publisherId)
 
         public int SizeNeeded => 8 + 1;
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
+            var span = writer.GetSpan(SizeNeeded);
             var offset = WireFormatting.WriteUInt16(span, Key);
             offset += WireFormatting.WriteUInt16(span.Slice(offset), ((ICommand)this).Version);
             offset += WireFormatting.WriteUInt32(span.Slice(offset), correlationId);
             offset += WireFormatting.WriteByte(span.Slice(offset), publisherId);
+            writer.Advance(offset);
             return offset;
         }
     }

diff --git a/RabbitMQ.Stream.Client/DeletePublisherResponse.cs b/RabbitMQ.Stream.Client/DeletePublisherResponse.cs
@@ -24,10 +24,11 @@ public DeletePublisherResponse(uint correlationId, ResponseCode responseCode)
 
         public ResponseCode ResponseCode => responseCode;
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
             throw new NotImplementedException();
         }
+
         internal static int Read(ReadOnlySequence<byte> frame, out DeletePublisherResponse command)
         {
             var offset = WireFormatting.ReadUInt16(frame, out _);

diff --git a/RabbitMQ.Stream.Client/Deliver.cs b/RabbitMQ.Stream.Client/Deliver.cs
@@ -27,7 +27,7 @@ private Deliver(byte subscriptionId, Chunk chunk)
 
         public byte SubscriptionId => subscriptionId;
 
-        public int Write(Span<byte> span)
+        public int Write(IBufferWriter<byte> writer)
         {
             throw new NotImplementedException();
         }

diff --git a/RabbitMQ.Stream.Client/HeartBeatRequest.cs b/RabbitMQ.Stream.Client/HeartBeatRequest.cs
@@ -2,7 +2,7 @@
 // 2.0, and the Mozilla Public License, version 2.0.
 // Copyright (c) 2007-2023 VMware, Inc.
 
-using System;
+using System.Buffers;
 
 namespace RabbitMQ.Stream.Client;
 
@@ -12,10 +12,12 @@ namespace RabbitMQ.Stream.Client;
 
     public int SizeNeeded => 4;
 
-    public int Write(Span<byte> span)
+    public int Write(IBufferWriter<byte> writer)
     {
+        var span = writer.GetSpan(SizeNeeded);
         var offset = WireFormatting.WriteUInt16(span, Key);
         offset += WireFormatting.WriteUInt16(span.Slice(offset), ((ICommand)this).Version);
+        writer.Advance(offset);
         return offset;
     }
 }