natsdotnet/src/NATS.Server/JetStream/Storage/MemStore.cs

using System.Text.Json;
using NATS.Server.JetStream.Models;

namespace NATS.Server.JetStream.Storage;

public sealed class MemStore : IStreamStore
{
    private sealed class SnapshotRecord
    {
        public ulong Sequence { get; init; }
        public string Subject { get; init; } = string.Empty;
        public string PayloadBase64 { get; init; } = string.Empty;
        public DateTime TimestampUtc { get; init; }
    }

    private readonly object _gate = new();
    private ulong _last;
    private readonly Dictionary<ulong, StoredMessage> _messages = new();

    public ValueTask<ulong> AppendAsync(string subject, ReadOnlyMemory<byte> payload, CancellationToken ct)
    {
        lock (_gate)
        {
            _last++;
            _messages[_last] = new StoredMessage
            {
                Sequence = _last,
                Subject = subject,
                Payload = payload,
                TimestampUtc = DateTime.UtcNow,
            };
            return ValueTask.FromResult(_last);
        }
    }

    public ValueTask<StoredMessage?> LoadAsync(ulong sequence, CancellationToken ct)
    {
        lock (_gate)
        {
            _messages.TryGetValue(sequence, out var msg);
            return ValueTask.FromResult(msg);
        }
    }

    public ValueTask<StoredMessage?> LoadLastBySubjectAsync(string subject, CancellationToken ct)
    {
        lock (_gate)
        {
            var match = _messages.Values
                .Where(m => string.Equals(m.Subject, subject, StringComparison.Ordinal))
                .OrderByDescending(m => m.Sequence)
                .FirstOrDefault();
            return ValueTask.FromResult(match);
        }
    }

    public ValueTask<IReadOnlyList<StoredMessage>> ListAsync(CancellationToken ct)
    {
        lock (_gate)
        {
            var messages = _messages.Values
                .OrderBy(m => m.Sequence)
                .ToArray();
            return ValueTask.FromResult<IReadOnlyList<StoredMessage>>(messages);
        }
    }

    public ValueTask<bool> RemoveAsync(ulong sequence, CancellationToken ct)
    {
        lock (_gate)
        {
            return ValueTask.FromResult(_messages.Remove(sequence));
        }
    }

    public ValueTask PurgeAsync(CancellationToken ct)
    {
        lock (_gate)
        {
            _messages.Clear();
            _last = 0;
            return ValueTask.CompletedTask;
        }
    }

    public ValueTask<byte[]> CreateSnapshotAsync(CancellationToken ct)
    {
        lock (_gate)
        {
            var snapshot = _messages
                .Values
                .OrderBy(x => x.Sequence)
                .Select(x => new SnapshotRecord
                {
                    Sequence = x.Sequence,
                    Subject = x.Subject,
                    PayloadBase64 = Convert.ToBase64String(x.Payload.ToArray()),
                    TimestampUtc = x.TimestampUtc,
                })
                .ToArray();
            return ValueTask.FromResult(JsonSerializer.SerializeToUtf8Bytes(snapshot));
        }
    }

    public ValueTask RestoreSnapshotAsync(ReadOnlyMemory<byte> snapshot, CancellationToken ct)
    {
        lock (_gate)
        {
            _messages.Clear();
            _last = 0;

            if (!snapshot.IsEmpty)
            {
                var records = JsonSerializer.Deserialize<SnapshotRecord[]>(snapshot.Span);
                if (records != null)
                {
                    foreach (var record in records)
                    {
                        _messages[record.Sequence] = new StoredMessage
                        {
                            Sequence = record.Sequence,
                            Subject = record.Subject,
                            Payload = Convert.FromBase64String(record.PayloadBase64),
                            TimestampUtc = record.TimestampUtc,
                        };
                        _last = Math.Max(_last, record.Sequence);
                    }
                }
            }

            return ValueTask.CompletedTask;
        }
    }

    public ValueTask<ApiStreamState> GetStateAsync(CancellationToken ct)
    {
        lock (_gate)
        {
            return ValueTask.FromResult(new ApiStreamState
            {
                Messages = (ulong)_messages.Count,
                FirstSeq = _messages.Count == 0 ? 0UL : _messages.Keys.Min(),
                LastSeq = _last,
                Bytes = (ulong)_messages.Values.Sum(m => m.Payload.Length),
            });
        }
    }

    public void TrimToMaxMessages(ulong maxMessages)
    {
        lock (_gate)
        {
            while ((ulong)_messages.Count > maxMessages)
            {
                var first = _messages.Keys.Min();
                _messages.Remove(first);
            }
        }
    }
}