refactor: refine InMemoryStream

Author: WeihanLi

src/WeihanLi.Common/Helpers/InMemoryStream.cs

@@ -2,95 +2,138 @@
 // Licensed under the Apache license.
 
 using System.Collections.Concurrent;
-using System.Runtime.CompilerServices;
-using WeihanLi.Common.Abstractions;
+using WeihanLi.Common.Models;
+using WeihanLi.Extensions;
 
 namespace WeihanLi.Common.Helpers;
 
-public interface IStream
+public class StreamMessage<T>
 {
-    Task AddMessageAsync(string streamName, StreamMessage message);
-    Task AcknowledgeMessageAsync(string streamName, string messageId);
-    Task<int> CountAsync(string streamName, string? start = null, string? end = null);
-    Task TrimAsync(string streamName, int maxSize);
-    Task TrimAsync(string streamName, string maxId);
-    IAsyncEnumerable<StreamMessage> ReadMessagesAsync(string streamName, int count, string? start = null, string? end = null, CancellationToken cancellationToken = default);
+    public required T Id { get; init; }
+    public DateTimeOffset Timestamp { get; set; }
+    public Dictionary<string, string> Fields { get => field; set => field = Guard.NotNull(value); } = [];
+    public Dictionary<string, object> Properties { get => field; set => field = Guard.NotNull(value); } = [];
 }
 
-public class InMemoryStream : IStream
+public class StreamInfo<T>
 {
-    private readonly ConcurrentDictionary<string, ConcurrentQueue<StreamMessage>> _streams = new();
+    public required T MinId { get; set; }
+    public DateTimeOffset MinTimestamp { get; set; }
+    public required T MaxId { get; set; }
+    public DateTimeOffset MaxTimestamp { get; set; }
+    public int Count { get; set; }
+}
 
-    public Task AddMessageAsync(string streamName, StreamMessage message)
-    {
-        var stream = _streams.GetOrAdd(streamName, _ => new ConcurrentQueue<StreamMessage>());
-        stream.Enqueue(message);
-        return Task.CompletedTask;
-    }
+public class StreamGroupInfo<T>
+{
+    public required string GroupName { get; set; }
+    public required T Offset { get; set; }
+}
+
+public interface IStream<T>
+{
+    string StreamName { get; }
+
+    Task AddAsync(T id, Dictionary<string, string> fields, DateTimeOffset? timestamp = null, Dictionary<string, object>? properties = null, CancellationToken cancellationToken = default);
+    IAsyncEnumerable<StreamMessage<T>> FetchAsync(T lastId, int count, Ordering order = default, CancellationToken cancellationToken = default);
 
-    public Task AcknowledgeMessageAsync(string streamName, string messageId) => Task.CompletedTask;
+    Task<int> CountAsync(T? min = default, T? max = default, RangeInclusion inclusion = default, CancellationToken cancellationToken = default);
 
-    public Task<int> CountAsync(string streamName, string? start = null, string? end = null)
+    Task<StreamInfo<T>> InfoAsync(CancellationToken cancellationToken = default);
+    Task<IReadOnlyCollection<StreamGroupInfo<T>>> GroupsAsync(CancellationToken cancellationToken = default);
+    Task<StreamGroupInfo<T>?> GroupInfoAsync(string groupName, CancellationToken cancellationToken = default);
+    Task AddGroupAsync(string groupName, T offset, CancellationToken cancellationToken = default);
+    Task<bool> RemoveGroupAsync(string groupName, CancellationToken cancellationToken = default);
+    Task AckAsync(string groupName, T id, CancellationToken cancellationToken = default);
+}
+
+public sealed class InMemoryStream<T>(string name, IComparer<T>? comparer = null) : IStream<T>
+{
+    private readonly PriorityQueue<StreamMessage<T>, T> _messages = new(comparer);
+    private readonly ConcurrentDictionary<string, StreamGroupInfo<T>> _groups = new();
+
+    public string StreamName => name;
+
+    public Task AckAsync(string groupName, T id, CancellationToken cancellationToken = default)
     {
-        if (_streams.TryGetValue(streamName, out var stream))
+        if (!_groups.TryGetValue(groupName, out var groupInfo))
         {
-            return Task.FromResult(stream.Count(message =>
-                (start == null || string.Compare(message.Id, start) >= 0) &&
-                (end == null || string.Compare(message.Id, end) <= 0)));
+            throw new InvalidOperationException($"Group [{groupName}] not exists");
         }
 
-        return Task.FromResult(0);
+        groupInfo.Offset = id;
+        return Task.CompletedTask;
     }
 
-    public Task TrimAsync(string streamName, int maxSize)
+    public Task AddAsync(T id, Dictionary<string, string> fields, DateTimeOffset? timestamp = null, Dictionary<string, object>? properties = null, CancellationToken cancellationToken = default)
     {
-        if (_streams.TryGetValue(streamName, out var stream))
+        var message = new StreamMessage<T>
         {
-            while (stream.Count > maxSize && stream.TryDequeue(out _))
-            {
-                // Remove messages until the stream size is within the limit
-            }
-        }
+            Id = id,
+            Fields = fields,
+            Timestamp = timestamp ?? DateTimeOffset.Now,
+            Properties = properties ?? new()
+        };
+        _messages.Enqueue(message, id);
+
         return Task.CompletedTask;
     }
 
