Consume clean exit fixes

sandbox/Example.JetStream.PullConsumer/Program.cs

@@ -1,5 +1,6 @@
+using System.Buffers;
 using System.Diagnostics;
-using Example.JetStream.PullConsumer;
+using System.Text;
 using Microsoft.Extensions.Logging;
 using NATS.Client.Core;
 using NATS.Client.JetStream;
@@ -20,12 +21,12 @@
 
 var consumer = await js.CreateConsumerAsync("s1", "c1");
 
-var idle = TimeSpan.FromSeconds(15);
-var expires = TimeSpan.FromSeconds(30);
+var idle = TimeSpan.FromSeconds(5);
+var expires = TimeSpan.FromSeconds(10);
 
 // int? maxMsgs = null;
 // int? maxBytes = 128;
-int? maxMsgs = 1000;
+int? maxMsgs = 10;
 int? maxBytes = null;
 
 void Report(int i, Stopwatch sw, string data)
@@ -41,7 +42,6 @@ void Report(int i, Stopwatch sw, string data)
     MaxBytes = maxBytes,
     Expires = expires,
     IdleHeartbeat = idle,
-    Serializer = new RawDataSerializer(),
 };
 
 var fetchOpts = new NatsJSFetchOpts
@@ -50,32 +50,39 @@ void Report(int i, Stopwatch sw, string data)
     MaxBytes = maxBytes,
     Expires = expires,
     IdleHeartbeat = idle,
-    Serializer = new RawDataSerializer(),
 };
 
 var nextOpts = new NatsJSNextOpts
 {
     Expires = expires,
     IdleHeartbeat = idle,
-    Serializer = new RawDataSerializer(),
 };
 
 var stopwatch = Stopwatch.StartNew();
 var count = 0;
 
+var cmd = args.Length > 0 ? args[0] : "consume";
+var cmdOpt = args.Length > 1 ? args[1] : "none";
+
 try
 {
-    if (args.Length > 0 && args[0] == "fetch")
+    if (cmd == "fetch")
     {
         while (!cts.Token.IsCancellationRequested)
         {
             try
             {
                 Console.WriteLine($"___\nFETCH {maxMsgs}");
                 await consumer.RefreshAsync(cts.Token);
-                await using var sub = await consumer.FetchAsync<RawData>(fetchOpts, cts.Token);
+                await using var sub = await consumer.FetchAsync<NatsMemoryOwner<byte>>(fetchOpts, cts.Token);
                 await foreach (var msg in sub.Msgs.ReadAllAsync(cts.Token))
                 {
+                    using (msg.Data)
+                    {
+                        var message = Encoding.ASCII.GetString(msg.Data.Span);
+                        Console.WriteLine($"Received: {message}");
+                    }
+
                     await msg.AckAsync(cancellationToken: cts.Token);
                     Report(++count, stopwatch, $"data: {msg.Data}");
                 }
@@ -91,16 +98,22 @@ void Report(int i, Stopwatch sw, string data)
             }
         }
     }
-    else if (args.Length > 0 && args[0] == "fetch-all")
+    else if (cmd == "fetch-all")
     {
         while (!cts.Token.IsCancellationRequested)
         {
             try
             {
                 Console.WriteLine($"___\nFETCH {maxMsgs}");
                 await consumer.RefreshAsync(cts.Token);
-                await foreach (var msg in consumer.FetchAllAsync<RawData>(fetchOpts, cts.Token))
+                await foreach (var msg in consumer.FetchAllAsync<NatsMemoryOwner<byte>>(fetchOpts, cts.Token))
                 {
+                    using (msg.Data)
+                    {
+                        var message = Encoding.ASCII.GetString(msg.Data.Span);
+                        Console.WriteLine($"Received: {message}");
+                    }
+
                     await msg.AckAsync(cancellationToken: cts.Token);
                     Report(++count, stopwatch, $"data: {msg.Data}");
                 }
@@ -116,16 +129,22 @@ void Report(int i, Stopwatch sw, string data)
             }
         }
     }
-    else if (args.Length > 0 && args[0] == "next")
+    else if (cmd == "next")
     {
         while (!cts.Token.IsCancellationRequested)
         {
             try
             {
                 Console.WriteLine("___\nNEXT");
-                var next = await consumer.NextAsync<RawData>(nextOpts, cts.Token);
+                var next = await consumer.NextAsync<NatsMemoryOwner<byte>>(nextOpts, cts.Token);
                 if (next is { } msg)
                 {
+                    using (msg.Data)
+                    {
+                        var message = Encoding.ASCII.GetString(msg.Data.Span);
+                        Console.WriteLine($"Received: {message}");
+                    }
+
                     await msg.AckAsync(cancellationToken: cts.Token);
                     Report(++count, stopwatch, $"data: {msg.Data}");
                 }
@@ -141,21 +160,48 @@ void Report(int i, Stopwatch sw, string data)
             }
         }
     }
-    else if (args.Length > 0 && args[0] == "consume")
+    else if (cmd == "consume")
     {
         while (!cts.Token.IsCancellationRequested)
         {
             try
             {
                 Console.WriteLine("___\nCONSUME");
-                await using var sub = await consumer.ConsumeAsync<RawData>(
-                    consumeOpts,
-                    cts.Token);
+                await using var sub = await consumer.ConsumeAsync<NatsMemoryOwner<byte>>(consumeOpts);
 
-                await foreach (var msg in sub.Msgs.ReadAllAsync(cts.Token))
+                cts.Token.Register(() =>
+                {
+                    sub.DisposeAsync().GetAwaiter().GetResult();
+                });
+
+                var stopped = false;
+                await foreach (var msg in sub.Msgs.ReadAllAsync())
                 {
+                    using (msg.Data)
+                    {
+                        var message = Encoding.ASCII.GetString(msg.Data.Span);
+                        Console.WriteLine($"Received: {message}");
+                        if (message == "stop")
+                        {
+                            Console.WriteLine("Stopping consumer...");
+                            sub.Stop();
+                            stopped = true;
+                        }
+                    }
+
                     await msg.AckAsync(cancellationToken: cts.Token);
                     Report(++count, stopwatch, $"data: {msg.Data}");
+
+                    if (cmdOpt == "with-pause")
+                    {
+                        await Task.Delay(1_000);
+                    }
+                }
+
+                if (stopped)
+                {
+                    Console.WriteLine("Stopped consumer.");
+                    break;
                 }
             }
             catch (NatsJSProtocolException e)
@@ -169,15 +215,21 @@ void Report(int i, Stopwatch sw, string data)
             }
         }
     }
