ElsaWorkflow/src/Elsa.CustomServer/Nats/NatsTriggerHostedService.cs

using System.Collections.Concurrent;
using System.Text;
using System.Text.Json;
using Elsa.CustomServer.Activities;
using Elsa.Workflows.Runtime;
using Elsa.Workflows.Runtime.Filters;
using Elsa.Workflows.Runtime.Requests;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using NATS.Client.JetStream;
using NATS.Client.JetStream.Models;

namespace Elsa.CustomServer.Nats;

// Consumes NATS JetStream on behalf of NatsTrigger-typed workflow triggers.
// Each stored trigger carries its own (Url, User, Password, Stream, Subject)
// stimulus; the service spins up one durable pull consumer per unique stimulus
// and dispatches a workflow per received message. Consumers use a deterministic
// DurableName so multiple Elsa instances form a load-balanced group.
public sealed class NatsTriggerHostedService(
    NatsConnectionCache connectionCache,
    IServiceProvider services,
    ILogger<NatsTriggerHostedService> log)
    : BackgroundService
{
    private readonly ConcurrentDictionary<NatsTriggerStimulus, CancellationTokenSource> _subscriptions = new();

    protected override async Task ExecuteAsync(CancellationToken stoppingToken)
    {
        try
        {
            await RefreshSubscriptionsAsync(stoppingToken);

            // Idle until cancellation. Refreshes happen out-of-band via NatsTriggerIndexChangedHandler.
            await Task.Delay(Timeout.Infinite, stoppingToken);
        }
        catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
        {
        }
    }

    public async Task RefreshSubscriptionsAsync(CancellationToken ct)
    {
        using var scope = services.CreateScope();
        var triggerStore = scope.ServiceProvider.GetRequiredService<ITriggerStore>();
        var stored = await triggerStore.FindManyAsync(new TriggerFilter { Name = NatsTrigger.StimulusName }, ct);

        var desired = new HashSet<NatsTriggerStimulus>();
        foreach (var t in stored)
        {
            var stim = TryDeserializeStimulus(t.Payload);
            if (stim is null || string.IsNullOrWhiteSpace(stim.Url) || string.IsNullOrWhiteSpace(stim.Stream) || string.IsNullOrWhiteSpace(stim.Subject))
                continue;
            desired.Add(stim);
        }

        // Stop consumers for triggers that are no longer indexed.
        foreach (var key in _subscriptions.Keys.ToList())
        {
            if (desired.Contains(key)) continue;
            if (_subscriptions.TryRemove(key, out var cts))
            {
                log.LogInformation("Stopping NATS consumer for url={Url} stream={Stream} subject={Subject}", key.Url, key.Stream, key.Subject);
                cts.Cancel();
                cts.Dispose();
            }
        }

        // Start consumers for newly indexed triggers.
        foreach (var key in desired)
        {
            if (_subscriptions.ContainsKey(key)) continue;
            var cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
            if (_subscriptions.TryAdd(key, cts))
            {
                log.LogInformation("Starting NATS consumer for url={Url} stream={Stream} subject={Subject}", key.Url, key.Stream, key.Subject);
                _ = Task.Run(() => ConsumeLoopAsync(key, cts.Token), CancellationToken.None);
            }
        }
    }

    private static NatsTriggerStimulus? TryDeserializeStimulus(object? payload)
    {
        if (payload is null) return null;
        if (payload is NatsTriggerStimulus s) return s;

        // Persistence round-trips the payload through JSON, so it typically comes back as JsonElement.
        try
        {
            var json = payload as string ?? JsonSerializer.Serialize(payload);
            return JsonSerializer.Deserialize<NatsTriggerStimulus>(json, new JsonSerializerOptions
            {
                PropertyNameCaseInsensitive = true,
            });
        }
        catch
        {
            return null;
        }
    }

    private async Task ConsumeLoopAsync(NatsTriggerStimulus stim, CancellationToken ct)
    {
        // Deterministic consumer name so every Elsa instance joins the same group.
        // Sanitize because NATS consumer names disallow '.', '*', '>', whitespace.
        var consumerName = $"elsa-{Sanitize(stim.Stream)}-{Sanitize(stim.Subject)}";

        // Reconnect-forever outer loop — catches higher-level errors (server restart, permissions).
        // NatsConnection handles TCP-level reconnects itself.
        while (!ct.IsCancellationRequested)
        {
            try
            {
                var nats = connectionCache.Get(stim.Url, stim.User, stim.Password);
                var js = new NatsJSContext(nats);

                var consumer = await js.CreateOrUpdateConsumerAsync(stim.Stream, new ConsumerConfig
                {
                    // NATS.Net 2.5 sends the create request to $JS.API.CONSUMER.CREATE.<stream>.<name>,
                    // pulling <name> from Name. Server compares it against DurableName in the payload
                    // and returns "consumer name in subject does not match durable name in request" if
                    // Name is missing. Set both to the same value.
                    Name = consumerName,
                    DurableName = consumerName,
                    AckPolicy = ConsumerConfigAckPolicy.Explicit,
                    DeliverPolicy = ConsumerConfigDeliverPolicy.New, // only applied on first-time creation
                    FilterSubject = stim.Subject,
                    MaxAckPending = 256,
                    AckWait = TimeSpan.FromSeconds(30),
                    MaxDeliver = 5,
                }, ct);

                await foreach (var msg in consumer.ConsumeAsync<byte[]>(cancellationToken: ct))
                {
                    try
                    {
                        await DispatchAsync(stim, msg, ct);
                        await msg.AckAsync(cancellationToken: ct);
                    }
                    catch (OperationCanceledException) when (ct.IsCancellationRequested)
                    {
                        throw;
                    }
                    catch (Exception ex)
                    {
                        log.LogError(ex, "NATS dispatch failed for stream={Stream} subject={Subject} seq={Seq}; NAKing", stim.Stream, msg.Subject, msg.Metadata?.Sequence.Stream);
                        try
                        {
                            await msg.NakAsync(delay: TimeSpan.FromSeconds(5), cancellationToken: ct);
                        }
                        catch (Exception nakEx)
                        {
                            log.LogWarning(nakEx, "NAK failed for stream={Stream} subject={Subject}", stim.Stream, msg.Subject);
                        }
                    }
                }
            }
            catch (OperationCanceledException) when (ct.IsCancellationRequested)
            {
                return;
            }
            catch (Exception ex)
            {
                log.LogError(ex, "NATS consume loop crashed for url={Url} stream={Stream} subject={Subject}; retrying in 5s", stim.Url, stim.Stream, stim.Subject);
                try { await Task.Delay(TimeSpan.FromSeconds(5), ct); }
                catch (OperationCanceledException) { return; }
            }
        }
    }

    private async Task DispatchAsync(NatsTriggerStimulus stim, NatsJSMsg<byte[]> msg, CancellationToken ct)
    {
        using var scope = services.CreateScope();
        var dispatcher = scope.ServiceProvider.GetRequiredService<IWorkflowDispatcher>();

        var body = msg.Data is { Length: > 0 } bytes ? Encoding.UTF8.GetString(bytes) : string.Empty;
        var headers = new Dictionary<string, string>();
        if (msg.Headers is { } h)
        {
            foreach (var key in h.Keys)
                headers[key] = h[key].ToString();
        }

        var request = new DispatchTriggerWorkflowsRequest(NatsTrigger.StimulusName, stim)
        {
            Input = new Dictionary<string, object>
            {
                [NatsTrigger.BodyInputKey] = body,
                [NatsTrigger.SubjectInputKey] = msg.Subject,
                [NatsTrigger.HeadersInputKey] = headers,
            },
        };

        await dispatcher.DispatchAsync(request, options: null!, cancellationToken: ct);
        log.LogInformation("NATS → dispatched trigger stream={Stream} subject={Subject} seq={Seq}",
            stim.Stream, msg.Subject, msg.Metadata?.Sequence.Stream);
    }

    private static string Sanitize(string value) =>
        new string(value.Select(c => char.IsLetterOrDigit(c) ? c : '_').ToArray());
}