natsnet/dotnet/src/ZB.MOM.NatsNet.Server/NatsServer.Gateways.ReplyMap.cs

// Copyright 2018-2026 The NATS Authors
// Licensed under the Apache License, Version 2.0

using System.Collections.Concurrent;
using System.Text;
using System.Threading.Channels;

namespace ZB.MOM.NatsNet.Server;

public sealed partial class NatsServer
{
    private readonly ConcurrentDictionary<GwReplyMapping, object> _gwReplyMappings = new();
    private readonly Channel<TimeSpan> _gwReplyMapTtlUpdates = Channel.CreateBounded<TimeSpan>(1);
    private int _gwReplyMapWorkerRunning;

    internal void StoreRouteByHash(string serverIdHash, ClientConnection route)
    {
        if (!_gateway.Enabled || string.IsNullOrWhiteSpace(serverIdHash))
            return;

        _gateway.RoutesIdByHash[serverIdHash] = route;
    }

    internal void RemoveRouteByHash(string serverIdHash)
    {
        if (!_gateway.Enabled || string.IsNullOrWhiteSpace(serverIdHash))
            return;

        _gateway.RoutesIdByHash.TryRemove(serverIdHash, out _);
    }

    internal (ClientConnection? Route, bool PerAccount) GetRouteByHash(byte[] hash, byte[] accountName)
    {
        if (hash.Length == 0)
            return (null, false);

        var id = Encoding.ASCII.GetString(hash);
        var perAccount = false;
        var accountKey = Encoding.ASCII.GetString(accountName);
        if (_accRouteByHash.TryGetValue(accountKey, out var accountRouteEntry))
        {
            if (accountRouteEntry == null)
            {
                id += accountKey;
                perAccount = true;
            }
            else if (accountRouteEntry is int routeIndex)
            {
                id += routeIndex.ToString();
            }
        }

        if (_gateway.RoutesIdByHash.TryGetValue(id, out var route))
            return (route, perAccount);

        if (!perAccount && _gateway.RoutesIdByHash.TryGetValue($"{Encoding.ASCII.GetString(hash)}0", out var noPoolRoute))
        {
            lock (noPoolRoute)
            {
                if (noPoolRoute.Route?.NoPool == true)
                    return (noPoolRoute, false);
            }
        }

        return (null, perAccount);
    }

    internal void TrackGWReply(ClientConnection? client, Account? account, byte[] reply, byte[] routedReply)
    {
        GwReplyMapping? mapping = null;
        object? locker = null;

        if (account != null)
        {
            mapping = account.GwReplyMapping;
            locker = account;
        }
        else if (client != null)
        {
            mapping = client.GwReplyMapping;
            locker = client;
        }

        if (mapping == null || locker == null || reply.Length == 0 || routedReply.Length == 0)
            return;

        var ttl = _gateway.RecSubExp <= TimeSpan.Zero ? TimeSpan.FromSeconds(2) : _gateway.RecSubExp;
        lock (locker)
        {
            var wasEmpty = mapping.Mapping.Count == 0;
            var maxMappedLen = Math.Min(routedReply.Length, GatewayHandler.GwSubjectOffset + reply.Length);
            var mappedSubject = Encoding.ASCII.GetString(routedReply, 0, maxMappedLen);
            var key = mappedSubject.Length > GatewayHandler.GwSubjectOffset
                ? mappedSubject[GatewayHandler.GwSubjectOffset..]
                : mappedSubject;

            mapping.Mapping[key] = new GwReplyMap
            {
                Ms = mappedSubject,
                Exp = DateTime.UtcNow.Add(ttl).Ticks,
            };

            if (wasEmpty)
            {
                Interlocked.Exchange(ref mapping.Check, 1);
                _gwReplyMappings[mapping] = locker;
                if (Interlocked.CompareExchange(ref _gwReplyMapWorkerRunning, 1, 0) == 0)
                {
                    if (!_gwReplyMapTtlUpdates.Writer.TryWrite(ttl))
                    {
                        while (_gwReplyMapTtlUpdates.Reader.TryRead(out _)) { }
                        _gwReplyMapTtlUpdates.Writer.TryWrite(ttl);
                    }
                    StartGWReplyMapExpiration();
                }
            }
        }
    }

    internal void StartGWReplyMapExpiration()
    {
        _ = StartGoRoutine(() =>
        {
            var ttl = TimeSpan.Zero;
            var token = _quitCts.Token;

            while (!token.IsCancellationRequested)
            {
                try
                {
                    if (ttl == TimeSpan.Zero)
                    {
                        ttl = _gwReplyMapTtlUpdates.Reader.ReadAsync(token).AsTask().GetAwaiter().GetResult();
                    }

                    Task.Delay(ttl, token).GetAwaiter().GetResult();
                }
                catch (OperationCanceledException)
                {
                    break;
                }

                while (_gwReplyMapTtlUpdates.Reader.TryRead(out var nextTtl))
                    ttl = nextTtl;

                var nowTicks = DateTime.UtcNow.Ticks;
                var hasMappings = false;
                foreach (var entry in _gwReplyMappings.ToArray())
                {
                    var mapping = entry.Key;
                    var mapLocker = entry.Value;
                    lock (mapLocker)
                    {
                        foreach (var key in mapping.Mapping.Keys.ToArray())
                        {
                            if (mapping.Mapping[key].Exp <= nowTicks)
                                mapping.Mapping.Remove(key);
                        }

                        if (mapping.Mapping.Count == 0)
                        {
                            Interlocked.Exchange(ref mapping.Check, 0);
                            _gwReplyMappings.TryRemove(mapping, out _);
                        }
                        else
                        {
                            hasMappings = true;
                        }
                    }
                }

                if (!hasMappings && Interlocked.CompareExchange(ref _gwReplyMapWorkerRunning, 0, 1) == 1)
                    ttl = TimeSpan.Zero;
            }
        });
    }
}