-    else if (args.Length > 0 && args[0] == "consume-all")
+    else if (cmd == "consume-all")
     {
         while (!cts.Token.IsCancellationRequested)
         {
             try
             {
                 Console.WriteLine("___\nCONSUME-ALL");
-                await foreach (var msg in consumer.ConsumeAllAsync<RawData>(consumeOpts, cts.Token))
+                await foreach (var msg in consumer.ConsumeAllAsync<NatsMemoryOwner<byte>>(consumeOpts, cts.Token))
                 {
+                    using (msg.Data)
+                    {
+                        var message = Encoding.ASCII.GetString(msg.Data.Span);
+                        Console.WriteLine($"Received: {message}");
+                    }
+
                     await msg.AckAsync(cancellationToken: cts.Token);
                     Report(++count, stopwatch, $"data: {msg.Data}");
                 }

src/NATS.Client.JetStream/INatsJSConsume.cs

@@ -2,18 +2,28 @@
 
 namespace NATS.Client.JetStream;
 
-/// <summary>
-/// Interface to manage a <c>consume()</c> operation on a consumer.
-/// </summary>
-public interface INatsJSConsume : IAsyncDisposable
-{
-    void Stop();
-}
-
 /// <summary>
 /// Interface to extract messages from a <c>consume()</c> operation on a consumer.
 /// </summary>
-public interface INatsJSConsume<T> : INatsJSConsume
+public interface INatsJSConsume<T> : IAsyncDisposable
 {
+    /// <summary>
+    /// Messages received from the consumer.
+    /// </summary>
     ChannelReader<NatsJSMsg<T?>> Msgs { get; }
+
+    /// <summary>
+    /// Stop the consumer gracefully.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// This will wait for any inflight messages to be processed before stopping.
+    /// </para>
+    /// <para>
+    /// Disposing would stop consuming immediately. This might leave messages behind
+    /// without acknowledgement. Which is fine, messages will be scheduled for redelivery,
+    /// however, it might not be the desired behavior.
+    /// </para>
+    /// </remarks>
+    void Stop();
 }

src/NATS.Client.JetStream/INatsJSFetch.cs

@@ -2,18 +2,13 @@
 
 namespace NATS.Client.JetStream;
 
-/// <summary>
-/// Interface to manage a <c>fetch()</c> operation on a consumer.
-/// </summary>
-public interface INatsJSFetch : IAsyncDisposable
-{
-    void Stop();
-}
-
 /// <summary>
 /// Interface to extract messages from a <c>fetch()</c> operation on a consumer.
 /// </summary>
-public interface INatsJSFetch<T> : INatsJSFetch
+public interface INatsJSFetch<T> : IAsyncDisposable
 {
+    /// <summary>
+    /// User messages received from the consumer.
+    /// </summary>
     ChannelReader<NatsJSMsg<T?>> Msgs { get; }
 }