Files
natsnet/dotnet/src/ZB.MOM.NatsNet.Server/JetStream/JetStreamEngine.cs

150 lines
4.2 KiB
C#

namespace ZB.MOM.NatsNet.Server;
internal sealed class JetStreamEngine(JetStream state)
{
private readonly JetStream _state = state;
private static readonly TimeSpan MinUsageUpdateWindow = TimeSpan.FromMilliseconds(250);
internal void SetStarted()
{
_state.Lock.EnterWriteLock();
try
{
_state.Started = DateTime.UtcNow;
}
finally
{
_state.Lock.ExitWriteLock();
}
}
internal bool IsEnabled() => Interlocked.CompareExchange(ref _state.Disabled, 0, 0) == 0;
internal void SetJetStreamStandAlone(bool isStandAlone)
{
_state.Lock.EnterWriteLock();
try
{
_state.StandAlone = isStandAlone;
}
finally
{
_state.Lock.ExitWriteLock();
}
}
internal bool IsShuttingDown()
{
_state.Lock.EnterReadLock();
try
{
return _state.ShuttingDown;
}
finally
{
_state.Lock.ExitReadLock();
}
}
internal Exception? DisableJetStream(JsAccount? account)
{
if (account?.Account is not Account a)
return new InvalidOperationException("jetstream not enabled for account");
_state.Lock.EnterWriteLock();
try
{
_state.Accounts.Remove(a.Name);
}
finally
{
_state.Lock.ExitWriteLock();
}
account.Delete();
return null;
}
internal bool WouldExceedLimits(StorageType storageType, int size)
{
var total = storageType == StorageType.MemoryStorage
? Interlocked.Read(ref _state.MemUsed)
: Interlocked.Read(ref _state.StoreUsed);
var max = storageType == StorageType.MemoryStorage
? _state.Config.MaxMemory
: _state.Config.MaxStore;
return total + size > max;
}
internal bool LimitsExceeded(StorageType storageType) => WouldExceedLimits(storageType, 0);
internal static string TierName(int replicas) => $"R{(replicas <= 0 ? 1 : replicas)}";
internal static bool IsSameTier(StreamConfig cfgA, StreamConfig cfgB) =>
cfgA.Replicas == cfgB.Replicas;
internal static Dictionary<string, JetStreamAccountLimits> DiffCheckedLimits(
Dictionary<string, JetStreamAccountLimits> a,
Dictionary<string, JetStreamAccountLimits> b)
{
var diff = new Dictionary<string, JetStreamAccountLimits>(StringComparer.Ordinal);
foreach (var (tier, oldLimit) in a)
{
b.TryGetValue(tier, out var newLimit);
newLimit ??= new JetStreamAccountLimits();
diff[tier] = new JetStreamAccountLimits
{
MaxMemory = newLimit.MaxMemory - oldLimit.MaxMemory,
MaxStore = newLimit.MaxStore - oldLimit.MaxStore,
};
}
foreach (var (tier, newLimit) in b)
{
if (a.ContainsKey(tier))
continue;
diff[tier] = new JetStreamAccountLimits
{
MaxMemory = newLimit.MaxMemory,
MaxStore = newLimit.MaxStore,
};
}
return diff;
}
internal static (ulong Mem, ulong Store) ReservedStorage(
Dictionary<string, object?> streamAssignments,
string tier)
{
ulong mem = 0;
ulong store = 0;
foreach (var assignment in streamAssignments.Values.OfType<StreamAssignmentView>())
{
var cfg = assignment.Config;
if (!string.IsNullOrEmpty(tier) && !string.Equals(tier, TierName(cfg.Replicas), StringComparison.Ordinal))
continue;
if (cfg.MaxBytes <= 0)
continue;
if (cfg.Storage == StorageType.FileStorage)
store += (ulong)cfg.MaxBytes;
else if (cfg.Storage == StorageType.MemoryStorage)
mem += (ulong)cfg.MaxBytes;
}
return (mem, store);
}
internal static bool ShouldSendUsageUpdate(DateTime lastUpdateUtc) =>
DateTime.UtcNow - lastUpdateUtc >= MinUsageUpdateWindow;
}
internal sealed class StreamAssignmentView
{
public required StreamConfig Config { get; init; }
}