From ece3ab31c4009cb87b8d72f8e1224ef9cbafa207 Mon Sep 17 00:00:00 2001 From: ereali-aneo Date: Tue, 16 Sep 2025 17:27:13 +0200 Subject: [PATCH 1/2] feat add Priority to Nats queue --- Adaptors/Nats/src/Nats.cs | 14 ++ Adaptors/Nats/src/PullQueueStorage.cs | 145 +++++++++++++----- Adaptors/Nats/src/PushQueueStorage.cs | 66 +++++--- Adaptors/Nats/src/QueueMessageHandler.cs | 6 +- .../modules/storage/queue/nats/outputs.tf | 2 + 5 files changed, 170 insertions(+), 63 deletions(-) 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..4ff232808 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,60 +100,127 @@ public int MaxPriority /// public async IAsyncEnumerable PullMessagesAsync(string partitionId, int nbMessages, - [EnumeratorCancellation] CancellationToken cancellationToken = default) + [EnumeratorCancellation] CancellationToken cancellationToken) { - INatsJSConsumer? consumer; - INatsJSStream? stream; + INatsJSConsumer[] prioryties; + INatsJSStream? stream; + // Create a fonc commune create stream + //var streamGestion = new StreamGestion(js_, options_); try { - stream = await js_.GetStreamAsync("armonik-stream") + stream = await js_.GetStreamAsync("armonik-stream", cancellationToken:cancellationToken) .ConfigureAwait(false); + for (var prio = 0; prio < MaxPriority; prio++) + { + var priority = prio + 1; + if (!stream.Info.Config.Subjects!.Contains(partitionId + priority)) + { + stream.Info.Config.Subjects!.Add(partitionId + priority); + await js_.UpdateStreamAsync(stream.Info.Config, cancellationToken: cancellationToken) + .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) + try + { + var subjects = Enumerable.Range(1, + MaxPriority) + .Select(i => partitionId + i) + .ToArray(); + var config = new StreamConfig + { + Name = "armonik-stream", + Storage = StreamConfigStorage.File, + Subjects = subjects, + Retention = StreamConfigRetention.Workqueue, + }; + stream = await js_.CreateStreamAsync(config, cancellationToken) + .ConfigureAwait(false); + + } + catch (NatsJSApiException) when (ex.Error.Code == 400) + { + stream = await js_.GetStreamAsync("armonik-stream", cancellationToken: cancellationToken) .ConfigureAwait(false); - } + } + catch (NatsJSApiException e2) + { + logger_.LogError(e2, + "{errorcode}", e2.Error.Code); + throw; + } + //end + } try { - consumer = await stream.GetConsumerAsync(partitionId) - .ConfigureAwait(false); + prioryties = await Task.WhenAll(Enumerable.Range(1, + MaxPriority) + .Select(async i => await stream.GetConsumerAsync(partitionId + i, cancellationToken) + .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); + prioryties = await Task.WhenAll(Enumerable.Range(1, + MaxPriority) + .Select(async i => await js_.CreateConsumerAsync("armonik-stream", + new ConsumerConfig(partitionId + i) + { + DurableName = partitionId + i, + AckWait = TimeSpan.FromSeconds(options_.AckWait), + AckPolicy = ConsumerConfigAckPolicy.Explicit, + FilterSubject = partitionId + i, + }, cancellationToken))); } - await foreach (var natsJSMsg in consumer.FetchAsync(new NatsJSFetchOpts - { - MaxMsgs = nbMessages, - })) + var i = nbMessages; + foreach (var consumer in prioryties.Reverse()) { - yield return new QueueMessageHandler(natsJSMsg, - js_, - options_.AckWait, - options_.AckExtendDeadlineStep, - logger_, - cancellationToken); + if (options_.WaitTimeSeconds < TimeSpan.FromSeconds(1)) + { + await foreach (var natsJSMsg in consumer.FetchNoWaitAsync(new NatsJSFetchOpts + { + MaxMsgs = i, + }, cancellationToken:cancellationToken) + .ConfigureAwait(false)) + { + yield return new QueueMessageHandler(natsJSMsg, + js_, + options_.AckWait, + options_.AckExtendDeadlineStep, + logger_, + cancellationToken); + --i; + if (i == 0) + { + break; + } + } + } + else + { + await foreach (var natsJSMsg in consumer.FetchAsync(new NatsJSFetchOpts + { + MaxMsgs = i, + Expires = options_.WaitTimeSeconds, + }, cancellationToken: cancellationToken) + .ConfigureAwait(false)) + { + yield return new QueueMessageHandler(natsJSMsg, + js_, + options_.AckWait, + options_.AckExtendDeadlineStep, + logger_, + cancellationToken); + --i; + if (i == 0) + { + break; + } + } + } } } } diff --git a/Adaptors/Nats/src/PushQueueStorage.cs b/Adaptors/Nats/src/PushQueueStorage.cs index fdc4fa6fb..3c10b693e 100644 --- a/Adaptors/Nats/src/PushQueueStorage.cs +++ b/Adaptors/Nats/src/PushQueueStorage.cs @@ -17,6 +17,8 @@ using System; using System.Collections.Generic; +using System.IO; +using System.Linq; using System.Text; using System.Threading; using System.Threading.Tasks; @@ -58,7 +60,8 @@ public PushQueueStorage(INatsJSContext js, /// public int MaxPriority - => int.MaxValue; + => int.Max(options_.MaxPriority, + 1); /// public Task Check(HealthCheckTag tag) @@ -85,41 +88,59 @@ public async Task PushMessagesAsync(IEnumerable messages, string partitionId, CancellationToken cancellationToken = default) { + INatsJSStream? stream; + try { 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)) + stream = await js_.GetStreamAsync("armonik-stream", cancellationToken: cancellationToken) + .ConfigureAwait(false); + for (var prio = 0; prio < MaxPriority; prio++) { - existing.Info.Config.Subjects!.Add(partitionId); - await js_.UpdateStreamAsync(existing.Info.Config) - .ConfigureAwait(false); + var priority = prio + 1; + if (!stream.Info.Config.Subjects!.Contains(partitionId + priority)) + { + stream.Info.Config.Subjects!.Add(partitionId + priority); + await js_.UpdateStreamAsync(stream.Info.Config, cancellationToken: cancellationToken) + .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); + try + { + var subjects = Enumerable.Range(1, + MaxPriority) + .Select(i => partitionId + i) + .ToArray(); + var config = new StreamConfig + { + Name = "armonik-stream", + Storage = StreamConfigStorage.File, + Subjects = subjects, + Retention = StreamConfigRetention.Workqueue, + }; + stream = await js_.CreateStreamAsync(config, cancellationToken) + .ConfigureAwait(false); + + } + catch (NatsJSApiException) when (ex.Error.Code == 400) + { + stream = await js_.GetStreamAsync("armonik-stream", cancellationToken: cancellationToken) + .ConfigureAwait(false); + } } + await Publish(messages, partitionId, cancellationToken); @@ -136,11 +157,12 @@ 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, + await js_.PublishAsync(partitionId + message.Options.Priority, 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/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" { From bb03b2e0d9a444ba35192b2386c1598e0acaecd0 Mon Sep 17 00:00:00 2001 From: ereali-aneo Date: Thu, 6 Nov 2025 10:20:27 +0100 Subject: [PATCH 2/2] refacto: create stream gestion --- Adaptors/Nats/src/PullQueueStorage.cs | 183 +++++++++++-------------- Adaptors/Nats/src/PushQueueStorage.cs | 60 ++------ Adaptors/Nats/src/StreamGestion.cs | 190 ++++++++++++++++++++++++++ 3 files changed, 283 insertions(+), 150 deletions(-) create mode 100644 Adaptors/Nats/src/StreamGestion.cs diff --git a/Adaptors/Nats/src/PullQueueStorage.cs b/Adaptors/Nats/src/PullQueueStorage.cs index 4ff232808..549e476b0 100644 --- a/Adaptors/Nats/src/PullQueueStorage.cs +++ b/Adaptors/Nats/src/PullQueueStorage.cs @@ -102,125 +102,102 @@ public async IAsyncEnumerable PullMessagesAsync(string int nbMessages, [EnumeratorCancellation] CancellationToken cancellationToken) { - INatsJSConsumer[] prioryties; - INatsJSStream? stream; - // Create a fonc commune create stream - //var streamGestion = new StreamGestion(js_, options_); - try - { - stream = await js_.GetStreamAsync("armonik-stream", cancellationToken:cancellationToken) - .ConfigureAwait(false); - for (var prio = 0; prio < MaxPriority; prio++) - { - var priority = prio + 1; - if (!stream.Info.Config.Subjects!.Contains(partitionId + priority)) - { - stream.Info.Config.Subjects!.Add(partitionId + priority); - await js_.UpdateStreamAsync(stream.Info.Config, cancellationToken: cancellationToken) - .ConfigureAwait(false); - } - } - } - catch (NatsJSApiException ex) when (ex.Error.Code == 404) - { - try - { - var subjects = Enumerable.Range(1, - MaxPriority) - .Select(i => partitionId + i) - .ToArray(); - var config = new StreamConfig - { - Name = "armonik-stream", - Storage = StreamConfigStorage.File, - Subjects = subjects, - Retention = StreamConfigRetention.Workqueue, - }; - stream = await js_.CreateStreamAsync(config, cancellationToken) - .ConfigureAwait(false); + var streamGestion = new StreamGestion(js_, + options_); - } - catch (NatsJSApiException) when (ex.Error.Code == 400) - { - stream = await js_.GetStreamAsync("armonik-stream", cancellationToken: cancellationToken) - .ConfigureAwait(false); - } + // Ensure stream and consumers exist for all priority levels + await streamGestion.EnsureStreamExistsAsync(partitionId, + cancellationToken) + .ConfigureAwait(false); - catch (NatsJSApiException e2) - { - logger_.LogError(e2, - "{errorcode}", e2.Error.Code); - throw; - } - //end - } - try - { - prioryties = await Task.WhenAll(Enumerable.Range(1, - MaxPriority) - .Select(async i => await stream.GetConsumerAsync(partitionId + i, cancellationToken) - .ConfigureAwait(false))); - } - catch (NatsJSApiException ex) when (ex.Error.Code == 404) + 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); + + // 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) { - prioryties = await Task.WhenAll(Enumerable.Range(1, - MaxPriority) - .Select(async i => await js_.CreateConsumerAsync("armonik-stream", - new ConsumerConfig(partitionId + i) - { - DurableName = partitionId + i, - AckWait = TimeSpan.FromSeconds(options_.AckWait), - AckPolicy = ConsumerConfigAckPolicy.Explicit, - FilterSubject = partitionId + i, - }, cancellationToken))); + yield return new QueueMessageHandler(msg, + js_, + options_.AckWait, + options_.AckExtendDeadlineStep, + logger_, + cancellationToken); } + } - var i = nbMessages; - foreach (var consumer in prioryties.Reverse()) + /// + /// 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)) { - MaxMsgs = i, - }, cancellationToken:cancellationToken) - .ConfigureAwait(false)) - { - yield return new QueueMessageHandler(natsJSMsg, - js_, - options_.AckWait, - options_.AckExtendDeadlineStep, - logger_, - cancellationToken); - --i; - if (i == 0) - { - break; - } + messages.Add((natsJSMsg, priority)); } } else { - await foreach (var natsJSMsg in consumer.FetchAsync(new NatsJSFetchOpts - { - MaxMsgs = i, - Expires = options_.WaitTimeSeconds, - }, cancellationToken: cancellationToken) - .ConfigureAwait(false)) - { - yield return new QueueMessageHandler(natsJSMsg, - js_, - options_.AckWait, - options_.AckExtendDeadlineStep, - logger_, - cancellationToken); - --i; - if (i == 0) - { - break; - } + 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 3c10b693e..1c58abffc 100644 --- a/Adaptors/Nats/src/PushQueueStorage.cs +++ b/Adaptors/Nats/src/PushQueueStorage.cs @@ -17,7 +17,6 @@ using System; using System.Collections.Generic; -using System.IO; using System.Linq; using System.Text; using System.Threading; @@ -88,8 +87,6 @@ public async Task PushMessagesAsync(IEnumerable messages, string partitionId, CancellationToken cancellationToken = default) { - INatsJSStream? stream; - try { await Publish(messages, @@ -99,51 +96,18 @@ await Publish(messages, } catch (Exception) { - try - { - stream = await js_.GetStreamAsync("armonik-stream", cancellationToken: cancellationToken) - .ConfigureAwait(false); - for (var prio = 0; prio < MaxPriority; prio++) - { - var priority = prio + 1; - if (!stream.Info.Config.Subjects!.Contains(partitionId + priority)) - { - stream.Info.Config.Subjects!.Add(partitionId + priority); - await js_.UpdateStreamAsync(stream.Info.Config, cancellationToken: cancellationToken) - .ConfigureAwait(false); - } - } - } - catch (NatsJSApiException ex) when (ex.Error.Code == 404) - { - try - { - var subjects = Enumerable.Range(1, - MaxPriority) - .Select(i => partitionId + i) - .ToArray(); - var config = new StreamConfig - { - Name = "armonik-stream", - Storage = StreamConfigStorage.File, - Subjects = subjects, - Retention = StreamConfigRetention.Workqueue, - }; - stream = await js_.CreateStreamAsync(config, cancellationToken) - .ConfigureAwait(false); - - } - catch (NatsJSApiException) when (ex.Error.Code == 400) - { - stream = await js_.GetStreamAsync("armonik-stream", cancellationToken: cancellationToken) - .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); } } @@ -162,7 +126,9 @@ private async Task Publish(IEnumerable messages, cancellationToken), async message => { - await js_.PublishAsync(partitionId + message.Options.Priority, + 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/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); + } + } +}