Merge branch 'codex/sublist-allocation-reduction'

This commit is contained in:
Joseph Doherty
2026-03-13 10:15:46 -04:00
9 changed files with 568 additions and 248 deletions

View File

@@ -0,0 +1,7 @@
namespace NATS.Server.Subscriptions;
internal readonly record struct RoutedSubKey(string RouteId, string Account, string Subject, string? Queue)
{
public static RoutedSubKey FromRemoteSubscription(RemoteSubscription sub)
=> new(sub.RouteId, sub.Account, sub.Subject, sub.Queue);
}

View File

@@ -12,11 +12,15 @@ public sealed class SubList : IDisposable
{
private const int CacheMax = 1024;
private const int CacheSweep = 256;
[ThreadStatic]
private static MatchBuilder? s_matchBuilder;
[ThreadStatic]
private static List<RoutedSubKey>? s_remoteSubRemovalKeys;
private readonly ReaderWriterLockSlim _lock = new();
private readonly TrieLevel _root = new();
private readonly SubListCacheSweeper _sweeper = new();
private readonly Dictionary<string, RemoteSubscription> _remoteSubs = new(StringComparer.Ordinal);
private readonly Dictionary<RoutedSubKey, RemoteSubscription> _remoteSubs = [];
private Dictionary<string, CachedResult>? _cache = new(StringComparer.Ordinal);
private uint _count;
private volatile bool _disposed;
@@ -31,8 +35,6 @@ public sealed class SubList : IDisposable
private readonly Dictionary<string, List<Action<bool>>> _queueRemoveNotifications = new(StringComparer.Ordinal);
private readonly record struct CachedResult(SubListResult Result, long Generation);
internal readonly record struct RoutedSubKeyInfo(string RouteId, string Account, string Subject, string? Queue);
public event Action<InterestChange>? InterestChanged;
public SubList()
@@ -178,7 +180,7 @@ public sealed class SubList : IDisposable
_lock.EnterWriteLock();
try
{
var key = BuildRoutedSubKey(sub.RouteId, sub.Account, sub.Subject, sub.Queue);
var key = RoutedSubKey.FromRemoteSubscription(sub);
var changed = false;
if (sub.IsRemoval)
{
@@ -223,7 +225,7 @@ public sealed class SubList : IDisposable
_lock.EnterWriteLock();
try
{
var key = BuildRoutedSubKey(sub.RouteId, sub.Account, sub.Subject, sub.Queue);
var key = RoutedSubKey.FromRemoteSubscription(sub);
if (!_remoteSubs.TryGetValue(key, out var existing))
return;
@@ -240,51 +242,36 @@ public sealed class SubList : IDisposable
}
}
internal static string BuildRoutedSubKey(string routeId, string account, string subject, string? queue)
=> $"{routeId}|{account}|{subject}|{queue}";
internal static string? GetAccNameFromRoutedSubKey(string routedSubKey)
=> GetRoutedSubKeyInfo(routedSubKey)?.Account;
internal static RoutedSubKeyInfo? GetRoutedSubKeyInfo(string routedSubKey)
{
if (string.IsNullOrWhiteSpace(routedSubKey))
return null;
var parts = routedSubKey.Split('|');
if (parts.Length != 4)
return null;
if (parts[0].Length == 0 || parts[1].Length == 0 || parts[2].Length == 0)
return null;
var queue = parts[3].Length == 0 ? null : parts[3];
return new RoutedSubKeyInfo(parts[0], parts[1], parts[2], queue);
}
public int RemoveRemoteSubs(string routeId)
{
_lock.EnterWriteLock();
try
{
var removalKeys = RentRemoteSubRemovalKeys();
var removed = 0;
foreach (var kvp in _remoteSubs.ToArray())
foreach (var (key, _) in _remoteSubs)
{
var info = GetRoutedSubKeyInfo(kvp.Key);
if (info == null || !string.Equals(info.Value.RouteId, routeId, StringComparison.Ordinal))
if (!string.Equals(key.RouteId, routeId, StringComparison.Ordinal))
continue;
if (_remoteSubs.Remove(kvp.Key))
removalKeys.Add(key);
}
foreach (var key in removalKeys)
{
if (_remoteSubs.Remove(key, out var removedSub))
{
removed++;
InterestChanged?.Invoke(new InterestChange(
InterestChangeKind.RemoteRemoved,
kvp.Value.Subject,
kvp.Value.Queue,
kvp.Value.Account));
removedSub.Subject,
removedSub.Queue,
removedSub.Account));
}
}
removalKeys.Clear();
if (removed > 0)
Interlocked.Increment(ref _generation);
@@ -301,30 +288,34 @@ public sealed class SubList : IDisposable
_lock.EnterWriteLock();
try
{
var removalKeys = RentRemoteSubRemovalKeys();
var removed = 0;
foreach (var kvp in _remoteSubs.ToArray())
foreach (var (key, _) in _remoteSubs)
{
var info = GetRoutedSubKeyInfo(kvp.Key);
if (info == null)
continue;
if (!string.Equals(info.Value.RouteId, routeId, StringComparison.Ordinal)
|| !string.Equals(info.Value.Account, account, StringComparison.Ordinal))
if (!string.Equals(key.RouteId, routeId, StringComparison.Ordinal)
|| !string.Equals(key.Account, account, StringComparison.Ordinal))
{
continue;
}
if (_remoteSubs.Remove(kvp.Key))
removalKeys.Add(key);
}
foreach (var key in removalKeys)
{
if (_remoteSubs.Remove(key, out var removedSub))
{
removed++;
InterestChanged?.Invoke(new InterestChange(
InterestChangeKind.RemoteRemoved,
kvp.Value.Subject,
kvp.Value.Queue,
kvp.Value.Account));
removedSub.Subject,
removedSub.Queue,
removedSub.Account));
}
}
removalKeys.Clear();
if (removed > 0)
Interlocked.Increment(ref _generation);
@@ -391,9 +382,9 @@ public sealed class SubList : IDisposable
}
else
{
var key = token.ToString();
if (!level.Nodes.TryGetValue(key, out node))
if (!TryGetLiteralNode(level, token, out _, out node))
{
var key = token.ToString();
node = new TrieNode();
level.Nodes[key] = node;
}
@@ -503,13 +494,20 @@ public sealed class SubList : IDisposable
}
else
{
level.Nodes.TryGetValue(token.ToString(), out node);
if (!TryGetLiteralNode(level, token, out var existingToken, out node))
return false;
pathList.Add((level, node, existingToken, isPwc: false, isFwc: false));
if (node.Next == null)
return false; // corrupted trie state
level = node.Next;
continue;
}
if (node == null)
return false; // not found
var tokenStr = token.ToString();
var tokenStr = isPwc ? "*" : ">";
pathList.Add((level, node, tokenStr, isPwc, isFwc));
if (node.Next == null)
return false; // corrupted trie state
@@ -587,22 +585,9 @@ public sealed class SubList : IDisposable
return cached.Result;
}
var plainSubs = new List<Subscription>();
var queueSubs = new List<List<Subscription>>();
MatchLevel(_root, tokens, 0, plainSubs, queueSubs);
SubListResult result;
if (plainSubs.Count == 0 && queueSubs.Count == 0)
{
result = SubListResult.Empty;
}
else
{
var queueSubsArr = new Subscription[queueSubs.Count][];
for (int i = 0; i < queueSubs.Count; i++)
queueSubsArr[i] = queueSubs[i].ToArray();
result = new SubListResult(plainSubs.ToArray(), queueSubsArr);
}
var builder = RentMatchBuilder();
MatchLevel(_root, tokens, 0, builder);
var result = builder.ToResult();
if (_cache != null)
{
@@ -681,6 +666,37 @@ public sealed class SubList : IDisposable
return removed;
}
private static bool TryGetLiteralNode(TrieLevel level, ReadOnlySpan<char> token, out string existingToken, out TrieNode node)
{
foreach (var (candidate, existingNode) in level.Nodes)
{
if (!SubjectMatch.TokenEquals(token, candidate))
continue;
existingToken = candidate;
node = existingNode;
return true;
}
existingToken = string.Empty;
node = null!;
return false;
}
private static MatchBuilder RentMatchBuilder()
{
var builder = s_matchBuilder ??= new MatchBuilder();
builder.Reset();
return builder;
}
private static List<RoutedSubKey> RentRemoteSubRemovalKeys()
{
var keys = s_remoteSubRemovalKeys ??= [];
keys.Clear();
return keys;
}
private bool HasExactQueueInterestNoLock(string subject, string queue)
{
var subs = new List<Subscription>();
@@ -827,6 +843,42 @@ public sealed class SubList : IDisposable
AddNodeToResults(pwc, plainSubs, queueSubs);
}
private static void MatchLevel(TrieLevel? level, string[] tokens, int tokenIndex, MatchBuilder builder)
{
TrieNode? pwc = null;
TrieNode? node = null;
for (int i = tokenIndex; i < tokens.Length; i++)
{
if (level == null)
return;
if (level.Fwc != null)
AddNodeToResults(level.Fwc, builder);
pwc = level.Pwc;
if (pwc != null)
MatchLevel(pwc.Next, tokens, i + 1, builder);
node = null;
if (level.Nodes.TryGetValue(tokens[i], out var found))
{
node = found;
level = node.Next;
}
else
{
level = null;
}
}
if (node != null)
AddNodeToResults(node, builder);
if (pwc != null)
AddNodeToResults(pwc, builder);
}
private static void AddNodeToResults(TrieNode node,
List<Subscription> plainSubs, List<List<Subscription>> queueSubs)
{
@@ -858,6 +910,19 @@ public sealed class SubList : IDisposable
}
}
private static void AddNodeToResults(TrieNode node, MatchBuilder builder)
{
builder.PlainSubs.AddRange(node.PlainSubs);
foreach (var (queueName, subs) in node.QueueSubs)
{
if (subs.Count == 0)
continue;
builder.AddQueueGroup(queueName, subs);
}
}
public SubListStats Stats()
{
_lock.EnterReadLock();
@@ -1373,4 +1438,47 @@ public sealed class SubList : IDisposable
public bool IsEmpty => PlainSubs.Count == 0 && QueueSubs.Count == 0 &&
(Next == null || (Next.Nodes.Count == 0 && Next.Pwc == null && Next.Fwc == null));
}
private sealed class MatchBuilder
{
private readonly Dictionary<string, int> _queueIndexes = new(StringComparer.Ordinal);
private readonly List<List<Subscription>> _queueGroups = [];
private int _queueGroupCount;
public List<Subscription> PlainSubs { get; } = [];
public void Reset()
{
PlainSubs.Clear();
_queueIndexes.Clear();
for (var i = 0; i < _queueGroupCount; i++)
_queueGroups[i].Clear();
_queueGroupCount = 0;
}
public void AddQueueGroup(string queueName, HashSet<Subscription> subs)
{
if (!_queueIndexes.TryGetValue(queueName, out var index))
{
index = _queueGroupCount++;
_queueIndexes[queueName] = index;
if (index == _queueGroups.Count)
_queueGroups.Add([]);
}
_queueGroups[index].AddRange(subs);
}
public SubListResult ToResult()
{
if (PlainSubs.Count == 0 && _queueGroupCount == 0)
return SubListResult.Empty;
var queueSubsArr = new Subscription[_queueGroupCount][];
for (var i = 0; i < _queueGroupCount; i++)
queueSubsArr[i] = _queueGroups[i].ToArray();
return new SubListResult(PlainSubs.ToArray(), queueSubsArr);
}
}
}

View File

@@ -249,6 +249,9 @@ public static class SubjectMatch
return tokens.Count == test.Count;
}
internal static bool TokenEquals(ReadOnlySpan<char> token, string candidate)
=> token.SequenceEqual(candidate);
private static bool TokensCanMatch(ReadOnlySpan<char> t1, ReadOnlySpan<char> t2)
{
if (t1.Length == 1 && (t1[0] == Pwc || t1[0] == Fwc))