feat(batch25): implement gateway reply map and inbound message pipeline

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

@@ -0,0 +1,177 @@
+// 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;
+            }
+        });
+    }
+}