Handle publish serialization exceptions (#144)

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

@@ -10,6 +10,7 @@ internal sealed class PublishCommand<T> : CommandBase<PublishCommand<T>>
     private NatsHeaders? _headers;
     private T? _value;
     private INatsSerializer? _serializer;
+    private Action<Exception>? _errorHandler;
     private CancellationToken _cancellationToken;
 
     private PublishCommand()
@@ -18,7 +19,7 @@ private PublishCommand()
 
     public override bool IsCanceled => _cancellationToken.IsCancellationRequested;
 
-    public static PublishCommand<T> Create(ObjectPool pool, string subject, string? replyTo, NatsHeaders? headers, T? value, INatsSerializer serializer, CancellationToken cancellationToken)
+    public static PublishCommand<T> Create(ObjectPool pool, string subject, string? replyTo, NatsHeaders? headers, T? value, INatsSerializer serializer, Action<Exception>? errorHandler, CancellationToken cancellationToken)
     {
         if (!TryRent(pool, out var result))
         {
@@ -30,14 +31,40 @@ public static PublishCommand<T> Create(ObjectPool pool, string subject, string?
         result._headers = headers;
         result._value = value;
         result._serializer = serializer;
+        result._errorHandler = errorHandler;
         result._cancellationToken = cancellationToken;
 
         return result;
     }
 
     public override void Write(ProtocolWriter writer)
     {
-        writer.WritePublish(_subject!, _replyTo, _headers, _value, _serializer!);
+        try
+        {
+            writer.WritePublish(_subject!, _replyTo, _headers, _value, _serializer!);
+        }
+        catch (Exception e)
+        {
+            if (_errorHandler is { } errorHandler)
+            {
+                ThreadPool.UnsafeQueueUserWorkItem(
+                    state =>
+                    {
+                        try
+                        {
+                            state.handler(state.exception);
+                        }
+                        catch
+                        {
+                            // ignore
+                        }
+                    },
+                    (handler: errorHandler, exception: e),
+                    preferLocal: false);
+            }
+
+            throw;
+        }
     }
 
     protected override void Reset()
@@ -46,6 +73,7 @@ protected override void Reset()
         _headers = default;
         _value = default;
         _serializer = null;
+        _errorHandler = default;
         _cancellationToken = default;
     }
 }

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

@@ -129,14 +129,27 @@ private async Task WriteLoopAsync()
                             continue;
                         }
 
-                        if (command is IBatchCommand batch)
+                        try
                         {
-                            count += batch.Write(protocolWriter);
+                            if (command is IBatchCommand batch)
+                            {
+                                count += batch.Write(protocolWriter);
+                            }
+                            else
+                            {
+                                command.Write(protocolWriter);
+                                count++;
+                            }
                         }
-                        else
+                        catch (Exception e)
                         {
-                            command.Write(protocolWriter);
-                            count++;
+                            // flag potential serialization exceptions
+                            if (command is IPromise promise)
+                            {
+                                promise.SetException(e);
+                            }
+
+                            throw;
                         }
 
                         if (command is IPromise p)

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

@@ -28,20 +28,20 @@ internal ValueTask PubPostAsync(string subject, string? replyTo = default, ReadO
         }
     }
 
