Introducing NATS stream provider

Author: galvesribeiro

Directory.Packages.props

@@ -77,7 +77,7 @@
     <PackageVersion Include="StackExchange.Redis" Version="2.8.16" />
     <PackageVersion Include="KubernetesClient" Version="15.0.1" />
     <PackageVersion Include="CassandraCSharpDriver" Version="3.20.1" />
-    <PackageVersion Include="NATS.Net" Version="2.5.9" />
+    <PackageVersion Include="NATS.Net" Version="2.6.4" />
     <!-- Test related packages -->
     <PackageVersion Include="coverlet.collector" Version="6.0.2" />
     <PackageVersion Include="coverlet.msbuild" Version="6.0.0" />

Orleans.slnx

@@ -83,6 +83,9 @@
   <Folder Name="/src/Extensions/Cassandra/">
     <Project Path="src/Cassandra/Orleans.Clustering.Cassandra/Orleans.Clustering.Cassandra.csproj" />
   </Folder>
+  <Folder Name="/src/Extensions/NATS/">
+    <Project Path="src\Orleans.Streaming.NATS\Orleans.Streaming.NATS.csproj" Type="Classic C#" />
+  </Folder>
   <Folder Name="/src/Extensions/Redis/">
     <Project Path="src/Redis/Orleans.Clustering.Redis/Orleans.Clustering.Redis.csproj" />
     <Project Path="src/Redis/Orleans.GrainDirectory.Redis/Orleans.GrainDirectory.Redis.csproj" />
@@ -114,11 +117,11 @@
     <Project Path="test/Grains/BenchmarkGrains/BenchmarkGrains.csproj" />
   </Folder>
   <Folder Name="/test/DistributedTests/">
-    <File Path="test/DistributedTests/README.md" />
     <Project Path="test/DistributedTests/DistributedTests.Client/DistributedTests.Client.csproj" />
     <Project Path="test/DistributedTests/DistributedTests.Common/DistributedTests.Common.csproj" />
     <Project Path="test/DistributedTests/DistributedTests.Grains/DistributedTests.Grains.csproj" />
     <Project Path="test/DistributedTests/DistributedTests.Server/DistributedTests.Server.csproj" />
+    <File Path="test/DistributedTests/README.md" />
   </Folder>
   <Folder Name="/test/Extensions/">
     <Project Path="test/Extensions/AWSUtils.Tests/AWSUtils.Tests.csproj" />
@@ -130,6 +133,7 @@
     <Project Path="test/Extensions/TesterAdoNet/Tester.AdoNet.csproj" />
     <Project Path="test/Extensions/TesterAzureUtils/Tester.AzureUtils.csproj" />
     <Project Path="test/Extensions/TesterZooKeeperUtils/Tester.ZooKeeperUtils.csproj" />
+    <Project Path="test\Extensions\NATS.Tests\NATS.Tests.csproj" Type="Classic C#" />
   </Folder>
   <Folder Name="/test/Grains/">
     <Project Path="test/Grains/TestFSharp/TestFSharp.fsproj" />
@@ -155,4 +159,4 @@
     <Project Path="test/Transactions/Orleans.Transactions.Azure.Test/Orleans.Transactions.Azure.Test.csproj" />
     <Project Path="test/Transactions/Orleans.Transactions.Tests/Orleans.Transactions.Tests.csproj" />
   </Folder>
-</Solution>
+</Solution>

src/Orleans.Streaming.NATS/NatsOptions.cs

@@ -1,4 +1,5 @@
-using NATS.Client.Core;
+using System.Text.Json;
+using NATS.Client.Core;
 
 namespace Orleans.Streaming.NATS;
 
@@ -44,4 +45,9 @@ public class NatsOptions
     /// The number of connections used to send stream messages to NATS JetStream.
     /// </summary>
     public int ProducerCount { get; set; } = 8;
-}
+
+    /// <summary>
+    /// System.Text.Json serializer options to be used by the NATS provider.
+    /// </summary>
+    public JsonSerializerOptions JsonSerializerOptions { get; set; } = default!;
+}

src/Orleans.Streaming.NATS/Providers/NatsAdapterFactory.cs

@@ -42,7 +42,6 @@ public NatsAdapterFactory(
         adapterCache = new SimpleQueueAdapterCache(cacheOptions, this.providerName, this.loggerFactory);
     }
 
-
     /// <summary> Init the factory.</summary>
     public virtual void Init()
     {
@@ -97,4 +96,4 @@ public static NatsAdapterFactory Create(IServiceProvider services, string name)
         factory.Init();
         return factory;
     }
-}
+}

src/Orleans.Streaming.NATS/Providers/NatsConnectionManager.cs

@@ -6,16 +6,16 @@
 using NATS.Client.Core;
 using NATS.Client.JetStream;
 using NATS.Client.JetStream.Models;
+using NATS.Client.Serializers.Json;
 
 namespace Orleans.Streaming.NATS;
 
 /// <summary>
