Do Serialisation in nats con to enable simplification

Signed-off-by: James Thompson <[email protected]>

Author: thompson-tomo

src/NATS.Client.Core/Commands/CommandWriter.cs

@@ -33,16 +33,13 @@ internal sealed class CommandWriter : IAsyncDisposable
     private readonly bool _trace;
     private readonly string _name;
     private readonly NatsConnection _connection;
-    private readonly ObjectPool _pool;
-    private readonly int _arrayPoolInitialSize;
     private readonly object _lock = new();
     private readonly CancellationTokenSource _cts;
     private readonly ConnectionStatsCounter _counter;
     private readonly Memory<byte> _consolidateMem = new byte[SendMemSize].AsMemory();
     private readonly TimeSpan _defaultCommandTimeout;
     private readonly Action<PingCommand> _enqueuePing;
     private readonly ProtocolWriter _protocolWriter;
-    private readonly HeaderWriter _headerWriter;
     private readonly Channel<int> _channelLock;
     private readonly Channel<int> _channelSize;
     private readonly PipeReader _pipeReader;
@@ -55,25 +52,19 @@ internal sealed class CommandWriter : IAsyncDisposable
     private CancellationTokenSource? _ctsReader;
     private volatile bool _disposed;
 
-    public CommandWriter(string name, NatsConnection connection, ObjectPool pool, NatsOpts opts, ConnectionStatsCounter counter, Action<PingCommand> enqueuePing, TimeSpan? overrideCommandTimeout = default)
+    public CommandWriter(string name, NatsConnection connection, NatsOpts opts, ConnectionStatsCounter counter, Action<PingCommand> enqueuePing, TimeSpan? overrideCommandTimeout = default)
     {
         _logger = opts.LoggerFactory.CreateLogger<CommandWriter>();
         _trace = _logger.IsEnabled(LogLevel.Trace);
         _name = name;
         _connection = connection;
-        _pool = pool;
-
-        // Derive ArrayPool rent size from buffer size to
-        // avoid defining another option.
-        _arrayPoolInitialSize = opts.WriterBufferSize / 256;
 
         _counter = counter;
         _defaultCommandTimeout = overrideCommandTimeout ?? opts.CommandTimeout;
         _enqueuePing = enqueuePing;
         _protocolWriter = new ProtocolWriter(opts.SubjectEncoding);
         _channelLock = Channel.CreateBounded<int>(1);
         _channelSize = Channel.CreateUnbounded<int>(new UnboundedChannelOptions { SingleWriter = true, SingleReader = true });
-        _headerWriter = new HeaderWriter(opts.HeaderEncoding);
         _cts = new CancellationTokenSource();
 
         var pipe = new Pipe(new PipeOptions(
@@ -294,49 +285,23 @@ public ValueTask PongAsync(CancellationToken cancellationToken = default)
         return default;
     }
 
-    public ValueTask PublishAsync<T>(string subject, T? value, NatsHeaders? headers, string? replyTo, INatsSerialize<T> serializer, CancellationToken cancellationToken)
+    public ValueTask PublishAsync(string subject, ReadOnlyMemory<byte>? headersBuffer, ReadOnlyMemory<byte> payloadBuffer, string? replyTo, CancellationToken cancellationToken)
     {
         if (_trace)
         {
             _logger.LogTrace(NatsLogEvents.Protocol, "PUB {Subject} {ReplyTo}", subject, replyTo);
         }
 
-        NatsPooledBufferWriter<byte>? headersBuffer = null;
-        if (headers != null)
-        {
-            if (!_pool.TryRent(out headersBuffer))
-                headersBuffer = new NatsPooledBufferWriter<byte>(_arrayPoolInitialSize);
-        }
-
-        NatsPooledBufferWriter<byte> payloadBuffer;
-        if (!_pool.TryRent(out payloadBuffer!))
-            payloadBuffer = new NatsPooledBufferWriter<byte>(_arrayPoolInitialSize);
-
         try
         {
-            if (headers != null)
-                _headerWriter.Write(headersBuffer!, headers);
-
-            if (value != null)
-                serializer.Serialize(payloadBuffer, value);
-
-            var size = payloadBuffer.WrittenMemory.Length + (headersBuffer?.WrittenMemory.Length ?? 0);
+            var size = payloadBuffer.Length + (headersBuffer?.Length ?? 0);
             if (_connection.ServerInfo is { } info && size > info.MaxPayload)
             {
                 throw new NatsPayloadTooLargeException($"Payload size {size} exceeds server's maximum payload size {info.MaxPayload}");
             }
         }
         catch
         {
-            payloadBuffer.Reset();
-            _pool.Return(payloadBuffer);
-
-            if (headersBuffer != null)
-            {
-                headersBuffer.Reset();
-                _pool.Return(headersBuffer);
-            }
-
             throw;
         }
 
@@ -361,21 +326,12 @@ public ValueTask PublishAsync<T>(string subject, T? value, NatsHeaders? headers,
                 throw new ObjectDisposedException(nameof(CommandWriter));
             }
 
-            _protocolWriter.WritePublish(_pipeWriter, subject, replyTo, headersBuffer?.WrittenMemory, payloadBuffer.WrittenMemory);
+            _protocolWriter.WritePublish(_pipeWriter, subject, replyTo, headersBuffer, payloadBuffer);
             EnqueueCommand();
         }
         finally
         {
             _semLock.Release();
-
-            payloadBuffer.Reset();
-            _pool.Return(payloadBuffer);
-
-            if (headersBuffer != null)
-            {
-                headersBuffer.Reset();
-                _pool.Return(headersBuffer);
-            }
         }
 
         return default;
@@ -815,7 +771,7 @@ private async ValueTask PongStateMachineAsync(bool lockHeld, CancellationToken c
 #if !NETSTANDARD
     [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))]
 #endif
-    private async ValueTask PublishStateMachineAsync(bool lockHeld, string subject, string? replyTo, NatsPooledBufferWriter<byte>? headersBuffer, NatsPooledBufferWriter<byte> payloadBuffer, CancellationToken cancellationToken)
+    private async ValueTask PublishStateMachineAsync(bool lockHeld, string subject, string? replyTo, ReadOnlyMemory<byte>? headersBuffer, ReadOnlyMemory<byte> payloadBuffer, CancellationToken cancellationToken)
     {
         try
         {
@@ -839,7 +795,7 @@ private async ValueTask PublishStateMachineAsync(bool lockHeld, string subject,
                     await _flushTask!.WaitAsync(_defaultCommandTimeout, cancellationToken).ConfigureAwait(false);
                 }
 
-                _protocolWriter.WritePublish(_pipeWriter, subject, replyTo, headersBuffer?.WrittenMemory, payloadBuffer.WrittenMemory);
+                _protocolWriter.WritePublish(_pipeWriter, subject, replyTo, headersBuffer, payloadBuffer);
                 EnqueueCommand();
             }
             catch (TimeoutException)
@@ -853,16 +809,9 @@ private async ValueTask PublishStateMachineAsync(bool lockHeld, string subject,
                 _semLock.Release();
             }
         }
-        finally
+        catch
         {
-            payloadBuffer.Reset();
-            _pool.Return(payloadBuffer);
-
-            if (headersBuffer != null)
-            {
-                headersBuffer.Reset();
-                _pool.Return(headersBuffer);
-            }
+            throw;
         }
     }
 

