[Journaling] Auto-retire & purge state machines

src/Orleans.Journaling/DurableNothing.cs

@@ -1,24 +1,22 @@
-using System.Buffers;
+using System.Buffers;
+using System.Diagnostics;
 using Microsoft.Extensions.DependencyInjection;
 
 namespace Orleans.Journaling;
 
 /// <summary>
 /// A durable object which does nothing, used for retiring other durable types.
 /// </summary>
-public interface IDurableNothing
+[DebuggerDisplay("DurableNothing")]
+internal sealed class DurableNothing : IDurableStateMachine
 {
-}
+    public string StateMachineKey { get; }
 
-/// <summary>
-/// A durable object which does nothing, used for retiring other durable types.
-/// </summary>
-internal sealed class DurableNothing : IDurableNothing, IDurableStateMachine
-{
     public DurableNothing([ServiceKey] string key, IStateMachineManager manager)
     {
         ArgumentNullException.ThrowIfNullOrEmpty(key);
         manager.RegisterStateMachine(key, this);
+        StateMachineKey = key;
     }
 
     void IDurableStateMachine.Reset(IStateMachineLogWriter storage) { }

src/Orleans.Journaling/HostingExtensions.cs

@@ -2,6 +2,7 @@
 using Microsoft.Extensions.DependencyInjection.Extensions;
 
 namespace Orleans.Journaling;
+
 public static class HostingExtensions
 {
     public static ISiloBuilder AddStateMachineStorage(this ISiloBuilder builder)
@@ -15,7 +16,6 @@ public static ISiloBuilder AddStateMachineStorage(this ISiloBuilder builder)
         builder.Services.TryAddKeyedScoped(typeof(IDurableValue<>), KeyedService.AnyKey, typeof(DurableValue<>));
         builder.Services.TryAddKeyedScoped(typeof(IPersistentState<>), KeyedService.AnyKey, typeof(DurableState<>));
         builder.Services.TryAddKeyedScoped(typeof(IDurableTaskCompletionSource<>), KeyedService.AnyKey, typeof(DurableTaskCompletionSource<>));
-        builder.Services.TryAddKeyedScoped(typeof(IDurableNothing), KeyedService.AnyKey, typeof(DurableNothing));
         return builder;
     }
 }

src/Orleans.Journaling/StateMachineManager.cs

@@ -25,6 +25,7 @@ internal sealed partial class StateMachineManager : IStateMachineManager, ILifec
     private readonly Task _workLoop;
     private ManagerState _state;
     private Task? _pendingWrite;
+    private bool _hasStateMachineToRetire;
     private ulong _nextStateMachineId = MinApplicationStateMachineId;
     private LogExtentBuilder? _currentLogSegment;
 
@@ -125,21 +126,38 @@ private async Task WorkLoop()
                             //       If the current log length is greater than the snapshot size, then take a snapshot instead of appending more log entries.
                             var isSnapshot = workItem.Type is WorkItemType.WriteSnapshot;
                             LogExtentBuilder? logSegment;
+
                             lock (_lock)
                             {
-                                if (isSnapshot && _currentLogSegment is { } existingSegment)
+                                if (isSnapshot)
                                 {
                                     // If there are pending writes, reset them since they will be captured by the snapshot instead.
                                     // If we did not do this, the log would begin with some writes which would be followed by a snapshot which also included those writes.
-                                    existingSegment.Reset();
-                                }
-                                else
-                                {
-                                    _currentLogSegment ??= new();
+                                    _currentLogSegment?.Reset();
+
+                                    if (_hasStateMachineToRetire) // We use this flag because the majooority of times, there wont be any state machine to retire.
+                                    {
+                                        // Since this is a snapshot, we use the opportunity to purge retired state machines.
+                                        foreach (var (id, machine) in _stateMachinesMap)
+                                        {
+                                            if (machine is DurableNothing retired)
+                                            {
+                                                var name = retired.StateMachineKey;
+
+                                                LogPurgingRetiredStateMachine(_logger, name, id);
+
+                                                _stateMachinesMap.Remove(id);
+                                                _stateMachineIds.Remove(name); // This will take effect when the snapshot is persisted below.
+                                            }
+                                        }
+                                    }
                                 }
 
+                                _currentLogSegment ??= new();
+
                                 // The map of state machine ids is itself stored as a durable state machine with the id 0.
                                 // This must be stored first, since it includes the identities of all other state machines, which are needed when replaying the log.
+                                // If we removed retired machines, this snapshot will persist that change.
                                 AppendUpdatesOrSnapshotStateMachine(_currentLogSegment, isSnapshot, 0, _stateMachineIds);
 
                                 foreach (var (id, stateMachine) in _stateMachinesMap)
@@ -366,7 +384,15 @@ private void OnSetStateMachineId(string name, ulong id)
             }
             else
             {
-                throw new InvalidOperationException($"State machine \"{name}\" (id: {id}) has not been registered on this state machine manager.");
+                // When a state machine is not found, we substitute it with a DurableNothing.
+                // It will be purged on the next snapshot.
+
+                LogRetiredStateMachineDetected(_logger, name, id);
+
+                _stateMachinesMap[id] = new DurableNothing(name, this);
+                _hasStateMachineToRetire = true;
+
+                // No need to call Reset, as DurableNothing does nothing.
             }
         }
     }
