Skip to content

feat: add Priority to Nats queue #1071

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ereali-aneo wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

feat: add Priority to Nats queue #1071

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ece3ab3
feat add Priority to Nats queue
ereali-aneo Sep 16, 2025
bb03b2e
refacto: create stream gestion
ereali-aneo Nov 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions Adaptors/Nats/src/Nats.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

using System;

namespace ArmoniK.Core.Adapters.Nats;

/// <summary>
Expand Down Expand Up @@ -51,4 +53,16 @@ internal class Nats
/// If DegreeOfParallelism is negative, no limit is enforced.
/// </summary>
public int DegreeOfParallelism { get; set; }


/// <summary>
/// Number of priority levels supported. Each priority level will create its own Nats subject.
/// </summary>
public int MaxPriority { get; set; } = 0;

/// <summary>
/// Nats long polling wait time in seconds (5).
/// Set to 0 in order to disable long polling.
/// </summary>
public TimeSpan WaitTimeSeconds { get; set; } = TimeSpan.FromSeconds(20);
}
138 changes: 92 additions & 46 deletions Adaptors/Nats/src/PullQueueStorage.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;
Expand Down Expand Up @@ -76,7 +77,8 @@ public Task Init(CancellationToken cancellationToken)

/// <inheritdoc />
public int MaxPriority
=> int.MaxValue;
=> int.Max(options_.MaxPriority,
1);

/// <inheritdoc />
/// <remarks>
Expand All @@ -98,60 +100,104 @@ public int MaxPriority
/// </remarks>
public async IAsyncEnumerable<IQueueMessageHandler> 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<string>(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,
logger_,
cancellationToken);
}
}

/// <summary>
/// Fetches messages from a single consumer with priority information.
/// </summary>
/// <param name="consumer">The NATS consumer to fetch from.</param>
/// <param name="priority">The priority level of this consumer.</param>
/// <param name="maxMessages">Maximum number of messages to fetch.</param>
/// <param name="cancellationToken">Cancellation token.</param>
/// <returns>List of messages with their priority.</returns>
private async Task<List<(NatsJSMsg<string> msg, int priority)>> FetchWithPriorityAsync(INatsJSConsumer consumer,
int priority,
int maxMessages,
CancellationToken cancellationToken)
{
var messages = new List<(NatsJSMsg<string>, int)>();

try
{
if (options_.WaitTimeSeconds < TimeSpan.FromSeconds(1))
{
await foreach (var natsJSMsg in consumer.FetchNoWaitAsync<string>(new NatsJSFetchOpts
{
MaxMsgs = maxMessages,
},
cancellationToken: cancellationToken)
.ConfigureAwait(false))
{
messages.Add((natsJSMsg, priority));
}
}
else
{
await foreach (var natsJSMsg in consumer.FetchAsync<string>(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;
}
}
50 changes: 19 additions & 31 deletions Adaptors/Nats/src/PushQueueStorage.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
Expand Down Expand Up @@ -58,7 +59,8 @@ public PushQueueStorage(INatsJSContext js,

/// <inheritdoc />
public int MaxPriority
=> int.MaxValue;
=> int.Max(options_.MaxPriority,
1);

/// <inheritdoc />
public Task<HealthCheckResult> Check(HealthCheckTag tag)
Expand Down Expand Up @@ -89,40 +91,23 @@ public async Task PushMessagesAsync(IEnumerable<MessageData> 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);
}
}

Expand All @@ -136,11 +121,14 @@ await Publish(messages,
private async Task Publish(IEnumerable<MessageData> 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
{
Expand Down
6 changes: 3 additions & 3 deletions Adaptors/Nats/src/QueueMessageHandler.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -121,7 +121,7 @@ await message_.AckAsync()
/// Extends the acknowledgement deadline for the current message.
/// </summary>
private async Task ModifyAckDeadline(CancellationToken cancellationToken)
=> await message_.AckProgressAsync()
=> await message_.AckProgressAsync(cancellationToken: cancellationToken)
.ConfigureAwait(false);

/// <summary>
Expand Down
Loading
Loading