feat(batch33): merge js-cluster-streams

This commit is contained in:
Joseph Doherty
2026-02-28 23:14:30 -05:00
14 changed files with 1881 additions and 1 deletions

View File

@@ -336,6 +336,29 @@ internal sealed class JetStreamCluster
if (streams.Count == 0)
InflightConsumers.Remove(accountName);
}
internal static (
string[] NewPeers,
string[] OldPeers,
Dictionary<string, bool> NewPeerSet,
Dictionary<string, bool> OldPeerSet) GenPeerInfo(string[] peers, int split)
{
var safeSplit = Math.Clamp(split, 0, peers.Length);
var oldPeers = peers.Take(safeSplit).ToArray();
var newPeers = peers.Skip(safeSplit).ToArray();
var oldSet = oldPeers.ToDictionary(p => p, _ => true, StringComparer.Ordinal);
var newSet = newPeers.ToDictionary(p => p, _ => true, StringComparer.Ordinal);
return (newPeers, oldPeers, newSet, oldSet);
}
internal static bool IsControlHdr(byte[]? headers)
{
if (headers == null || headers.Length == 0)
return false;
var text = System.Text.Encoding.ASCII.GetString(headers);
return text.Contains(NatsHeaderConstants.JsResponseType, StringComparison.OrdinalIgnoreCase);
}
}
// ============================================================================
@@ -452,6 +475,31 @@ internal sealed class RaftGroup
/// <summary>Internal Raft node — not serialized.</summary>
[JsonIgnore]
public IRaftNode? Node { get; set; }
internal bool IsMember(string id) =>
Peers.Contains(id, StringComparer.Ordinal);
internal void SetPreferred(NatsServer server)
{
if (Peers.Length == 0)
return;
if (Peers.Length == 1)
{
Preferred = Peers[0];
return;
}
var online = new List<string>(Peers.Length);
foreach (var peer in Peers)
{
var info = server.GetNodeInfo(peer);
if (info is { Offline: false })
online.Add(peer);
}
var candidates = online.Count > 0 ? online : [.. Peers];
Preferred = candidates[Random.Shared.Next(candidates.Count)];
}
}
// ============================================================================
@@ -482,6 +530,49 @@ internal sealed class StreamAssignment
[JsonIgnore] public bool Resetting { get; set; }
[JsonIgnore] public Exception? Error { get; set; }
[JsonIgnore] public UnsupportedStreamAssignment? Unsupported { get; set; }
internal StreamAssignment CopyGroup()
{
var clone = new StreamAssignment
{
Client = Client,
Created = Created,
ConfigJson = ConfigJson,
Config = Config,
Group = Group == null
? null
: new RaftGroup
{
Name = Group.Name,
Peers = [.. Group.Peers],
Storage = Group.Storage,
Cluster = Group.Cluster,
Preferred = Group.Preferred,
ScaleUp = Group.ScaleUp,
Node = Group.Node,
},
Sync = Sync,
Subject = Subject,
Reply = Reply,
Restore = Restore,
Consumers = Consumers,
Responded = Responded,
Recovering = Recovering,
Reassigning = Reassigning,
Resetting = Resetting,
Error = Error,
Unsupported = Unsupported,
};
return clone;
}
internal bool MissingPeers() =>
Group != null &&
Config != null &&
Group.Peers.Length < Math.Max(1, Config.Replicas);
internal string RecoveryKey() =>
$"{Client?.ServiceAccount() ?? string.Empty}:{Config?.Name ?? Subject ?? string.Empty}";
}
// ============================================================================
@@ -551,6 +642,45 @@ internal sealed class ConsumerAssignment
[JsonIgnore] public bool Recovering { get; set; }
[JsonIgnore] public Exception? Error { get; set; }
[JsonIgnore] public UnsupportedConsumerAssignment? Unsupported { get; set; }
internal ConsumerAssignment CopyGroup()
{
var clone = new ConsumerAssignment
{
Client = Client,
Created = Created,
Name = Name,
Stream = Stream,
ConfigJson = ConfigJson,
Config = Config,
Group = Group == null
? null
: new RaftGroup
{
Name = Group.Name,
Peers = [.. Group.Peers],
Storage = Group.Storage,
Cluster = Group.Cluster,
Preferred = Group.Preferred,
ScaleUp = Group.ScaleUp,
Node = Group.Node,
},
Subject = Subject,
Reply = Reply,
State = State,
Responded = Responded,
Recovering = Recovering,
Error = Error,
Unsupported = Unsupported,
};
return clone;
}
internal string StreamRecoveryKey() =>
$"{Client?.ServiceAccount() ?? string.Empty}:{Stream}";
internal string RecoveryKey() =>
$"{Client?.ServiceAccount() ?? string.Empty}:{Stream}:{Name}";
}
// ============================================================================