@@ -452,4 +478,14 @@ private sealed class StateMachineManagerState(
         Level = LogLevel.Error,
         Message = "Error processing work items.")]
     private static partial void LogErrorProcessingWorkItems(ILogger logger, Exception exception);
+
+    [LoggerMessage(
+        Level = LogLevel.Information,
+        Message = "State machine \"{Name}\" (id: {Id}) not found. Substituting a placeholder for graceful retirement.")]
+    private static partial void LogRetiredStateMachineDetected(ILogger logger, string name, ulong id);
+
+    [LoggerMessage(
+        Level = LogLevel.Information,
+        Message = "Purging retired state machine \"{Name}\" (id: {Id}).")]
+    private static partial void LogPurgingRetiredStateMachine(ILogger logger, string name, ulong id);
 }

test/Orleans.Journaling.Tests/StateMachineManagerTests.cs

@@ -256,4 +256,65 @@ public async Task StateMachineManager_LargeStateRecovery_Test()
             Assert.Equal($"Value {i}", recoveredDict[i]);
         }
     }
+
+    /// <summary>
+    /// Tests that a "retired" state machine (one that is no longer registered)
+    /// has its data (and itself) purged when the storage triggers a compaction.
+    /// </summary>
+    [Fact]
+    public async Task StateMachineManager_RetiredStateMachine_IsPurgedOnCompaction()
+    {
+        const string DictToKeepKey = "dictToKeep";
+        const string DictToRetireKey = "dictToRetire";
+
+        var sut1 = CreateTestSystem();
+
+        // We beging with 2 dictionaries, one of which we will retire by means of not registering it in the manger.
+        // This would be in the real-world developers removing it from the grain's ctor as a dependecy. 
+        var dictToKeep = new DurableDictionary<string, int>(DictToKeepKey, sut1.Manager, CodecProvider.GetCodec<string>(), CodecProvider.GetCodec<int>(), SessionPool);
+        var dictToRetire = new DurableDictionary<string, int>(DictToRetireKey, sut1.Manager, CodecProvider.GetCodec<string>(), CodecProvider.GetCodec<int>(), SessionPool);
+
+        await sut1.Lifecycle.OnStart();
+
+        dictToKeep.Add("a", 1);
+        dictToRetire.Add("b", 2);
+
+        await sut1.Manager.WriteStateAsync(CancellationToken.None);
+
+        var sut2 = CreateTestSystem(sut1.Storage);
+
+        // This time, we only register the dictionary we want to keep, this marks dictToRetire as retired.
+        var recoveredDictToKeep = new DurableDictionary<string, int>(DictToKeepKey, sut2.Manager, CodecProvider.GetCodec<string>(), CodecProvider.GetCodec<int>(), SessionPool);
+
+        await sut2.Lifecycle.OnStart();
+
+        // The manager should have recovered the state for dictToKeep, and created a DurableNothing placeholder for dictToRetire.
+        Assert.Equal(1, recoveredDictToKeep["a"]);
+
+        // Now, we trigger the compaction logic in VolatileStateMachineStorage by writing more than 10 times.
+        for (var i = 0; i < 11; i++)
+        {
+            recoveredDictToKeep["a"] = i;
+            await sut2.Manager.WriteStateAsync(CancellationToken.None);
+        }
+
+        // At this point, the manager has performed a snapshot, so it should have purged the dictToRetire data.
+        var sut3 = CreateTestSystem(sut2.Storage);
+
+        // So by registering both dictionaries again, we should see what state remains after the snapshot.
+        var finalDictToKeep = new DurableDictionary<string, int>(DictToKeepKey, sut3.Manager, CodecProvider.GetCodec<string>(), CodecProvider.GetCodec<int>(), SessionPool);
+        var finalDictToRetire = new DurableDictionary<string, int>(DictToRetireKey, sut3.Manager, CodecProvider.GetCodec<string>(), CodecProvider.GetCodec<int>(), SessionPool);
+
+        await sut3.Lifecycle.OnStart();
+
+        // The dictionary we kept should have the last value written to it.
+        Assert.Equal(10, finalDictToKeep["a"]); // The last value from the i=0..10 loop.
+
+        // The retired dictionary should now be empty because its state was purged during the compaction.
+        // Note that this is a new version of dictToRetire, since that is removed as a state machine, idea here
+        // is that if we can register a new dictToRetire with the same key, it means that the machine itself has been removed
+        // but also the data, otherwise previous machine would have had at least one entry i.e. ["b", 2]. The removal of the machine
+        // itself has also the nice benefit of being able to reuse old machine names.
+        Assert.Empty(finalDictToRetire);
+    }
 }