src/NATS.Client.Core/Commands/PriorityCommandWriter.cs

@@ -8,7 +8,7 @@ internal sealed class PriorityCommandWriter : IAsyncDisposable
 
     public PriorityCommandWriter(NatsConnection connection, ObjectPool pool, SocketConnectionWrapper socketConnection, NatsOpts opts, ConnectionStatsCounter counter, Action<PingCommand> enqueuePing)
     {
-        CommandWriter = new CommandWriter("init", connection, pool, opts, counter, enqueuePing);
+        CommandWriter = new CommandWriter("init", connection, opts, counter, enqueuePing);
         CommandWriter.Reset(socketConnection);
     }
 

src/NATS.Client.Core/Internal/NatsRecievedEvent.cs

@@ -0,0 +1,22 @@
+using System.Buffers;
+
+namespace NATS.Client.Core.Internal;
+
+public readonly record struct NatsRecievedEvent
+{
+    public NatsRecievedEvent(string subject, string? replyTo, ReadOnlySequence<byte>? headersBuffer, ReadOnlySequence<byte> payload)
+    {
+        Subject = subject;
+        ReplyTo = replyTo;
+        HeadersBuffer = headersBuffer;
+        Payload = payload;
+    }
+
+    public string Subject { get; }
+
+    public string? ReplyTo { get; }
+
+    public ReadOnlySequence<byte>? HeadersBuffer { get; }
+
+    public ReadOnlySequence<byte> Payload { get; }
+}

src/NATS.Client.Core/Internal/ReplyTask.cs

@@ -10,27 +10,23 @@ internal sealed class ReplyTask<T> : ReplyTaskBase, IDisposable
     private readonly object _gate;
     private readonly ReplyTaskFactory _factory;
     private readonly long _id;
-    private readonly NatsConnection _connection;
-    private readonly INatsDeserialize<T> _deserializer;
     private readonly TimeSpan _requestTimeout;
     private readonly TaskCompletionSource _tcs;