View File

@@ -816,6 +816,964 @@ internal sealed class JetStreamEngine(JetStream state)
!string.IsNullOrWhiteSpace(value) &&
value.Length >= 10 &&
value.StartsWith("A", StringComparison.Ordinal);
// -------------------------------------------------------------------------
// JetStream cluster stream methods (Batch 33 Group A)
// -------------------------------------------------------------------------
internal void MonitorCluster()
{
var meta = GetMetaGroup();
if (meta == null)
return;
SetMetaRecovering();
try
{
CheckClusterSize();
var (_, _, _, error) = MetaSnapshot();
if (error != null)
Server()?.Warnf("JetStream meta snapshot failed in MonitorCluster: {0}", error.Message);
}
finally
{
ClearMetaRecovering();
ClusterStoppedC()?.Writer.TryWrite(true);
}
}
internal void CheckClusterSize()
{
var server = Server();
var meta = GetMetaGroup();
if (server == null || meta == null)
return;
var activePeers = server.ActivePeers();
if (activePeers.Count == 0)
return;
var jsPeers = 0;
foreach (var peer in activePeers)
{
var info = server.GetNodeInfo(peer);
if (info?.Js == true)
jsPeers++;
}
if (jsPeers > 0 && jsPeers < meta.ClusterSize())
meta.AdjustClusterSize(jsPeers);
}
internal (StreamConfig Config, bool Ok) ClusterStreamConfig(string accountName, string streamName)
{
_state.Lock.EnterReadLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster)
return (new StreamConfig(), false);
if (!cluster.Streams.TryGetValue(accountName, out var accountStreams))
return (new StreamConfig(), false);
if (!accountStreams.TryGetValue(streamName, out var assignment))
return (new StreamConfig(), false);
if (assignment.Config == null)
return (new StreamConfig(), false);
return (assignment.Config.Clone(), true);
}
finally
{
_state.Lock.ExitReadLock();
}
}
internal (byte[] Snapshot, int Streams, int Consumers, Exception? Error) MetaSnapshot()
{
_state.Lock.EnterReadLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster)
return ([], 0, 0, null);
return EncodeMetaSnapshot(cluster.Streams);
}
finally
{
_state.Lock.ExitReadLock();
}
}
internal Exception? ApplyMetaSnapshot(byte[] buffer, RecoveryUpdates? updates, bool isRecovering)
{
var (decoded, error) = DecodeMetaSnapshot(buffer);
if (error != null || decoded == null)
return error;
_state.Lock.EnterWriteLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster)
return null;
foreach (var accountStreams in decoded.Values)
{
foreach (var assignment in accountStreams.Values)
{
SetStreamAssignmentRecovering(assignment);
if (assignment.Consumers == null)
continue;
foreach (var consumer in assignment.Consumers.Values)
{
SetConsumerAssignmentRecovering(consumer);
}
}
}
cluster.Streams = decoded;
if (!isRecovering && updates != null)
{
updates.AddStreams.Clear();
updates.UpdateStreams.Clear();
updates.RemoveStreams.Clear();
updates.UpdateConsumers.Clear();
updates.RemoveConsumers.Clear();
}
return null;
}
finally
{
_state.Lock.ExitWriteLock();
}
}
internal (Dictionary<string, Dictionary<string, StreamAssignment>>? Streams, Exception? Error) DecodeMetaSnapshot(byte[] buffer)
{
if (buffer.Length == 0)
{
return (
new Dictionary<string, Dictionary<string, StreamAssignment>>(StringComparer.Ordinal),
null
);
}
try
{
var writeable = System.Text.Json.JsonSerializer.Deserialize<List<WriteableStreamAssignment>>(buffer) ?? [];
var streams = new Dictionary<string, Dictionary<string, StreamAssignment>>(StringComparer.Ordinal);
foreach (var item in writeable)
{
var accountName = item.Client?.ServiceAccount() ?? string.Empty;
if (!streams.TryGetValue(accountName, out var accountStreams))
{
accountStreams = new Dictionary<string, StreamAssignment>(StringComparer.Ordinal);
streams[accountName] = accountStreams;
}
var assignment = new StreamAssignment
{
Client = item.Client,
Created = item.Created,
ConfigJson = item.ConfigJson,
Group = item.Group,
Sync = item.Sync,
Config = item.ConfigJson.ValueKind == System.Text.Json.JsonValueKind.Undefined
? null
: System.Text.Json.JsonSerializer.Deserialize<StreamConfig>(item.ConfigJson.GetRawText()),
Consumers = new Dictionary<string, ConsumerAssignment>(StringComparer.Ordinal),
};
if (item.Consumers.Count > 0)
{
foreach (var consumer in item.Consumers)
{
var streamName = !string.IsNullOrWhiteSpace(consumer.Stream)
? consumer.Stream
: assignment.Config?.Name ?? string.Empty;
var assignmentConsumer = new ConsumerAssignment
{
Client = consumer.Client,
Created = consumer.Created,
Name = consumer.Name,
Stream = streamName,
ConfigJson = consumer.ConfigJson,
Group = consumer.Group,
State = consumer.State,
Config = consumer.ConfigJson.ValueKind == System.Text.Json.JsonValueKind.Undefined
? null
: System.Text.Json.JsonSerializer.Deserialize<ConsumerConfig>(consumer.ConfigJson.GetRawText()),
};
assignment.Consumers[assignmentConsumer.Name] = assignmentConsumer;
}
}
var streamKey = assignment.Config?.Name ?? string.Empty;
accountStreams[streamKey] = assignment;
}
return (streams, null);
}
catch (Exception ex)
{
return (null, ex);
}
}
internal (byte[] Snapshot, int Streams, int Consumers, Exception? Error) EncodeMetaSnapshot(
Dictionary<string, Dictionary<string, StreamAssignment>> streams)
{
try
{
var streamCount = 0;
var consumerCount = 0;
var writeable = new List<WriteableStreamAssignment>();
foreach (var accountStreams in streams.Values)
{
foreach (var assignment in accountStreams.Values)
{
streamCount++;
var configJson = assignment.ConfigJson.ValueKind == System.Text.Json.JsonValueKind.Undefined &&
assignment.Config != null
? System.Text.Json.JsonSerializer.SerializeToElement(assignment.Config)
: assignment.ConfigJson;
var ws = new WriteableStreamAssignment
{
Client = assignment.Client,
Created = assignment.Created,
ConfigJson = configJson,
Group = assignment.Group,
Sync = assignment.Sync,
};
if (assignment.Consumers != null)
{
foreach (var consumer in assignment.Consumers.Values)
{
consumerCount++;
var consumerConfigJson =
consumer.ConfigJson.ValueKind == System.Text.Json.JsonValueKind.Undefined &&
consumer.Config != null
? System.Text.Json.JsonSerializer.SerializeToElement(consumer.Config)
: consumer.ConfigJson;
ws.Consumers.Add(new WriteableConsumerAssignment
{
Client = consumer.Client,
Created = consumer.Created,
Name = consumer.Name,
Stream = consumer.Stream,
ConfigJson = consumerConfigJson,
Group = consumer.Group,
State = consumer.State,
});
}
}
writeable.Add(ws);
}
}
if (writeable.Count == 0)
return ([], 0, 0, null);
var snapshot = System.Text.Json.JsonSerializer.SerializeToUtf8Bytes(writeable);
return (snapshot, streamCount, consumerCount, null);
}
catch (Exception ex)
{
return ([], 0, 0, ex);
}
}
internal Exception? CollectStreamAndConsumerChanges(
IRaftNodeCheckpoint checkpoint,
Dictionary<string, Dictionary<string, StreamAssignment>> streams)
{
try
{
foreach (var (appendEntry, error) in checkpoint.AppendEntriesSeq())
{
if (error != null)
return error;
foreach (var entry in appendEntry.Entries)
{
if (entry.Type != EntryType.EntryNormal || entry.Data.Length == 0)
continue;
var op = (EntryOp)entry.Data[0];
var payload = entry.Data.Length > 1 ? entry.Data.AsSpan(1).ToArray() : [];
switch (op)
{
case EntryOp.AssignStreamOp:
case EntryOp.UpdateStreamOp:
{
var sa = System.Text.Json.JsonSerializer.Deserialize<StreamAssignment>(payload);
if (sa?.Client == null || sa.Config == null)
break;
var account = sa.Client.ServiceAccount();
if (!streams.TryGetValue(account, out var accountStreams))
{
accountStreams = new Dictionary<string, StreamAssignment>(StringComparer.Ordinal);
streams[account] = accountStreams;
}
accountStreams[sa.Config.Name] = sa;
break;
}
case EntryOp.RemoveStreamOp:
{
var sa = System.Text.Json.JsonSerializer.Deserialize<StreamAssignment>(payload);
if (sa?.Client == null || sa.Config == null)
break;
if (streams.TryGetValue(sa.Client.ServiceAccount(), out var accountStreams))
{
accountStreams.Remove(sa.Config.Name);
}
break;
}
case EntryOp.AssignConsumerOp:
case EntryOp.AssignCompressedConsumerOp:
{
var ca = System.Text.Json.JsonSerializer.Deserialize<ConsumerAssignment>(payload);
if (ca?.Client == null || string.IsNullOrWhiteSpace(ca.Stream))
break;
var account = ca.Client.ServiceAccount();
if (!streams.TryGetValue(account, out var accountStreams))
break;
if (!accountStreams.TryGetValue(ca.Stream, out var streamAssignment))
break;
streamAssignment.Consumers ??= new Dictionary<string, ConsumerAssignment>(StringComparer.Ordinal);
streamAssignment.Consumers[ca.Name] = ca;
break;
}
case EntryOp.RemoveConsumerOp:
{
var ca = System.Text.Json.JsonSerializer.Deserialize<ConsumerAssignment>(payload);
if (ca?.Client == null || string.IsNullOrWhiteSpace(ca.Stream))
break;
var account = ca.Client.ServiceAccount();
if (!streams.TryGetValue(account, out var accountStreams))
break;
if (!accountStreams.TryGetValue(ca.Stream, out var streamAssignment))
break;
streamAssignment.Consumers?.Remove(ca.Name);
break;
}
}
}
}
return null;
}
catch (Exception ex)
{
return ex;
}
}
internal void SetStreamAssignmentRecovering(StreamAssignment assignment)
{
assignment.Responded = true;
assignment.Recovering = true;
assignment.Restore = null;
if (assignment.Group != null)
{
assignment.Group.Preferred = string.Empty;
assignment.Group.ScaleUp = false;
}
}
internal void SetConsumerAssignmentRecovering(ConsumerAssignment assignment)
{
assignment.Responded = true;
assignment.Recovering = true;
if (assignment.Group != null)
{
assignment.Group.Preferred = string.Empty;
assignment.Group.ScaleUp = false;
}
}
internal void ProcessAddPeer(string peer)
{
_state.Lock.EnterWriteLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster)
return;
if (!cluster.IsLeader())
return;
foreach (var accountStreams in cluster.Streams.Values)
{
foreach (var assignment in accountStreams.Values)
{
if (!assignment.MissingPeers())
continue;
var copy = assignment.CopyGroup();
copy.Group!.Peers = [.. copy.Group.Peers, peer];
assignment.Group = copy.Group;
if (assignment.Consumers == null)
continue;
foreach (var consumer in assignment.Consumers.Values)
{
if (consumer.Config?.Durable is { Length: > 0 } || (consumer.Group?.Peers.Length ?? 0) > 1)
consumer.Group!.Peers = assignment.Group.Peers;
}
}
}
}
finally
{
_state.Lock.ExitWriteLock();
}
}
internal void ProcessRemovePeer(string peer)
{
_state.Lock.EnterWriteLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster)
return;
foreach (var accountStreams in cluster.Streams.Values)
{
foreach (var assignment in accountStreams.Values)
{
if (assignment.Group == null || !assignment.Group.IsMember(peer))
continue;
RemovePeerFromStreamLocked(assignment, peer);
}
}
}
finally
{
_state.Lock.ExitWriteLock();
}
}
internal bool RemovePeerFromStream(StreamAssignment assignment, string peer)
{
_state.Lock.EnterWriteLock();
try
{
return RemovePeerFromStreamLocked(assignment, peer);
}
finally
{
_state.Lock.ExitWriteLock();
}
}
internal bool RemovePeerFromStreamLocked(StreamAssignment assignment, string peer)
{
if (assignment.Group == null || !assignment.Group.IsMember(peer))
return false;
assignment.Group.Peers = [.. assignment.Group.Peers.Where(p => !string.Equals(p, peer, StringComparison.Ordinal))];
assignment.Group.Preferred = string.Empty;
if (assignment.Consumers == null)
return true;
foreach (var consumer in assignment.Consumers.Values)
{
if (consumer.Group == null)
continue;
consumer.Group.Peers = [.. consumer.Group.Peers.Where(p => !string.Equals(p, peer, StringComparison.Ordinal))];
consumer.Group.Preferred = string.Empty;
}
return true;
}
internal bool HasPeerEntries(IEnumerable<Entry> entries)
{
foreach (var entry in entries)
{
if (entry.Type == EntryType.EntryAddPeer || entry.Type == EntryType.EntryRemovePeer)
return true;
}
return false;
}
internal (bool IsRecovering, bool DidSnapshot, Exception? Error) ApplyMetaEntries(
IReadOnlyList<Entry> entries,
RecoveryUpdates? updates)
{
var recovering = updates != null;
var didSnapshot = false;
try
{
foreach (var entry in entries)
{
if (entry.Type == EntryType.EntryCatchup)
{
recovering = true;
break;
}
if (entry.Type == EntryType.EntrySnapshot)
{
var error = ApplyMetaSnapshot(entry.Data, updates, recovering);
if (error != null)
return (recovering, didSnapshot, error);
didSnapshot = true;
continue;
}
if (entry.Type == EntryType.EntryRemovePeer)
{
ProcessRemovePeer(System.Text.Encoding.ASCII.GetString(entry.Data));
continue;
}
if (entry.Type == EntryType.EntryAddPeer)
{
ProcessAddPeer(System.Text.Encoding.ASCII.GetString(entry.Data));
continue;
}
if (entry.Type != EntryType.EntryNormal || entry.Data.Length == 0)
continue;
var op = (EntryOp)entry.Data[0];
var payload = entry.Data.Length > 1 ? entry.Data.AsSpan(1).ToArray() : [];
switch (op)
{
case EntryOp.AssignStreamOp:
{
var sa = System.Text.Json.JsonSerializer.Deserialize<StreamAssignment>(payload);
if (sa != null)
SetStreamAssignmentRecovering(sa);
break;
}
case EntryOp.AssignConsumerOp:
case EntryOp.AssignCompressedConsumerOp:
{
var ca = System.Text.Json.JsonSerializer.Deserialize<ConsumerAssignment>(payload);
if (ca != null)
SetConsumerAssignmentRecovering(ca);
break;
}
}
}
return (recovering, didSnapshot, null);
}
catch (Exception ex)
{
return (recovering, didSnapshot, ex);
}
}
internal (IRaftNode? Node, Exception? Error) CreateRaftGroup(
string accountName,
RaftGroup group,
bool recovering,
StorageType storage)
{
var server = Server();
if (server == null)
return (null, new InvalidOperationException("jetstream server unavailable"));
if (group.Peers.Length <= 1)
return (null, null);
var meta = GetMetaGroup();
if (meta == null || !group.IsMember(meta.ID()))
return (null, null);
var existing = server.LookupRaftNode(group.Name);
if (existing != null && existing.State() != RaftState.Closed)
{
group.Node = existing;
return (existing, null);
}
group.SetPreferred(server);
var config = new RaftConfig
{
Name = group.Name,
Recovering = recovering,
ScaleUp = group.ScaleUp,
Store = Path.Combine(_state.Config.StoreDir, accountName, "_js_", group.Name),
Log = null,
};
var (node, error) = server.StartRaftNode(accountName, config);
if (error != null || node == null)
return (null, error);
group.Node = node;
if (!string.IsNullOrEmpty(group.Preferred) && node.ID() == group.Preferred && node.Term() == 0)
node.CampaignImmediately();
_ = storage;
return (node, null);
}
internal static (
string[] NewPeers,
string[] OldPeers,
Dictionary<string, bool> NewPeerSet,
Dictionary<string, bool> OldPeerSet) GenPeerInfo(string[] peers, int split)
{
var safeSplit = Math.Clamp(split, 0, peers.Length);
var oldPeers = peers.Take(safeSplit).ToArray();
var newPeers = peers.Skip(safeSplit).ToArray();
var oldSet = oldPeers.ToDictionary(p => p, _ => true, StringComparer.Ordinal);
var newSet = newPeers.ToDictionary(p => p, _ => true, StringComparer.Ordinal);
return (newPeers, oldPeers, newSet, oldSet);
}
internal void MonitorStream(NatsStream? stream, StreamAssignment assignment, bool sendSnapshot)
{
var node = assignment.Group?.Node;
if (node == null)
return;
var isLeader = node.Leader();
ProcessStreamLeaderChange(stream, isLeader);
if (sendSnapshot && stream != null && isLeader)
{
var state = stream.State();
var snap = System.Text.Json.JsonSerializer.SerializeToUtf8Bytes(state);
node.SendSnapshot(snap);
}
}
internal static bool IsControlHdr(byte[]? headers)
{
if (headers == null || headers.Length == 0)
return false;
var text = System.Text.Encoding.ASCII.GetString(headers);
return text.Contains(NatsHeaderConstants.JsResponseType, StringComparison.OrdinalIgnoreCase);
}
internal (ulong MaxApplied, Exception? Error) ApplyStreamEntries(
NatsStream? stream,
CommittedEntry committed,
bool isRecovering)
{
try
{
if (stream == null)
return (0, null);
ulong maxApplied = 0;
foreach (var entry in committed.Entries)
{
if (entry.Type != EntryType.EntryNormal || entry.Data.Length == 0)
continue;
var op = (EntryOp)entry.Data[0];
if (op == EntryOp.StreamMsgOp || op == EntryOp.CompressedStreamMsgOp)
{
var error = ApplyStreamMsgOp(stream, entry, isRecovering);
if (error != null)
return (maxApplied, error);
maxApplied = Math.Max(maxApplied, committed.Index);
}
}
return (maxApplied, null);
}
catch (Exception ex)
{
return (0, ex);
}
}
internal Exception? ApplyStreamMsgOp(NatsStream stream, Entry entry, bool isRecovering)
{
if (isRecovering && stream.SkipBatchIfRecovering())
return null;
var payloadLength = Math.Max(0, entry.Data.Length - 1);
Interlocked.Increment(ref stream.Msgs);
Interlocked.Add(ref stream.Bytes, payloadLength);
return null;
}
internal void ProcessStreamLeaderChange(NatsStream? stream, bool isLeader)
{
if (stream == null)
return;
stream.SetLeader(isLeader, term: 0);
if (!isLeader)
stream.ResetClusteredState();
}
internal StreamAssignment? StreamAssignment(string accountName, string streamName)
{
_state.Lock.EnterReadLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster)
return null;
if (!cluster.Streams.TryGetValue(accountName, out var accountStreams))
return null;
return accountStreams.TryGetValue(streamName, out var assignment) ? assignment : null;
}
finally
{
_state.Lock.ExitReadLock();
}
}
internal StreamAssignment? StreamAssignmentOrInflight(string accountName, string streamName)
{
_state.Lock.EnterReadLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster)
return null;
if (cluster.Streams.TryGetValue(accountName, out var accountStreams) &&
accountStreams.TryGetValue(streamName, out var current))
return current;
if (cluster.InflightStreams.TryGetValue(accountName, out var inflight) &&
inflight.TryGetValue(streamName, out var info))
return info.Assignment;
return null;
}
finally
{
_state.Lock.ExitReadLock();
}
}
internal IEnumerable<StreamAssignment> StreamAssignmentsOrInflightSeq(string accountName)
{
_state.Lock.EnterReadLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster)
return [];
var seen = new HashSet<string>(StringComparer.Ordinal);
var items = new List<StreamAssignment>();
if (cluster.Streams.TryGetValue(accountName, out var accountStreams))
{
foreach (var (streamName, assignment) in accountStreams)
{
seen.Add(streamName);
items.Add(assignment);
}
}
if (cluster.InflightStreams.TryGetValue(accountName, out var inflightStreams))
{
foreach (var (streamName, inflight) in inflightStreams)
{
if (seen.Contains(streamName) || inflight.Assignment == null)
continue;
items.Add(inflight.Assignment);
}
}
return items;
}
finally
{
_state.Lock.ExitReadLock();
}
}
internal IEnumerable<(string Account, StreamAssignment Assignment)> StreamAssignmentsOrInflightSeqAllAccounts()
{
_state.Lock.EnterReadLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster)
return [];
var results = new List<(string Account, StreamAssignment Assignment)>();
foreach (var account in cluster.Streams.Keys.Union(cluster.InflightStreams.Keys, StringComparer.Ordinal))
{
var seen = new HashSet<string>(StringComparer.Ordinal);
if (cluster.Streams.TryGetValue(account, out var accountStreams))
{
foreach (var (streamName, assignment) in accountStreams)
{
seen.Add(streamName);
results.Add((account, assignment));
}
}
if (cluster.InflightStreams.TryGetValue(account, out var inflightStreams))
{
foreach (var (streamName, inflight) in inflightStreams)
{
if (seen.Contains(streamName) || inflight.Assignment == null)
continue;
results.Add((account, inflight.Assignment));
}
}
}
return results;
}
finally
{
_state.Lock.ExitReadLock();
}
}
internal void ProcessStreamAssignment(StreamAssignment assignment)
{
_state.Lock.EnterWriteLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster || assignment.Client == null)
return;
var accountName = assignment.Client.ServiceAccount();
if (!cluster.Streams.TryGetValue(accountName, out var accountStreams))
{
accountStreams = new Dictionary<string, StreamAssignment>(StringComparer.Ordinal);
cluster.Streams[accountName] = accountStreams;
}
var streamName = assignment.Config?.Name ?? assignment.Subject ?? string.Empty;
assignment.Responded = true;
accountStreams[streamName] = assignment;
cluster.RemoveInflightStreamProposal(accountName, streamName);
}
finally
{
_state.Lock.ExitWriteLock();
}
}
internal void ProcessUpdateStreamAssignment(StreamAssignment assignment)
{
ProcessStreamAssignment(assignment);
}
internal void ProcessClusterUpdateStream(StreamAssignment assignment)
{
ProcessUpdateStreamAssignment(assignment);
}
internal void ProcessClusterCreateStream(StreamAssignment assignment)
{
ProcessStreamAssignment(assignment);
}
internal void ProcessStreamRemoval(StreamAssignment assignment)
{
_state.Lock.EnterWriteLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster || assignment.Client == null)
return;
var accountName = assignment.Client.ServiceAccount();
var streamName = assignment.Config?.Name ?? assignment.Subject ?? string.Empty;
if (cluster.Streams.TryGetValue(accountName, out var accountStreams))
{
accountStreams.Remove(streamName);
if (accountStreams.Count == 0)
cluster.Streams.Remove(accountName);
}
cluster.RemoveInflightStreamProposal(accountName, streamName);
}
finally
{
_state.Lock.ExitWriteLock();
}
}
internal void ProcessClusterDeleteStream(StreamAssignment assignment)
{
ProcessStreamRemoval(assignment);
}
internal void ProcessConsumerAssignment(ConsumerAssignment assignment)
{
_state.Lock.EnterWriteLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster || assignment.Client == null)
return;
var accountName = assignment.Client.ServiceAccount();
if (!cluster.Streams.TryGetValue(accountName, out var accountStreams))
return;
if (!accountStreams.TryGetValue(assignment.Stream, out var streamAssignment))
return;
streamAssignment.Consumers ??= new Dictionary<string, ConsumerAssignment>(StringComparer.Ordinal);
assignment.Responded = true;
streamAssignment.Consumers[assignment.Name] = assignment;
cluster.RemoveInflightConsumerProposal(accountName, assignment.Stream, assignment.Name);
}
finally
{
_state.Lock.ExitWriteLock();
}
}
internal void ProcessConsumerRemoval(ConsumerAssignment assignment)
{
_state.Lock.EnterWriteLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster || assignment.Client == null)
return;
var accountName = assignment.Client.ServiceAccount();
if (cluster.Streams.TryGetValue(accountName, out var accountStreams) &&
accountStreams.TryGetValue(assignment.Stream, out var streamAssignment))
{
streamAssignment.Consumers?.Remove(assignment.Name);
}
cluster.RemoveInflightConsumerProposal(accountName, assignment.Stream, assignment.Name);
}
finally
{
_state.Lock.ExitWriteLock();
}
}
internal void ProcessClusterCreateConsumer(ConsumerAssignment assignment)
{
ProcessConsumerAssignment(assignment);
}
internal void ProcessClusterDeleteConsumer(ConsumerAssignment assignment)
{
ProcessConsumerRemoval(assignment);
}
internal ConsumerAssignment? ConsumerAssignment(string accountName, string streamName, string consumerName)
{
_state.Lock.EnterReadLock();
try
{
if (_state.Cluster is not JetStreamCluster cluster)
return null;
if (!cluster.Streams.TryGetValue(accountName, out var accountStreams))
return null;
if (!accountStreams.TryGetValue(streamName, out var streamAssignment))
return null;
if (streamAssignment.Consumers == null)
return null;
return streamAssignment.Consumers.TryGetValue(consumerName, out var assignment) ? assignment : null;
}
finally
{
_state.Lock.ExitReadLock();
}
}
}
internal sealed class StreamAssignmentView

