feat: execute post-baseline jetstream parity plan

src/NATS.Server/JetStream/Storage/FileStore.cs

@@ -1,3 +1,4 @@
+using System.Text;
 using System.Text.Json;
 using NATS.Server.JetStream.Models;
 
@@ -5,12 +6,23 @@ namespace NATS.Server.JetStream.Storage;
 
 public sealed class FileStore : IStreamStore, IAsyncDisposable
 {
+    private readonly FileStoreOptions _options;
     private readonly string _dataFilePath;
     private readonly Dictionary<ulong, StoredMessage> _messages = new();
+    private readonly Dictionary<ulong, BlockPointer> _index = new();
     private ulong _last;
+    private int _blockCount;
+    private long _activeBlockBytes;
+    private long _writeOffset;
+
+    public int BlockCount => _messages.Count == 0 ? 0 : Math.Max(_blockCount, 1);
 
     public FileStore(FileStoreOptions options)
     {
+        _options = options;
+        if (_options.BlockSizeBytes <= 0)
+            _options.BlockSizeBytes = 64 * 1024;
+
         Directory.CreateDirectory(options.Directory);
         _dataFilePath = Path.Combine(options.Directory, "messages.jsonl");
         LoadExisting();
@@ -18,6 +30,8 @@ public sealed class FileStore : IStreamStore, IAsyncDisposable
 
     public async ValueTask<ulong> AppendAsync(string subject, ReadOnlyMemory<byte> payload, CancellationToken ct)
     {
+        PruneExpired(DateTime.UtcNow);
+
         _last++;
         var stored = new StoredMessage
         {
@@ -36,6 +50,9 @@ public sealed class FileStore : IStreamStore, IAsyncDisposable
             TimestampUtc = stored.TimestampUtc,
         });
         await File.AppendAllTextAsync(_dataFilePath, line + Environment.NewLine, ct);
+
+        var recordBytes = Encoding.UTF8.GetByteCount(line + Environment.NewLine);
+        TrackBlockForRecord(recordBytes, stored.Sequence);
         return _last;
     }
 
@@ -54,6 +71,14 @@ public sealed class FileStore : IStreamStore, IAsyncDisposable
         return ValueTask.FromResult(match);
     }
 
+    public ValueTask<IReadOnlyList<StoredMessage>> ListAsync(CancellationToken ct)
+    {
+        var messages = _messages.Values
+            .OrderBy(m => m.Sequence)
+            .ToArray();
+        return ValueTask.FromResult<IReadOnlyList<StoredMessage>>(messages);
+    }
+
     public ValueTask<bool> RemoveAsync(ulong sequence, CancellationToken ct)
     {
         var removed = _messages.Remove(sequence);
@@ -65,7 +90,11 @@ public sealed class FileStore : IStreamStore, IAsyncDisposable
     public ValueTask PurgeAsync(CancellationToken ct)
     {
         _messages.Clear();
+        _index.Clear();
         _last = 0;
+        _blockCount = 0;
+        _activeBlockBytes = 0;
+        _writeOffset = 0;
         if (File.Exists(_dataFilePath))
             File.Delete(_dataFilePath);
         return ValueTask.CompletedTask;
@@ -90,7 +119,11 @@ public sealed class FileStore : IStreamStore, IAsyncDisposable
     public ValueTask RestoreSnapshotAsync(ReadOnlyMemory<byte> snapshot, CancellationToken ct)
     {
         _messages.Clear();
+        _index.Clear();
         _last = 0;
+        _blockCount = 0;
+        _activeBlockBytes = 0;
+        _writeOffset = 0;
 
         if (!snapshot.IsEmpty)
         {
@@ -159,29 +192,83 @@ public sealed class FileStore : IStreamStore, IAsyncDisposable
                 Sequence = record.Sequence,
                 Subject = record.Subject ?? string.Empty,
                 Payload = Convert.FromBase64String(record.PayloadBase64 ?? string.Empty),
+                TimestampUtc = record.TimestampUtc,
             };
 
             _messages[message.Sequence] = message;
             if (message.Sequence > _last)
                 _last = message.Sequence;
+
+            var recordBytes = Encoding.UTF8.GetByteCount(line + Environment.NewLine);
+            TrackBlockForRecord(recordBytes, message.Sequence);
         }
+
+        PruneExpired(DateTime.UtcNow);
     }
 
     private void RewriteDataFile()
     {
-        var lines = new List<string>(_messages.Count);
+        Directory.CreateDirectory(Path.GetDirectoryName(_dataFilePath)!);
+        _index.Clear();
+        _blockCount = 0;
+        _activeBlockBytes = 0;
+        _writeOffset = 0;
+
+        using var stream = new FileStream(_dataFilePath, FileMode.Create, FileAccess.Write, FileShare.Read);
+        using var writer = new StreamWriter(stream, Encoding.UTF8);
+
         foreach (var message in _messages.OrderBy(kv => kv.Key).Select(kv => kv.Value))
         {
-            lines.Add(JsonSerializer.Serialize(new FileRecord
+            var line = JsonSerializer.Serialize(new FileRecord
             {
                 Sequence = message.Sequence,
                 Subject = message.Subject,
                 PayloadBase64 = Convert.ToBase64String(message.Payload.ToArray()),
                 TimestampUtc = message.TimestampUtc,
-            }));
+            });
+
+            writer.WriteLine(line);
+            var recordBytes = Encoding.UTF8.GetByteCount(line + Environment.NewLine);
+            TrackBlockForRecord(recordBytes, message.Sequence);
         }
 
-        File.WriteAllLines(_dataFilePath, lines);
+        writer.Flush();
+    }
+
+    private void TrackBlockForRecord(int recordBytes, ulong sequence)
+    {
+        if (_blockCount == 0)
+            _blockCount = 1;
+
+        if (_activeBlockBytes > 0 && _activeBlockBytes + recordBytes > _options.BlockSizeBytes)
+        {
+            _blockCount++;
+            _activeBlockBytes = 0;
+        }
+
+        _index[sequence] = new BlockPointer(_blockCount, _writeOffset);
+        _activeBlockBytes += recordBytes;
+        _writeOffset += recordBytes;
+    }
+
+    private void PruneExpired(DateTime nowUtc)
+    {
+        if (_options.MaxAgeMs <= 0)
+            return;
+
+        var cutoff = nowUtc.AddMilliseconds(-_options.MaxAgeMs);
+        var expired = _messages
+            .Where(kv => kv.Value.TimestampUtc < cutoff)
+            .Select(kv => kv.Key)
+            .ToArray();
+
+        if (expired.Length == 0)
+            return;
+
+        foreach (var sequence in expired)
+            _messages.Remove(sequence);
+
+        RewriteDataFile();
     }
 
     private sealed class FileRecord
@@ -191,4 +278,6 @@ public sealed class FileStore : IStreamStore, IAsyncDisposable
         public string? PayloadBase64 { get; init; }
         public DateTime TimestampUtc { get; init; }
     }
+
+    private readonly record struct BlockPointer(int BlockId, long Offset);
 }