changed capacity behavior of buffer collection GetOrCreate() and renamed to EnsureBuffer() to make it explicit

Author: ledjon-behluli

DurableStateMachines.Tests/DurableRingBufferCollectionTests.cs

@@ -1,5 +1,4 @@
 using Microsoft.Extensions.DependencyInjection;
-using System.Diagnostics;
 
 namespace DurableStateMachines.Tests;
 
@@ -8,11 +7,18 @@ public class DurableRingBufferCollectionTests(TestFixture fixture)
 {
     public interface IDurableRingBufferCollectionGrain : IGrainWithStringKey
     {
+        // Workaround methods to survive deactivation
+        Task<Dictionary<string, int>> GetAllCapacities();
+        Task SetAllCapacities(Dictionary<string, int> capacities);
+
+        // Collection methods
         Task<int> GetBuffersCount();
         Task<List<string>> GetKeys();
         Task<bool> ContainsBuffer(string key);
         Task<bool> RemoveBuffer(string key);
         Task ClearAll();
+
+        // Buffer methods
         Task Enqueue(string key, string value);
         Task<TryValue<string?>> TryDequeue(string key);
         Task SetBufferCapacity(string key, int capacity);
@@ -27,6 +33,25 @@ public class DurableRingBufferCollectionGrain(
             : DurableGrain, IDurableRingBufferCollectionGrain
     {
         private const int DefaultCapacity = 1;
+        private readonly Dictionary<string, int> _capacities = [];
+
+        private IDurableRingBuffer<string> GetBuffer(string key)
+        {
+            var capacity = _capacities.GetValueOrDefault(key, DefaultCapacity);
+            return state.EnsureBuffer(key, capacity);
+        }
+
+        public Task<Dictionary<string, int>> GetAllCapacities() => Task.FromResult(new Dictionary<string, int>(_capacities));
+
+        public Task SetAllCapacities(Dictionary<string, int> capacities)
+        {
+            _capacities.Clear();
+            foreach (var (key, value) in capacities)
+            {
+                _capacities[key] = value;
+            }
+            return Task.CompletedTask;
+        }
 
         public Task<int> GetBuffersCount() => Task.FromResult(state.Count);
         public Task<List<string>> GetKeys() => Task.FromResult(state.Keys.ToList());
@@ -37,26 +62,29 @@ public async Task<bool> RemoveBuffer(string key)
             var removed = state.Remove(key);
             if (removed)
             {
+                _capacities.Remove(key);
                 await WriteStateAsync();
             }
             return removed;
         }
 
         public async Task ClearAll()
         {
+            _capacities.Clear();
             state.Clear();
+
             await WriteStateAsync();
         }
 
         public async Task Enqueue(string key, string value)
         {
-            state.GetOrCreate(key, DefaultCapacity).Enqueue(value);
+            GetBuffer(key).Enqueue(value);
             await WriteStateAsync();
         }
 
         public async Task<TryValue<string?>> TryDequeue(string key)
         {
-            var success = state.GetOrCreate(key, DefaultCapacity).TryDequeue(out var item);
+            var success = GetBuffer(key).TryDequeue(out var item);
             if (success)
             {
                 await WriteStateAsync();
@@ -66,19 +94,20 @@ public async Task Enqueue(string key, string value)
 
         public async Task SetBufferCapacity(string key, int capacity)
         {
-            state.GetOrCreate(key, DefaultCapacity).SetCapacity(capacity); 
+            _capacities[key] = capacity;
+            GetBuffer(key); // This will use the now just set '_capacities[key]', as it does EnsureBuffer(key, capacity).
             await WriteStateAsync();
         }
 
         public async Task ClearBuffer(string key)
         {
-            state.GetOrCreate(key, DefaultCapacity).Clear();
+            GetBuffer(key).Clear();
             await WriteStateAsync();
         }
 
-        public Task<int> GetBufferCapacity(string key) => Task.FromResult(state.GetOrCreate(key, DefaultCapacity).Capacity);
-        public Task<int> GetBufferItemsCount(string key) => Task.FromResult(state.GetOrCreate(key, DefaultCapacity).Count);
-        public Task<List<string>> GetAllBufferItems(string key) => Task.FromResult(state.GetOrCreate(key, DefaultCapacity).ToList());
+        public Task<int> GetBufferCapacity(string key) => Task.FromResult(GetBuffer(key).Capacity);
+        public Task<int> GetBufferItemsCount(string key) => Task.FromResult(GetBuffer(key).Count);
+        public Task<List<string>> GetAllBufferItems(string key) => Task.FromResult(GetBuffer(key).ToList());
     }
 
     private IDurableRingBufferCollectionGrain GetGrain(string key) => fixture.Cluster.Client.GetGrain<IDurableRingBufferCollectionGrain>(key);
@@ -119,15 +148,15 @@ public async Task GetOrCreate()
         const string keyA = "BufferA";
 
         // We create a buffer implicitly by setting its capacity.
-        // This calls GetOrCreate(keyA, DefaultCapacity) which creates the buffer and sets capacity to 10.
         await grain.SetBufferCapacity(keyA, 10);
         Assert.Equal(10, await grain.GetBufferCapacity(keyA));
 
-        // Than we enqueue an item. This calls GetOrCreate(keyA, DefaultCapacity).
-        // Because the buffer already exists, the capacity parameter should be ignored.
+        // Than we enqueue an item. This calls GetBuffer(keyA).
+        // Because the grain now remembers the capacity is 10, it will be used,
+        // and the durable state's capacity will not be changed.
         await grain.Enqueue(keyA, "item1");
 
-        // The capacity should have NOT changed to the default.
+        // The capacity should have NOT changed.
         Assert.Equal(10, await grain.GetBufferCapacity(keyA));
         Assert.Equal(1, await grain.GetBufferItemsCount(keyA));
 
@@ -188,7 +217,9 @@ public async Task Persistence()
         await grain.Enqueue(KeyB, "b1");
         await grain.RemoveBuffer(KeyC);
 
+        var capacities1 = await grain.GetAllCapacities();
         await DeactivateGrain(grain);
+        await grain.SetAllCapacities(capacities1);
 
         Assert.Equal(2, await grain.GetBuffersCount());
         Assert.True(await grain.ContainsBuffer(KeyA));
@@ -203,7 +234,9 @@ public async Task Persistence()
         await grain.TryDequeue(KeyA); // Should only remove item "a1" not "BufferA"
         await grain.RemoveBuffer(KeyB);
 
+        var capacities2 = await grain.GetAllCapacities();
         await DeactivateGrain(grain);
+        await grain.SetAllCapacities(capacities2);
 
         Assert.Equal(1, await grain.GetBuffersCount());
         Assert.True(await grain.ContainsBuffer(KeyA));
@@ -228,7 +261,9 @@ public async Task Restore()
             await grain.Enqueue(key, $"item-{i}-2");
         }
 
+        var capacities = await grain.GetAllCapacities();
         await DeactivateGrain(grain); // To trigger a restore from the snapshot
+        await grain.SetAllCapacities(capacities);
 
         Assert.Equal(numBuffers, await grain.GetBuffersCount());
 

DurableStateMachines.Tests/DurableRingBufferTests.cs

@@ -103,6 +103,9 @@ public async Task BasicOperations()
     {
         var grain = GetGrain("basic");
 
+        await Assert.ThrowsAsync<ArgumentOutOfRangeException>(() => grain.SetCapacity(-1));
+        await Assert.ThrowsAsync<ArgumentOutOfRangeException>(() => grain.SetCapacity(0));
+
         await grain.SetCapacity(3);
 
         await grain.Enqueue("one");

DurableStateMachines/DurableRingBufferCollection.cs

@@ -25,12 +25,23 @@ public interface IDurableRingBufferCollection<TKey, TValue> where TKey : notnull
     IReadOnlyCollection<TKey> Keys { get; }
 
     /// <summary>
-    /// Gets or creates the ring buffer associated with the specified key.
+    /// Ensures that a ring buffer associated with the specified key exists and is configured with the given capacity.
     /// </summary>
-    /// <param name="key">The key of the ring buffer to get or create.</param>
-    /// <param name="capacity">The capacity to use when creating a new buffer, if it does not exist.</param>
-    /// <returns>A proxy to the ring buffer that ensures all operations are durable.</returns>
-    IDurableRingBuffer<TValue> GetOrCreate(TKey key, int capacity);
+    /// <param name="key">The key of the ring buffer to ensure.</param>
+    /// <param name="capacity">The desired capacity for the ring buffer.</param>
+    /// <returns>A durable proxy to the ring buffer, which will have the specified capacity after this call.</returns>
+    /// <remarks>
+    /// <para>
+    /// This method provides a convenient way to get a buffer and set its capacity in a single, atomic operation.
+    /// If a buffer for the given <paramref name="key"/> does not exist, it will be created with the specified <paramref name="capacity"/>.
+    /// If the buffer already exists, its capacity will be overwritten with the new value (if its different).
+    /// </para>
+    /// <para><strong>
+    /// Decreasing the capacity on an existing buffer may result in data loss, if the number of items
+    /// currently in the buffer exceeds the new capacity.
+    /// </strong></para>
+    /// </remarks>
+    IDurableRingBuffer<TValue> EnsureBuffer(TKey key, int capacity);
 
     /// <summary>
     /// Determines whether the collection contains a ring buffer with the specified key.
@@ -84,11 +95,12 @@ public DurableRingBufferCollection(
     public int Count => _proxies.Count;
     public IReadOnlyCollection<TKey> Keys => _proxies.Keys;
 
-    public IDurableRingBuffer<TValue> GetOrCreate(TKey key, int capacity)
+    public IDurableRingBuffer<TValue> EnsureBuffer(TKey key, int capacity)
     {
-        var proxy = GetOrCreateProxy(key, out var created);
+        ArgumentOutOfRangeException.ThrowIfNegativeOrZero(capacity, nameof(capacity));
 
-        if (created)
+        var proxy = GetOrCreateProxy(key);
+        if (proxy.Capacity != capacity)
         {
             proxy.SetCapacity(capacity);
         }
@@ -208,7 +220,7 @@ void ApplySnapshot(ref Reader<ReadOnlySequenceInput> reader)
                 var capacity = (int)reader.ReadVarUInt32();
                 var itemCount = (int)reader.ReadVarUInt32();
 
-                var buffer = GetOrCreateProxy(key, out _).Buffer;
+                var buffer = GetOrCreateProxy(key).Buffer;
                 buffer.SetCapacity(capacity);
 
                 for (var j = 0; j < itemCount; j++)
@@ -351,23 +363,18 @@ private void ClearBuffer(TKey key)
     }
 
     private bool ApplyRemove(TKey key) => _proxies.Remove(key, out _);
-    private bool ApplySetBufferCapacity(TKey key, int capacity) => GetOrCreateProxy(key, out _).Buffer.SetCapacity(capacity);
-    private void ApplyEnqueueBufferItem(TKey key, TValue item) => GetOrCreateProxy(key, out _).Buffer.Enqueue(item);
-    private bool ApplyTryDequeueBufferItem(TKey key, out TValue item) => GetOrCreateProxy(key, out _).Buffer.TryDequeue(out item!);
-    private bool ApplyClearBuffer(TKey key) => GetOrCreateProxy(key, out _).Buffer.Clear();
+    private bool ApplySetBufferCapacity(TKey key, int capacity) => GetOrCreateProxy(key).Buffer.SetCapacity(capacity);
+    private void ApplyEnqueueBufferItem(TKey key, TValue item) => GetOrCreateProxy(key).Buffer.Enqueue(item);
+    private bool ApplyTryDequeueBufferItem(TKey key, out TValue item) => GetOrCreateProxy(key).Buffer.TryDequeue(out item!);
+    private bool ApplyClearBuffer(TKey key) => GetOrCreateProxy(key).Buffer.Clear();
     private void ApplyClear() => _proxies.Clear();
 
-    internal RingBufferProxy GetOrCreateProxy(TKey key, out bool created)
+    internal RingBufferProxy GetOrCreateProxy(TKey key)
     {
-        created = false;
-
         if (!_proxies.TryGetValue(key, out var proxy))
         {
             proxy = new RingBufferProxy(key, this);
-
             _proxies[key] = proxy;
-
-            created = true;
         }
 
         return proxy;
@@ -448,7 +455,7 @@ public KeyValuePair<TKey, TValue[]>[] Items
 
             foreach (var key in collection.Keys)
             {
-                var proxy = collection.GetOrCreateProxy(key, out _);
+                var proxy = collection.GetOrCreateProxy(key);
                 result[i++] = new KeyValuePair<TKey, TValue[]>(key, [.. proxy]);
             }