-    internal ValueTask PubModelPostAsync<T>(string subject, T? data, INatsSerializer serializer, string? replyTo = default, NatsHeaders? headers = default, CancellationToken cancellationToken = default)
+    internal ValueTask PubModelPostAsync<T>(string subject, T? data, INatsSerializer serializer, string? replyTo = default, NatsHeaders? headers = default, Action<Exception>? errorHandler = default, CancellationToken cancellationToken = default)
     {
         headers?.SetReadOnly();
 
         if (ConnectionState == NatsConnectionState.Open)
         {
-            var command = PublishCommand<T>.Create(_pool, subject, replyTo, headers, data, serializer, cancellationToken);
+            var command = PublishCommand<T>.Create(_pool, subject, replyTo, headers, data, serializer, errorHandler, cancellationToken);
             return EnqueueCommandAsync(command);
         }
         else
         {
-            return WithConnectAsync(subject, replyTo, headers, data, serializer, cancellationToken, static (self, s, r, h, d, ser, c) =>
+            return WithConnectAsync(subject, replyTo, headers, data, serializer, errorHandler, cancellationToken, static (self, s, r, h, d, ser, eh, c) =>
             {
-                var command = PublishCommand<T>.Create(self._pool, s, r, h, d, ser, c);
+                var command = PublishCommand<T>.Create(self._pool, s, r, h, d, ser, eh, c);
                 return self.EnqueueCommandAsync(command);
             });
         }

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

@@ -33,7 +33,7 @@ public ValueTask PublishAsync<T>(string subject, T? data, NatsHeaders? headers =
         }
         else
         {
-            return PubModelPostAsync<T>(subject, data, serializer, replyTo, headers, cancellationToken);
+            return PubModelPostAsync<T>(subject, data, serializer, replyTo, headers, opts?.ErrorHandler, cancellationToken);
         }
     }
 

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

@@ -49,7 +49,7 @@ internal async ValueTask<INatsSub<TReply>> RequestSubAsync<TRequest, TReply>(
         }
         else
         {
-            await PubModelPostAsync(subject, data, serializer, replyTo, headers, cancellationToken).ConfigureAwait(false);
+            await PubModelPostAsync(subject, data, serializer, replyTo, headers, requestOpts?.ErrorHandler, cancellationToken).ConfigureAwait(false);
         }
 
         return sub;

src/NATS.Client.Core/NatsConnection.cs

@@ -747,6 +747,12 @@ private async ValueTask WithConnectAsync<T1, T2, T3, T4, T5, T6>(T1 item1, T2 it
         await coreAsync(this, item1, item2, item3, item4, item5, item6).ConfigureAwait(false);
     }
 
+    private async ValueTask WithConnectAsync<T1, T2, T3, T4, T5, T6, T7>(T1 item1, T2 item2, T3 item3, T4 item4, T5 item5, T6 item6, T7 item7, Func<NatsConnection, T1, T2, T3, T4, T5, T6, T7, ValueTask> coreAsync)
+    {
+        await ConnectAsync().ConfigureAwait(false);
+        await coreAsync(this, item1, item2, item3, item4, item5, item6, item7).ConfigureAwait(false);
+    }
+
     private async ValueTask<T> WithConnectAsync<T>(Func<NatsConnection, ValueTask<T>> coreAsync)
     {
         await ConnectAsync().ConfigureAwait(false);

src/NATS.Client.Core/NatsPubOpts.cs

@@ -9,4 +9,13 @@ public record NatsPubOpts
     /// Default value is false, and calls to PublishAsync will complete after the publish command has been written to the Command Channel
     /// </summary>
     public bool? WaitUntilSent { get; init; }
+
+    /// <summary>
+    /// Optional callback to handle serialization exceptions.
+    /// </summary>
+    /// <remarks>
+    /// When <c>WaitUntilSent</c> is set to <c>false</c> serialization exceptions won't propagate
+    /// to the caller but this callback will be called with the exception thrown by the serializer.
+    /// </remarks>
+    public Action<Exception>? ErrorHandler { get; init; }
 }

src/NATS.Client.JetStream/Internal/NatsJSConsume.cs

@@ -165,6 +165,7 @@ internal override IEnumerable<ICommand> GetReconnectCommands(int sid)
             headers: default,
             value: request,
             serializer: NatsJsonSerializer.Default,
+            errorHandler: default,
             cancellationToken: default);
     }
 

src/NATS.Client.JetStream/Internal/NatsJSFetch.cs

@@ -151,6 +151,7 @@ internal override IEnumerable<ICommand> GetReconnectCommands(int sid)
             headers: default,
             value: request,
             serializer: NatsJsonSerializer.Default,
+            errorHandler: default,
             cancellationToken: default);
     }
 

tests/NATS.Client.Core.Tests/SerializerTest.cs

@@ -0,0 +1,72 @@
+using System.Buffers;
+
+namespace NATS.Client.Core.Tests;
+
+public class SerializerTest
+{
+    private readonly ITestOutputHelper _output;
+
+    public SerializerTest(ITestOutputHelper output) => _output = output;
+
+    [Fact]
+    public async Task Serializer_exceptions()
+    {
+        await using var server = NatsServer.Start();
+        await using var nats = server.CreateClientConnection();
+
+        await Assert.ThrowsAsync<TestSerializerException>(async () =>
+        {
+            var signal = new WaitSignal<Exception>();
+
+            var opts = new NatsPubOpts
+            {
+                Serializer = new TestSerializer(),
+                WaitUntilSent = false,
+                ErrorHandler = e =>
+                {
+                    signal.Pulse(e);
+                },
+            };
+
+            await nats.PublishAsync(
+                "foo",
+                0,
+                opts: opts);
+
+            throw await signal;
+        });
+
+        await Assert.ThrowsAsync<TestSerializerException>(async () =>
+        {
+            await nats.PublishAsync(
+                "foo",
+                0,
+                opts: new NatsPubOpts { Serializer = new TestSerializer(), WaitUntilSent = true });
+        });
+
+        // Check that our connection isn't affected by the exceptions
+        await using var sub = await nats.SubscribeAsync<int>("foo");
+
+        var rtt = await nats.PingAsync();
+        Assert.True(rtt > TimeSpan.Zero);
+
+        await nats.PublishAsync("foo", 1);
+
+        var result = (await sub.Msgs.ReadAsync()).Data;
+
+        Assert.Equal(1, result);
+    }
+}
+
+public class TestSerializer : INatsSerializer
+{
+    public int Serialize<T>(ICountableBufferWriter bufferWriter, T? value) => throw new TestSerializerException();
+
+    public T? Deserialize<T>(in ReadOnlySequence<byte> buffer) => throw new TestSerializerException();
+
+    public object? Deserialize(in ReadOnlySequence<byte> buffer, Type type) => throw new TestSerializerException();
+}
+
+public class TestSerializerException : Exception
+{
+}