-/// Wrapper around a NATS and JetStream APIs
+/// Wrapper around NATS and JetStream APIs
 /// </summary>
 internal sealed class NatsConnectionManager
 {
-    private const string AckPayload = "+ACK";
-    private static readonly NatsJsonContextSerializer<NatsStreamMessage> Serializer;
+    private static readonly byte[] AckPayload = "+ACK"u8.ToArray();
     private readonly string _providerName;
     private readonly NatsOpts _natsClientOptions;
     private readonly NatsConnection _natsConnection;
@@ -25,20 +25,35 @@ internal sealed class NatsConnectionManager
     private readonly NatsJSContext[] _producerNatsContexts;
     private readonly NatsJSContext _natsContext;
 
-    static NatsConnectionManager()
-    {
-        Serializer = new NatsJsonContextSerializer<NatsStreamMessage>(NatsSerializerContext.Default);
-    }
-
     [GeneratedActivatorConstructor]
     public NatsConnectionManager(string providerName, ILoggerFactory loggerFactory, NatsOptions options)
     {
         this._providerName = providerName;
         this._loggerFactory = loggerFactory;
         this._logger = this._loggerFactory.CreateLogger<NatsConnectionManager>();
         this._options = options;
-        this._natsClientOptions =
-            this._options.NatsClientOptions ?? NatsOpts.Default with { Name = $"Orleans-{this._providerName}" };
+        this._options.JsonSerializerOptions.TypeInfoResolverChain.Add(NatsSerializerContext.Default);
+        if (this._options.NatsClientOptions is null)
+        {
+            this._options.NatsClientOptions = NatsOpts.Default with
+            {
+                Name = $"Orleans-{this._providerName}",
+                SerializerRegistry =
+                new NatsJsonContextOptionsSerializerRegistry(this._options.JsonSerializerOptions)
+            };
+        }
+        else
+        {
+            this._options.NatsClientOptions = this._options.NatsClientOptions with
+            {
+                Name = string.IsNullOrWhiteSpace(this._options.NatsClientOptions.Name)
+                    ? $"Orleans-{this._providerName}"
+                    : this._options.NatsClientOptions.Name,
+                SerializerRegistry = new NatsJsonContextOptionsSerializerRegistry(this._options.JsonSerializerOptions)
+            };
+        }
+
+        this._natsClientOptions = this._options.NatsClientOptions;
         this._natsConnection = new NatsConnection(this._natsClientOptions);
         this._natsContext = new NatsJSContext(this._natsConnection);
 
@@ -93,6 +108,7 @@ public async Task Initialize(CancellationToken cancellationToken = default)
             {
                 var streamConfig = new StreamConfig(this._options.StreamName, [$"{this._providerName}.>"])
                 {
+                    Retention = StreamConfigRetention.Workqueue,
                     SubjectTransform = new SubjectTransform
                     {
                         Src = $"{this._providerName}.*.*",
@@ -140,7 +156,10 @@ public async Task EnqueueMessage(NatsStreamMessage message, CancellationToken ca
 
         var context = this._producerNatsContexts[Math.Abs(id.GetHashCode()) % this._producerNatsContexts.Length];
 
-        var ack = await context.TryPublishAsync(subject, message, Serializer,
+        var ack = await context.TryPublishAsync(
+            subject,
+            message,
+            this._natsClientOptions.SerializerRegistry.GetSerializer<NatsStreamMessage>(),
             cancellationToken: cancellationToken);
 
         if (ack.Success)
@@ -165,14 +184,15 @@ public NatsStreamConsumer CreateConsumer(uint partition) =>
             this._options.StreamName,
             partition,
             this._options.BatchSize,
-            Serializer);
+            this._natsClientOptions.SerializerRegistry.GetDeserializer<NatsStreamMessage>());
 
     /// <summary>
     /// Acknowledge messages on a subject in a NATS JetStream stream
     /// </summary>
     /// <param name="subject">The ReplyTo subject</param>
     public async Task AcknowledgeMessages(string subject)
     {
-        await this._natsConnection.PublishAsync(subject, AckPayload);
+        await this._natsConnection
+            .PublishAsync(subject, AckPayload, serializer: NatsRawSerializer<byte[]>.Default);
     }
-}
+}

src/Orleans.Streaming.NATS/Providers/NatsQueueAdapterReceiver.cs

@@ -130,18 +130,16 @@ public async Task MessagesDeliveredAsync(IList<IBatchContainer> messages)
 
         if (messages.Count == 0) return;
 
-        var natsMessages = messages
-            .Cast<NatsBatchContainer>().Select(b => b.ReplyTo).ToArray();
-
         var tasks = new List<Task>();
 
-        foreach (var message in natsMessages)
+        foreach (var message in messages)
         {
-            if (message is null) continue;
-
-            tasks.Add(this._nats.AcknowledgeMessages(message));
+            if (message is NatsBatchContainer natsMessage && !string.IsNullOrWhiteSpace(natsMessage.ReplyTo))
+            {
+                tasks.Add(this._nats.AcknowledgeMessages(natsMessage.ReplyTo));
+            }
         }
 
-        await  Task.WhenAll(tasks);
+        await Task.WhenAll(tasks);
     }
-}
+}

src/Orleans.Streaming.NATS/Providers/NatsStreamConsumer.cs

@@ -19,15 +19,15 @@ internal sealed class NatsStreamConsumer(
     string stream,
     uint partition,
     int batchSize,
-    NatsJsonContextSerializer<NatsStreamMessage> serializer)
+    INatsDeserialize<NatsStreamMessage> serializer)
 {
     private readonly ILogger _logger = loggerFactory.CreateLogger<NatsStreamConsumer>();
 
     private readonly ConsumerConfig _config = new($"orleans-{provider}-{stream}-{partition}")
     {
         FilterSubject = $"{provider}.{partition}.>",
         MaxBatch = batchSize,
-        DeliverPolicy = ConsumerConfigDeliverPolicy.LastPerSubject,
+        DeliverPolicy = ConsumerConfigDeliverPolicy.All,
         MaxAckPending = batchSize
     };
 
@@ -77,4 +77,4 @@ public async Task Initialize(CancellationToken cancellationToken = default)
 
         this._consumer = consumer;
     }
-}
+}