-    public Task TrimAsync(string streamName, string maxId)
+    public Task AddGroupAsync(string groupName, T offset, CancellationToken cancellationToken = default)
     {
-        if (_streams.TryGetValue(streamName, out var stream))
+        if (_groups.ContainsKey(groupName))
         {
-            var messages = stream.ToArray();
-            stream.Clear();
-            foreach (var message in messages)
-            {
-                if (string.Compare(message.Id, maxId) > 0)
-                {
-                    stream.Enqueue(message);
-                }
-            }
+            throw new InvalidOperationException($"Group [{groupName}] not exists");
         }
+
+        _groups[groupName] = new StreamGroupInfo<T>()
+        {
+            GroupName = groupName,
+            Offset = offset
+        };
         return Task.CompletedTask;
     }
 
-    public async IAsyncEnumerable<StreamMessage> ReadMessagesAsync(string streamName, int count, string? start = null, string? end = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    public Task<int> CountAsync(T? min = default, T? max = default, RangeInclusion inclusion = default, CancellationToken cancellationToken = default)
     {
-        while (!cancellationToken.IsCancellationRequested)
+        // TODO: support min/max filter
+        return _messages.Count.WrapTask();
+    }
+
+    public IAsyncEnumerable<StreamMessage<T>> FetchAsync(T lastId, int count, Ordering order = Ordering.Ascending, CancellationToken cancellationToken = default)
+    {
+        throw new NotImplementedException();
+    }
+
+    public Task<StreamGroupInfo<T>?> GroupInfoAsync(string groupName, CancellationToken cancellationToken = default)
+    {
+        if (_groups.TryGetValue(groupName, out var groupInfo))
         {
-            if (_streams.TryGetValue(streamName, out var stream))
-            {
-                while (stream.TryDequeue(out var message))
-                {
-                    yield return message;
-                }
-            }
-            await Task.Delay(100, cancellationToken);
+            return Task.FromResult<StreamGroupInfo<T>?>(groupInfo);
         }
+
+        return Task.FromResult<StreamGroupInfo<T>?>(null);
     }
-}
 
-public class StreamMessage : IProperties
-{
-    public required string Id { get; set; }
-    public string Data { get; set; }
-    public DateTimeOffset Timestamp { get; set; }
-    public IDictionary<string, object?> Properties { get; set; } = new Dictionary<string, object?>();
+    public Task<IReadOnlyCollection<StreamGroupInfo<T>>> GroupsAsync(CancellationToken cancellationToken = default)
+    {
+        return _groups.Values.AsReadOnly().WrapTask();
+    }
+
+    public Task<StreamInfo<T>> InfoAsync(CancellationToken cancellationToken = default)
+    {
+        var streamInfo = new StreamInfo<T>
+        {
+            // TODO: update min/max from messages
+            MinId = default,
+            MaxId = default,
+            Count = _messages.Count
+        };
+        return streamInfo.WrapTask();
+    }
+
+    public Task<bool> RemoveGroupAsync(string groupName, CancellationToken cancellationToken = default)
+    {
+        return _groups.TryRemove(groupName, out _).WrapTask();
+    }
 }

src/WeihanLi.Common/Models/Ordering.cs

@@ -0,0 +1,18 @@
+// Copyright (c) Weihan Li. All rights reserved.
+// Licensed under the Apache license.
+
+namespace WeihanLi.Common.Models;
+
+public enum Ordering
+{
+    Ascending = 0,
+    Descending = 1
+}
+
+[Flags]
+public enum RangeInclusion
+{
+    None = 0,
+    IncludeLowerBound = 1,
+    IncludeUpperBound = 2
+}