diff --git a/Adaptors/Nats/src/Nats.cs b/Adaptors/Nats/src/Nats.cs
index 8c4fa036a..c48661821 100644
--- a/Adaptors/Nats/src/Nats.cs
+++ b/Adaptors/Nats/src/Nats.cs
@@ -15,6 +15,8 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see .
+using System;
+
namespace ArmoniK.Core.Adapters.Nats;
///
@@ -51,4 +53,16 @@ internal class Nats
/// If DegreeOfParallelism is negative, no limit is enforced.
///
public int DegreeOfParallelism { get; set; }
+
+
+ ///
+ /// Number of priority levels supported. Each priority level will create its own Nats subject.
+ ///
+ public int MaxPriority { get; set; } = 0;
+
+ ///
+ /// Nats long polling wait time in seconds (5).
+ /// Set to 0 in order to disable long polling.
+ ///
+ public TimeSpan WaitTimeSeconds { get; set; } = TimeSpan.FromSeconds(20);
}
diff --git a/Adaptors/Nats/src/PullQueueStorage.cs b/Adaptors/Nats/src/PullQueueStorage.cs
index eef7d7a0c..549e476b0 100644
--- a/Adaptors/Nats/src/PullQueueStorage.cs
+++ b/Adaptors/Nats/src/PullQueueStorage.cs
@@ -17,6 +17,7 @@
using System;
using System.Collections.Generic;
+using System.Linq;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;
@@ -76,7 +77,8 @@ public Task Init(CancellationToken cancellationToken)
///
public int MaxPriority
- => int.MaxValue;
+ => int.Max(options_.MaxPriority,
+ 1);
///
///
@@ -98,55 +100,46 @@ public int MaxPriority
///
public async IAsyncEnumerable PullMessagesAsync(string partitionId,
int nbMessages,
- [EnumeratorCancellation] CancellationToken cancellationToken = default)
+ [EnumeratorCancellation] CancellationToken cancellationToken)
{
- INatsJSConsumer? consumer;
- INatsJSStream? stream;
- try
- {
- stream = await js_.GetStreamAsync("armonik-stream")
- .ConfigureAwait(false);
- }
- catch (NatsJSApiException ex) when (ex.Error.Code == 404)
- {
- var config = new StreamConfig
- {
- Name = "armonik-stream",
- Storage = StreamConfigStorage.File,
- Subjects = new[]
- {
- partitionId,
- },
- Retention = StreamConfigRetention.Workqueue,
- };
- stream = await js_.CreateStreamAsync(config)
- .ConfigureAwait(false);
- }
+ var streamGestion = new StreamGestion(js_,
+ options_);
- try
- {
- consumer = await stream.GetConsumerAsync(partitionId)
- .ConfigureAwait(false);
- }
- catch (NatsJSApiException ex) when (ex.Error.Code == 404)
- {
- consumer = await js_.CreateConsumerAsync("armonik-stream",
- new ConsumerConfig(partitionId)
- {
- DurableName = partitionId,
- AckWait = TimeSpan.FromSeconds(options_.AckWait),
- AckPolicy = ConsumerConfigAckPolicy.Explicit,
- FilterSubject = partitionId,
- })
- .ConfigureAwait(false);
- }
+ // Ensure stream and consumers exist for all priority levels
+ await streamGestion.EnsureStreamExistsAsync(partitionId,
+ cancellationToken)
+ .ConfigureAwait(false);
+
+ var consumers = await streamGestion.EnsureConsumersExistAsync(partitionId,
+ cancellationToken)
+ .ConfigureAwait(false);
+
+ // Fetch from all consumers in parallel to reduce latency
+ var fetchTasks = consumers.Select((consumer,
+ index) =>
+ {
+ // Priority is in reverse order: index 0 = lowest priority, index MaxPriority-1 = highest priority
+ var priority = index + 1;
+ return FetchWithPriorityAsync(consumer,
+ priority,
+ nbMessages,
+ cancellationToken);
+ })
+ .ToArray();
+
+ var allResults = await Task.WhenAll(fetchTasks)
+ .ConfigureAwait(false);
- await foreach (var natsJSMsg in consumer.FetchAsync(new NatsJSFetchOpts
- {
- MaxMsgs = nbMessages,
- }))
+ // Flatten all results and sort by priority (highest first), then take only what we need
+ var sortedMessages = allResults.SelectMany(x => x)
+ .OrderByDescending(x => x.priority)
+ .Take(nbMessages)
+ .ToList();
+
+ // Yield messages in priority order
+ foreach (var (msg, _) in sortedMessages)
{
- yield return new QueueMessageHandler(natsJSMsg,
+ yield return new QueueMessageHandler(msg,
js_,
options_.AckWait,
options_.AckExtendDeadlineStep,
@@ -154,4 +147,57 @@ public async IAsyncEnumerable PullMessagesAsync(string
cancellationToken);
}
}
+
+ ///
+ /// Fetches messages from a single consumer with priority information.
+ ///
+ /// The NATS consumer to fetch from.
+ /// The priority level of this consumer.
+ /// Maximum number of messages to fetch.
+ /// Cancellation token.
+ /// List of messages with their priority.
+ private async Task msg, int priority)>> FetchWithPriorityAsync(INatsJSConsumer consumer,
+ int priority,
+ int maxMessages,
+ CancellationToken cancellationToken)
+ {
+ var messages = new List<(NatsJSMsg, int)>();
+
+ try
+ {
+ if (options_.WaitTimeSeconds < TimeSpan.FromSeconds(1))
+ {
+ await foreach (var natsJSMsg in consumer.FetchNoWaitAsync(new NatsJSFetchOpts
+ {
+ MaxMsgs = maxMessages,
+ },
+ cancellationToken: cancellationToken)
+ .ConfigureAwait(false))
+ {
+ messages.Add((natsJSMsg, priority));
+ }
+ }
+ else
+ {
+ await foreach (var natsJSMsg in consumer.FetchAsync(new NatsJSFetchOpts
+ {
+ MaxMsgs = maxMessages,
+ Expires = options_.WaitTimeSeconds,
+ },
+ cancellationToken: cancellationToken)
+ .ConfigureAwait(false))
+ {
+ messages.Add((natsJSMsg, priority));
+ }
+ }
+ }
+ catch (Exception ex)
+ {
+ logger_.LogDebug(ex,
+ "No messages available from consumer at priority {Priority}",
+ priority);
+ }
+
+ return messages;
+ }
}
diff --git a/Adaptors/Nats/src/PushQueueStorage.cs b/Adaptors/Nats/src/PushQueueStorage.cs
index fdc4fa6fb..1c58abffc 100644
--- a/Adaptors/Nats/src/PushQueueStorage.cs
+++ b/Adaptors/Nats/src/PushQueueStorage.cs
@@ -17,6 +17,7 @@
using System;
using System.Collections.Generic;
+using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
@@ -58,7 +59,8 @@ public PushQueueStorage(INatsJSContext js,
///
public int MaxPriority
- => int.MaxValue;
+ => int.Max(options_.MaxPriority,
+ 1);
///
public Task Check(HealthCheckTag tag)
@@ -89,40 +91,23 @@ public async Task PushMessagesAsync(IEnumerable messages,
{
await Publish(messages,
partitionId,
- cancellationToken);
+ cancellationToken)
+ .ConfigureAwait(false);
}
catch (Exception)
{
- try
- {
- var existing = await js_.GetStreamAsync("armonik-stream")
- .ConfigureAwait(false);
- if (!existing.Info.Config.Subjects!.Contains(partitionId))
- {
- existing.Info.Config.Subjects!.Add(partitionId);
- await js_.UpdateStreamAsync(existing.Info.Config)
- .ConfigureAwait(false);
- }
- }
- catch (NatsJSApiException ex) when (ex.Error.Code == 404)
- {
- var config = new StreamConfig
- {
- Name = "armonik-stream",
- Storage = StreamConfigStorage.File,
- Subjects = new[]
- {
- partitionId,
- },
- Retention = StreamConfigRetention.Workqueue,
- };
- await js_.CreateStreamAsync(config)
- .ConfigureAwait(false);
- }
+ // Stream or subjects might not exist, ensure they're created
+ var streamGestion = new StreamGestion(js_,
+ options_);
+ await streamGestion.EnsureStreamExistsAsync(partitionId,
+ cancellationToken)
+ .ConfigureAwait(false);
+ // Retry publishing after ensuring stream exists
await Publish(messages,
partitionId,
- cancellationToken);
+ cancellationToken)
+ .ConfigureAwait(false);
}
}
@@ -136,11 +121,14 @@ await Publish(messages,
private async Task Publish(IEnumerable messages,
string partitionId,
CancellationToken cancellationToken = default)
- => await messages.ParallelForEach(new ParallelTaskOptions(options_.DegreeOfParallelism,
+ => await messages.OrderByDescending(m => m.Options.Priority)
+ .ParallelForEach(new ParallelTaskOptions(options_.DegreeOfParallelism,
cancellationToken),
async message =>
{
- await js_.PublishAsync(partitionId,
+ var subject = StreamGestion.GetSubjectName(partitionId,
+ message.Options.Priority);
+ await js_.PublishAsync(subject,
Encoding.UTF8.GetBytes(message.TaskId),
headers: new NatsHeaders
{
diff --git a/Adaptors/Nats/src/QueueMessageHandler.cs b/Adaptors/Nats/src/QueueMessageHandler.cs
index ba6a740bb..17535afe2 100644
--- a/Adaptors/Nats/src/QueueMessageHandler.cs
+++ b/Adaptors/Nats/src/QueueMessageHandler.cs
@@ -101,13 +101,13 @@ await autoExtendAckDeadline_.Stop()
case QueueMessageStatus.Failed:
case QueueMessageStatus.Running:
case QueueMessageStatus.Postponed:
- await message_.NakAsync()
+ await message_.NakAsync(cancellationToken: CancellationToken.None)
.ConfigureAwait(false);
break;
case QueueMessageStatus.Cancelled:
case QueueMessageStatus.Processed:
case QueueMessageStatus.Poisonous:
- await message_.AckAsync()
+ await message_.AckAsync(cancellationToken: CancellationToken.None)
.ConfigureAwait(false);
break;
default:
@@ -121,7 +121,7 @@ await message_.AckAsync()
/// Extends the acknowledgement deadline for the current message.
///
private async Task ModifyAckDeadline(CancellationToken cancellationToken)
- => await message_.AckProgressAsync()
+ => await message_.AckProgressAsync(cancellationToken: cancellationToken)
.ConfigureAwait(false);
///
diff --git a/Adaptors/Nats/src/StreamGestion.cs b/Adaptors/Nats/src/StreamGestion.cs
new file mode 100644
index 000000000..21eea7083
--- /dev/null
+++ b/Adaptors/Nats/src/StreamGestion.cs
@@ -0,0 +1,190 @@
+// This file is part of the ArmoniK project
+//
+// Copyright (C) ANEO, 2021-2025. All rights reserved.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published
+// by the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see .
+
+using System;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+
+using NATS.Client.JetStream;
+using NATS.Client.JetStream.Models;
+
+namespace ArmoniK.Core.Adapters.Nats;
+
+///
+/// Helper class to manage NATS JetStream streams and consumers with priority support.
+///
+internal class StreamGestion
+{
+ private const string StreamName = "armonik-stream";
+ private const string SubjectSeparator = ":";
+
+ private readonly INatsJSContext js_;
+ private readonly Nats options_;
+
+ ///
+ /// Initializes a new instance of the class.
+ ///
+ /// The NATS JetStream context.
+ /// The NATS configuration options.
+ public StreamGestion(INatsJSContext js,
+ Nats options)
+ {
+ js_ = js;
+ options_ = options;
+ }
+
+ ///
+ /// Gets the maximum priority level. Minimum is 1.
+ ///
+ public int MaxPriority
+ => int.Max(options_.MaxPriority,
+ 1);
+
+ ///
+ /// Generates a subject name for a given partition and priority.
+ ///
+ /// The partition identifier.
+ /// The priority level (1-based).
+ /// The subject name (e.g., "partition:1").
+ public static string GetSubjectName(string partitionId,
+ int priority)
+ => partitionId + SubjectSeparator + priority;
+
+ ///
+ /// Ensures the stream exists and all priority subjects for the partition are registered.
+ ///
+ /// The partition identifier.
+ /// Cancellation token.
+ /// The stream instance.
+ public async Task EnsureStreamExistsAsync(string partitionId,
+ CancellationToken cancellationToken)
+ {
+ try
+ {
+ // Try to get existing stream
+ var stream = await js_.GetStreamAsync(StreamName,
+ cancellationToken: cancellationToken)
+ .ConfigureAwait(false);
+
+ // Ensure all priority subjects are registered
+ var updated = false;
+ for (var priority = 1; priority <= MaxPriority; priority++)
+ {
+ var subject = GetSubjectName(partitionId,
+ priority);
+ if (!stream.Info.Config.Subjects!.Contains(subject))
+ {
+ stream.Info.Config.Subjects!.Add(subject);
+ updated = true;
+ }
+ }
+
+ if (updated)
+ {
+ stream = await js_.UpdateStreamAsync(stream.Info.Config,
+ cancellationToken: cancellationToken)
+ .ConfigureAwait(false);
+ }
+
+ return stream;
+ }
+ catch (NatsJSApiException ex) when (ex.Error.Code == 404)
+ {
+ // Stream doesn't exist, try to create it
+ try
+ {
+ var subjects = Enumerable.Range(1,
+ MaxPriority)
+ .Select(i => GetSubjectName(partitionId,
+ i))
+ .ToArray();
+
+ var config = new StreamConfig
+ {
+ Name = StreamName,
+ Storage = StreamConfigStorage.File,
+ Subjects = subjects,
+ Retention = StreamConfigRetention.Workqueue,
+ };
+
+ return await js_.CreateStreamAsync(config,
+ cancellationToken)
+ .ConfigureAwait(false);
+ }
+ catch (NatsJSApiException ex2) when (ex2.Error.Code == 400)
+ {
+ // Race condition: another thread created the stream, fetch it
+ return await js_.GetStreamAsync(StreamName,
+ cancellationToken: cancellationToken)
+ .ConfigureAwait(false);
+ }
+ }
+ }
+
+ ///
+ /// Ensures consumers exist for all priority levels of a partition.
+ ///
+ /// The partition identifier.
+ /// Cancellation token.
+ /// Array of consumers, one per priority level (ordered by priority 1, 2, 3, ...).
+ public async Task EnsureConsumersExistAsync(string partitionId,
+ CancellationToken cancellationToken)
+ {
+ try
+ {
+ // Try to get existing consumers
+ return await Task.WhenAll(Enumerable.Range(1,
+ MaxPriority)
+ .Select(async priority =>
+ {
+ var consumerName = GetSubjectName(partitionId,
+ priority);
+ return await js_.GetConsumerAsync(StreamName,
+ consumerName,
+ cancellationToken)
+ .ConfigureAwait(false);
+ }))
+ .ConfigureAwait(false);
+ }
+ catch (NatsJSApiException ex) when (ex.Error.Code == 404)
+ {
+ // One or more consumers don't exist, create them all
+ return await Task.WhenAll(Enumerable.Range(1,
+ MaxPriority)
+ .Select(async priority =>
+ {
+ var consumerName = GetSubjectName(partitionId,
+ priority);
+ var subject = GetSubjectName(partitionId,
+ priority);
+
+ return await js_.CreateConsumerAsync(StreamName,
+ new ConsumerConfig(consumerName)
+ {
+ DurableName = consumerName,
+ AckWait = TimeSpan.FromSeconds(options_.AckWait),
+ AckPolicy = ConsumerConfigAckPolicy.Explicit,
+ FilterSubject = subject,
+ },
+ cancellationToken)
+ .ConfigureAwait(false);
+ }))
+ .ConfigureAwait(false);
+ }
+ }
+}
diff --git a/terraform/modules/storage/queue/nats/outputs.tf b/terraform/modules/storage/queue/nats/outputs.tf
index 629ab3096..7425feee0 100644
--- a/terraform/modules/storage/queue/nats/outputs.tf
+++ b/terraform/modules/storage/queue/nats/outputs.tf
@@ -3,6 +3,8 @@ output "generated_env_vars" {
"Components__QueueAdaptorSettings__ClassName" = "ArmoniK.Core.Adapters.Nats.QueueBuilder"
"Components__QueueAdaptorSettings__AdapterAbsolutePath" = "/adapters/queue/nats/ArmoniK.Core.Adapters.Nats.dll"
"Nats__Url" = "${var.queue_envs.host}:4222"
+ "Nats__MaxPriority" = "9"
+ "Nats__WaitTimeSeconds" = "00:00:00"
})
}
output "core_mounts" {