View File

@@ -45,6 +45,9 @@ internal sealed class NatsStream : IDisposable
/// <summary>IRaftNode — stored as object to avoid cross-dependency on Raft session.</summary>
private object? _node;
private StreamAssignment? _assignment;
private bool _migrating;
private bool _recovering;
public NatsStream(Account account, StreamConfig config, DateTime created)
{
@@ -79,6 +82,7 @@ internal sealed class NatsStream : IDisposable
{
Store = store,
IsMirror = cfg.Mirror != null,
_assignment = sa,
};
return stream;
}
@@ -317,6 +321,97 @@ internal sealed class NatsStream : IDisposable
finally { _mu.ExitReadLock(); }
}
public RaftGroup? RaftGroup()
{
_mu.EnterReadLock();
try { return _assignment?.Group; }
finally { _mu.ExitReadLock(); }
}
public IRaftNode? RaftNode()
{
_mu.EnterReadLock();
try { return _node as IRaftNode; }
finally { _mu.ExitReadLock(); }
}
public void RemoveNode()
{
_mu.EnterWriteLock();
try
{
if (_node is IRaftNode raft)
raft.Delete();
_node = null;
}
finally
{
_mu.ExitWriteLock();
}
}
public void WaitOnConsumerAssignments(CancellationToken cancellationToken = default)
{
if (cancellationToken.IsCancellationRequested)
return;
var stopAt = DateTime.UtcNow.AddSeconds(2);
while (DateTime.UtcNow < stopAt)
{
cancellationToken.ThrowIfCancellationRequested();
if (!_recovering)
break;
Thread.Sleep(50);
}
}
public bool IsMigrating()
{
_mu.EnterReadLock();
try { return _migrating; }
finally { _mu.ExitReadLock(); }
}
public bool ResetClusteredState(Exception? cause = null)
{
_mu.EnterWriteLock();
try
{
_recovering = true;
_isLeader = false;
_leaderTerm = 0;
_migrating = false;
if (cause != null && _node is IRaftNode raft)
raft.StepDown();
return true;
}
finally
{
_mu.ExitWriteLock();
}
}
public bool SkipBatchIfRecovering()
{
_mu.EnterReadLock();
try { return _recovering; }
finally { _mu.ExitReadLock(); }
}
public bool ShouldSendLostQuorum()
{
_mu.EnterReadLock();
try
{
var replicas = Math.Max(1, Config.Replicas);
return replicas > 1 && _node is IRaftNode raft && raft.Leaderless();
}
finally
{
_mu.ExitReadLock();
}
}
/// <summary>
/// Seals the stream so that no new messages can be stored.
/// Mirrors <c>stream.seal</c> in server/stream.go.

View File

@@ -0,0 +1,38 @@
namespace ZB.MOM.NatsNet.Server;
public sealed partial class NatsServer
{
internal int Replicas(StreamConfig? config)
{
if (config == null)
return 1;
return Math.Max(1, config.Replicas);
}
internal void SendStreamLostQuorumAdvisory(Account? account, string stream, string[]? peers = null)
{
_ = account;
_ = peers;
Noticef("JetStream stream lost quorum advisory for stream {0}", stream);
}
internal void SendStreamLeaderElectAdvisory(Account? account, string stream, string leader)
{
_ = account;
Noticef("JetStream stream leader elect advisory for stream {0}, leader {1}", stream, leader);
}
internal bool RemoveStream(Account? account, string streamName)
{
if (account == null)
return false;
var (stream, _) = account.LookupStream(streamName);
if (stream == null)
return false;
stream.Delete();
return true;
}
}