-    private NatsMsg<T> _msg;
+    private NatsRecievedEvent _msg;
 
-    public ReplyTask(ReplyTaskFactory factory, long id, string subject, NatsConnection connection, INatsDeserialize<T> deserializer, TimeSpan requestTimeout)
+    public ReplyTask(ReplyTaskFactory factory, long id, string subject, TimeSpan requestTimeout)
     {
         _factory = factory;
         _id = id;
         Subject = subject;
-        _connection = connection;
-        _deserializer = deserializer;
         _requestTimeout = requestTimeout;
         _tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
         _gate = new object();
     }
 
     public string Subject { get; }
 
-    public async ValueTask<NatsMsg<T>> GetResultAsync(CancellationToken cancellationToken)
+    public async ValueTask<NatsRecievedEvent> GetResultAsync(CancellationToken cancellationToken)
     {
         try
         {
@@ -53,7 +49,7 @@ public override void SetResult(string? replyTo, ReadOnlySequence<byte> payload,
     {
         lock (_gate)
         {
-            _msg = NatsMsg<T>.Build(Subject, replyTo, headersBuffer, payload, _connection, _connection.HeaderParser, _deserializer);
+            _msg = new NatsRecievedEvent(Subject, replyTo, headersBuffer, payload);
         }
 
         _tcs.TrySetResult();
@@ -73,7 +69,6 @@ internal sealed class ReplyTaskFactory
     private readonly string _inboxPrefixString;
     private readonly NatsConnection _connection;
     private readonly ConcurrentDictionary<long, ReplyTaskBase> _replies;
-    private readonly INatsSerializerRegistry _serializerRegistry;
     private readonly TimeSpan _requestTimeout;
     private readonly int _subjectMaxLength;
     private readonly bool _allocSubject;
@@ -86,14 +81,12 @@ public ReplyTaskFactory(NatsConnection connection)
         _inboxPrefix = Encoding.UTF8.GetBytes(_inboxPrefixString);
         _subjectMaxLength = _inboxPrefix.Length + 20; // 20 digits for long
         _allocSubject = _subjectMaxLength < 128;
-        _serializerRegistry = _connection.Opts.SerializerRegistry;
         _requestTimeout = _connection.Opts.RequestTimeout;
         _replies = new ConcurrentDictionary<long, ReplyTaskBase>();
     }
 
-    public ReplyTask<TReply> CreateReplyTask<TReply>(INatsDeserialize<TReply>? deserializer, TimeSpan? requestTimeout)
+    public ReplyTask<TReply> CreateReplyTask<TReply>(TimeSpan? requestTimeout)
     {
-        deserializer ??= _serializerRegistry.GetDeserializer<TReply>();
         var id = Interlocked.Increment(ref _nextId);
 
         string subject;
@@ -117,7 +110,7 @@ public ReplyTask<TReply> CreateReplyTask<TReply>(INatsDeserialize<TReply>? deser
             subject = _inboxPrefixString + id;
         }
 
-        var rt = new ReplyTask<TReply>(this, id, subject, _connection, deserializer, requestTimeout ?? _requestTimeout);
+        var rt = new ReplyTask<TReply>(this, id, subject, requestTimeout ?? _requestTimeout);
         _replies.TryAdd(id, rt);
         return rt;
     }

src/NATS.Client.Core/NatsConnection.Publish.cs

@@ -1,3 +1,5 @@
+using System.Diagnostics;
+using NATS.Client.Core.Commands;
 using NATS.Client.Core.Internal;
 
 namespace NATS.Client.Core;
@@ -14,9 +16,7 @@ public ValueTask PublishAsync(string subject, NatsHeaders? headers = default, st
             try
             {
                 headers?.SetReadOnly();
-                return ConnectionState != NatsConnectionState.Open
-                    ? ConnectAndPublishAsync(subject, default, headers, replyTo, NatsRawSerializer<byte[]>.Default, cancellationToken)
-                    : CommandWriter.PublishAsync(subject, default, headers, replyTo, NatsRawSerializer<byte[]>.Default, cancellationToken);
+                return PerformPublishAsync(subject, headers, default, NatsRawSerializer<byte[]>.Default, replyTo, cancellationToken);
             }
             catch (Exception ex)
             {
@@ -26,9 +26,7 @@ public ValueTask PublishAsync(string subject, NatsHeaders? headers = default, st
         }
 
         headers?.SetReadOnly();
-        return ConnectionState != NatsConnectionState.Open
-            ? ConnectAndPublishAsync(subject, default, headers, replyTo, NatsRawSerializer<byte[]>.Default, cancellationToken)
-            : CommandWriter.PublishAsync(subject, default, headers, replyTo, NatsRawSerializer<byte[]>.Default, cancellationToken);
+        return PerformPublishAsync(subject, headers, default, NatsRawSerializer<byte[]>.Default, replyTo, cancellationToken);
     }
 
     /// <inheritdoc />
@@ -42,9 +40,7 @@ public ValueTask PublishAsync<T>(string subject, T? data, NatsHeaders? headers =
             {
                 serializer ??= Opts.SerializerRegistry.GetSerializer<T>();
                 headers?.SetReadOnly();
-                return ConnectionState != NatsConnectionState.Open
-                    ? ConnectAndPublishAsync(subject, data, headers, replyTo, serializer, cancellationToken)
-                    : CommandWriter.PublishAsync(subject, data, headers, replyTo, serializer, cancellationToken);
+                return PerformPublishAsync(subject, headers, data, serializer, replyTo, cancellationToken);
             }
             catch (Exception ex)
             {
@@ -55,18 +51,71 @@ public ValueTask PublishAsync<T>(string subject, T? data, NatsHeaders? headers =
 
         serializer ??= Opts.SerializerRegistry.GetSerializer<T>();
         headers?.SetReadOnly();
-        return ConnectionState != NatsConnectionState.Open
-            ? ConnectAndPublishAsync(subject, data, headers, replyTo, serializer, cancellationToken)
-            : CommandWriter.PublishAsync(subject, data, headers, replyTo, serializer, cancellationToken);
+        return PerformPublishAsync<T>(subject, headers, data, serializer, replyTo, cancellationToken);
     }
 
     /// <inheritdoc />
     public ValueTask PublishAsync<T>(in NatsMsg<T> msg, INatsSerialize<T>? serializer = default, NatsPubOpts? opts = default, CancellationToken cancellationToken = default) =>
         PublishAsync(msg.Subject, msg.Data, msg.Headers, msg.ReplyTo, serializer, opts, cancellationToken);
 
-    private async ValueTask ConnectAndPublishAsync<T>(string subject, T? data, NatsHeaders? headers, string? replyTo, INatsSerialize<T> serializer, CancellationToken cancellationToken)
+    private async ValueTask ConnectAndPublishAsync(string subject, ReadOnlyMemory<byte>? headers, ReadOnlyMemory<byte> data, string? replyTo, CancellationToken cancellationToken)
     {
         await ConnectAsync().AsTask().WaitAsync(cancellationToken).ConfigureAwait(false);
-        await CommandWriter.PublishAsync(subject, data, headers, replyTo, serializer, cancellationToken).ConfigureAwait(false);
+        await CommandWriter.PublishAsync(subject, headers, data, replyTo, cancellationToken).ConfigureAwait(false);
+    }
+
+    private ValueTask PerformPublishAsync<T>(string subject, NatsHeaders? headers, T? data, INatsSerialize<T> serializer, string? replyTo = default, CancellationToken cancellationToken = default)
+    {
+        NatsPooledBufferWriter<byte>? headersBuffer = null;
+        NatsPooledBufferWriter<byte>? payloadBuffer = null;
+
+        ValueTask task;
+
+        try
+        {
+            if (headers != null)
+            {
+                headers?.SetReadOnly();
+                if (!_pool.TryRent(out headersBuffer))
+                    headersBuffer = new NatsPooledBufferWriter<byte>(_arrayPoolInitialSize);
+            }
+
+            if (!_pool.TryRent(out payloadBuffer!))
+                payloadBuffer = new NatsPooledBufferWriter<byte>(_arrayPoolInitialSize);
+
+            if (headers != null)
+                _headerWriter.Write(headersBuffer!, headers);
+
+            if (data != null)
+                serializer.Serialize(payloadBuffer, data);
+            task = DoPublishAsync(subject, headersBuffer?.WrittenMemory, payloadBuffer.WrittenMemory, replyTo, cancellationToken);
+        }
+        catch (Exception ex)
+        {
+            throw;
+        }
+        finally
+        {
+            if (payloadBuffer != null)
+            {
+                payloadBuffer.Reset();
+                _pool.Return(payloadBuffer);
+            }
+
+            if (headersBuffer != null)
+            {
+                headersBuffer.Reset();
+                _pool.Return(headersBuffer);
+            }
+        }
+
+        return task;
+    }
+
+    private ValueTask DoPublishAsync(string subject, ReadOnlyMemory<byte>? headers = default, ReadOnlyMemory<byte> data = default, string? replyTo = default, CancellationToken cancellationToken = default)
+    {
+        return ConnectionState != NatsConnectionState.Open
+            ? ConnectAndPublishAsync(subject, headers, data, replyTo, cancellationToken)
+            : CommandWriter.PublishAsync(subject, headers, data, replyTo, cancellationToken);
     }
 }