fix(journaling): serialize recovery retries and preserve persistence failures

docs/site/src/content/docs/grains/journaling/runtime-behavior.md

@@ -102,6 +102,34 @@ storage outcome, so the caller reconciles that outcome before retrying the comma
 An initialization failure preserves stored data for diagnosis. Restore the required format/codec registration
 or repair the backing data before creating a fresh manager or retrying activation.
 
+## Custom state lifecycle
+
+Custom <xref:Orleans.Journaling.IStateMachine> implementations share the manager's single logical execution thread.
+Resolve write codecs through <xref:Orleans.Journaling.IJournaledStateManager.GetRequiredCommandCodec*> to use the
+owning manager's configured format, including before recovery of an empty journal. Delegating managers forward
+codec resolution to that owner. Grain-bound managers resolve codecs from the activation's services;
+standalone owners use shared application services.
+
+<xref:Orleans.Journaling.IStateMachine.ValidateWrite*> and <xref:Orleans.Journaling.IStateMachine.ValidateDelete*>
+perform pure validation in the requesting caller's context. An admission rejection leaves the manager healthy.
+Deletion validates all states again in serialized execution, then calls
+<xref:Orleans.Journaling.IStateMachine.OnDeleteStarted*> on every state before awaiting storage deletion.
+Successful deletion resets states before completing callers.
+
+Before each append or snapshot capture, the manager checks <xref:Orleans.Journaling.IStateMachine.IsWritePrepared>
+and awaits <xref:Orleans.Journaling.IStateMachine.PrepareWriteAsync*> for each unprepared state. Preparation
+establishes that state's readiness and retains valid state-owned prerequisites across rechecks. Every await
+is followed by another all-state readiness pass. The final successful pass flows directly into synchronous
+capture in the same work-loop continuation. This also applies to writes which flush only committed entries
+or produce zero bytes. Preparation uses the manager's shutdown token; caller cancellation only ends that caller's wait.
+
+After storage acknowledges captured bytes, <xref:Orleans.Journaling.IStateMachine.OnWriteCompleted*>
+performs durable-completion bookkeeping. A zero-byte write completes without this callback.
+An admitted preparation, validation, capture, or storage failure fences the manager, records the original
+exception, and calls <xref:Orleans.Journaling.IStateMachine.OnFaulted*> on every registered state before
+faulting current and queued waiters. Notification failures are logged while the original failure remains
+the operation's outcome. Idle shutdown completes normally; cancellation during admitted work is terminal.
+
 ## Compaction
 
 Each provider reports when its journal crosses a configured storage threshold. The next `WriteStateAsync`:

src/Orleans.Journaling/IJournaledStateManager.cs

@@ -41,6 +41,20 @@ public interface IJournaledStateManager : IAsyncDisposable
     /// <returns><see langword="true"/> if the state machine is registered; otherwise, <see langword="false"/>.</returns>
     bool TryGetStateMachine(string name, [NotNullWhen(true)] out IStateMachine? stateMachine);
 
+    /// <summary>
+    /// Resolves a command codec for this manager's configured write journal format.
+    /// </summary>
+    /// <typeparam name="TCodec">The command codec service type.</typeparam>
+    /// <returns>The codec registered for this manager's write format.</returns>
+    /// <remarks>
+    /// Codecs are resolved from the owning activation's services for grain-bound managers and from shared
+    /// application services for standalone owners. They are available before recovery, including for an empty journal. Delegating managers forward
+    /// this call to their owning manager. The default implementation throws <see cref="NotSupportedException"/>.
+    /// </remarks>
+    /// <exception cref="NotSupportedException">The manager does not support command codec resolution.</exception>
+    TCodec GetRequiredCommandCodec<TCodec>() where TCodec : notnull
+        => throw new NotSupportedException("This journaled state manager does not support write command codec resolution.");
+
     /// <summary>
     /// Persists pending changes from the registered state machines to the journal.
     /// </summary>

src/Orleans.Journaling/IStateMachine.cs

@@ -14,7 +14,8 @@ namespace Orleans.Journaling;
 /// both apply the mutation locally and emit the corresponding command to the journal).
 /// </item>
 /// <item>
-/// When the application requests a write, the journaled state manager calls
+/// When the application requests a write, the journaled state manager validates the request,
+/// prepares all states, and calls
 /// <see cref="WritePendingEntries"/> (and occasionally <see cref="WriteSnapshot"/>) to materialize
 /// the pending changes, then flushes the journal to durable storage.
 /// </item>
@@ -67,6 +68,72 @@ public interface IStateMachine
     /// </remarks>
     void OnRecoveryCompleted() { }
 
+    /// <summary>
+    /// Gets whether this state has the prerequisites required for synchronous write capture.
+    /// The default is <see langword="true"/>.
+    /// </summary>
+    /// <remarks>
+    /// This synchronous check is pure. The manager evaluates it inside the admitted operation's
+    /// failure boundary, including writes which only flush committed entries or produce zero bytes.
+    /// An unrecoverable state-local error can be reported by throwing, which fences the manager.
+    /// </remarks>
+    bool IsWritePrepared => true;
+
+    /// <summary>
+    /// Acquires state-owned prerequisites for synchronous write capture when <see cref="IsWritePrepared"/> is false.
+    /// The default implementation completes synchronously.
+    /// </summary>
+    /// <param name="cancellationToken">The manager operation and shutdown token.</param>
+    /// <returns>A task which completes when this state is prepared.</returns>
+    /// <remarks>
+    /// Completion must establish this state's readiness. Retain valid prepared resources across readiness
+    /// rechecks. After preparation awaits, the manager rechecks every state and captures synchronously
+    /// in the same continuation as the final ready pass. Caller wait cancellation leaves preparation running.
+    /// Preparation failures fence the manager.
+    /// </remarks>
+    ValueTask PrepareWriteAsync(CancellationToken cancellationToken) => default;
+
+    /// <summary>
+    /// Validates a write request in the public caller's context before it is queued.
+    /// The default implementation accepts the request.
+    /// </summary>
+    /// <remarks>
+    /// Validation is pure. Throwing rejects this request and leaves the manager healthy.
+    /// </remarks>
+    void ValidateWrite() { }
+
+    /// <summary>
+    /// Validates deletion at public request admission and again during serialized execution.
+    /// The default implementation accepts deletion.
+    /// </summary>
+    /// <remarks>
+    /// Validation is pure. An admission failure rejects the request and leaves the manager healthy.
+    /// An execution-time failure fences the manager. All states pass execution-time validation
+    /// before the manager calls <see cref="OnDeleteStarted"/> on any state.
+    /// </remarks>
+    void ValidateDelete() { }
+
+    /// <summary>
+    /// Notifies the state that deletion is starting, after all execution-time validation succeeds
+    /// and before the storage operation begins. The default implementation performs no action.
+    /// </summary>
+    /// <remarks>
+    /// A successful storage deletion is followed by <see cref="Reset"/> before deletion waiters complete.
+    /// </remarks>
+    void OnDeleteStarted() { }
+
+    /// <summary>
+    /// Notifies the state of the manager's first terminal failure, before current and queued operation waiters fault.
+    /// The default implementation performs no action.
+    /// </summary>
+    /// <param name="exception">The original failure recorded by the manager.</param>
+    /// <remarks>
+    /// The manager is already fenced when this callback runs. Every registered state is notified even if
+    /// another notification throws; notification errors are logged and the original failure is preserved.
+    /// Idle shutdown completes normally. Cancellation during admitted preparation or storage work is terminal.
+    /// </remarks>
+    void OnFaulted(Exception exception) { }
+
     /// <summary>
     /// Writes pending state changes to the journal.
     /// </summary>

src/Orleans.Journaling/JournaledStateManager.cs

@@ -100,6 +100,9 @@ internal static IJournalStorage CreateStorage(IJournalStorageProvider storagePro
 
     internal IServiceProvider ServiceProvider => _grainContext is { } context ? context.ActivationServices : _shared.ServiceProvider;
 
+    public TCodec GetRequiredCommandCodec<TCodec>() where TCodec : notnull
+        => JournalFormatServices.GetRequiredCommandCodec<TCodec>(ServiceProvider, _shared.JournalFormatKey);
+
     public bool TryGetStateMachine(string name, [NotNullWhen(true)] out IStateMachine? stateMachine)
     {
         ArgumentException.ThrowIfNullOrEmpty(name);
@@ -253,6 +256,31 @@ private async Task WorkLoop()
                             case AppendJournalWorkItem:
                             case WriteSnapshotWorkItem:
                                 {
+                                    // Keep the final readiness pass and synchronous capture in this continuation.
+                                    bool prepared;
+                                    do
+                                    {
+                                        prepared = true;
+                                        foreach (var (name, state) in _states)
+                                        {
+                                            if (state.IsWritePrepared)
+                                            {
+                                                continue;
+                                            }
+
+                                            await state.PrepareWriteAsync(_shutdownCancellation.Token).ConfigureAwait(true);
+                                            if (!state.IsWritePrepared)
+                                            {
+                                                throw new InvalidOperationException(
+                                                    $"Journaled state '{name}' completed write preparation without becoming prepared.");
+                                            }
+
+                                            prepared = false;
+                                            break;
+                                        }
+                                    }
+                                    while (!prepared);
+
                                     // TODO: decide whether it's best to snapshot or append. Eg, by summing the size of the most recent snapshots and the current journal length.
                                     //       If the current journal length is greater than the snapshot size, then take a snapshot instead of appending more journal entries.
                                     var isSnapshot = workItem is WriteSnapshotWorkItem
@@ -447,6 +475,16 @@ private async Task WorkLoop()
 
                             case DeleteStateWorkItem:
                                 {
+                                    foreach (var state in _states.Values)
+                                    {
+                                        state.ValidateDelete();
+                                    }
+
+                                    foreach (var state in _states.Values)
+                                    {
+                                        state.OnDeleteStarted();
+                                    }
+
                                     // Clear storage.
                                     await DeleteStorageAsync(_shutdownCancellation.Token).ConfigureAwait(true);
 
@@ -548,6 +586,10 @@ private async Task WorkLoop()
                     }
                 }
             }
+            catch (OperationCanceledException) when (_shutdownCancellation.IsCancellationRequested)
+            {
+                return;
+            }
             catch (Exception exception)
             {
                 Fence(exception);
@@ -560,12 +602,30 @@ private void Fence(Exception exception)
     {
         lock (_lock)
         {
+            if (_state is ManagerState.Fenced)
+            {
+                return;
+            }
+
             _state = ManagerState.Fenced;
             _failure = exception;
         }
 
         try
         {
+            // Fencing prevents registration, so callbacks can use the stable registry outside the lock.
+            foreach (var (name, state) in _states)
+            {
+                try
+                {
+                    state.OnFaulted(exception);
+                }
+                catch (Exception notificationException)
+                {
+                    LogErrorNotifyingFaultedState(_shared.Logger, notificationException, name);
+                }
+            }
+
             if (!_shutdownCancellation.IsCancellationRequested)
             {
                 LogErrorProcessingWorkItems(_shared.Logger, exception);
@@ -660,6 +720,11 @@ public async ValueTask DeleteStateAsync(CancellationToken cancellationToken = de
         lock (_lock)
         {
             ThrowIfStateOperationsUnavailable();
+            foreach (var state in _states.Values)
+            {
+                state.ValidateDelete();
+            }
+
             task = EnqueueOrGetPendingWorkItem<DeleteStateWorkItem>(out didEnqueue);
         }
 
@@ -849,6 +914,11 @@ public async ValueTask WriteStateAsync(CancellationToken cancellationToken = def
         lock (_lock)
         {
             ThrowIfStateOperationsUnavailable();
+            foreach (var state in _states.Values)
+            {
+                state.ValidateWrite();
+            }
+
             var isSnapshot = _migrationSnapshotRequired || _storage.IsCompactionRequested;
             operation = isSnapshot ? JournalingInstruments.OperationSnapshot : JournalingInstruments.OperationAppend;
             pendingWrite = isSnapshot
@@ -1307,6 +1377,11 @@ void IStateMachine.WritePendingEntries(JournalStreamWriter writer) { }
         Message = "Error processing work items.")]
     private static partial void LogErrorProcessingWorkItems(ILogger logger, Exception exception);
 
+    [LoggerMessage(
+        Level = LogLevel.Error,
+        Message = "Error notifying journaled state \"{Name}\" of a terminal failure.")]
+    private static partial void LogErrorNotifyingFaultedState(ILogger logger, Exception exception, string name);
+
     [LoggerMessage(
         Level = LogLevel.Information,
         Message = "State \"{Name}\" was not found. I have substituted a placeholder for graceful time-based retirement.")]

src/api/Orleans.Journaling/Orleans.Journaling.cs

@@ -224,6 +224,7 @@ public partial interface IJournaledStateManager : System.IAsyncDisposable
         long PendingWriteByteCount { get; }
 
         System.Threading.Tasks.ValueTask DeleteStateAsync(System.Threading.CancellationToken cancellationToken = default);
+        TCodec GetRequiredCommandCodec<TCodec>();
         System.Threading.Tasks.ValueTask InitializeAsync(System.Threading.CancellationToken cancellationToken = default);
         void RegisterStateMachine(string name, IStateMachine stateMachine);
         System.Threading.Tasks.ValueTask System.IAsyncDisposable.DisposeAsync();
@@ -305,10 +306,17 @@ public partial interface IPreservedJournalEntry
 
     public partial interface IStateMachine
     {
+        bool IsWritePrepared { get; }
+
+        void OnDeleteStarted();
+        void OnFaulted(System.Exception exception);
         void OnRecoveryCompleted();
         void OnWriteCompleted();
+        System.Threading.Tasks.ValueTask PrepareWriteAsync(System.Threading.CancellationToken cancellationToken);
         void ReplayEntry(JournalEntry entry, JournalReplayContext context);
         void Reset(JournalStreamWriter writer);
+        void ValidateDelete();
+        void ValidateWrite();
         void WritePendingEntries(JournalStreamWriter writer);
         void WriteSnapshot(JournalStreamWriter writer);
     }

test/Orleans.Journaling.Tests/JournaledGrainCompositionTests.cs

@@ -336,7 +336,12 @@ private static CompositionSiloBuilder CreateBuilder()
         builder.Services.AddSingleton(TimeProvider.System);
         builder.Services.AddKeyedSingleton<TimeProvider>(KeyedService.AnyKey, static (services, _) => services.GetRequiredService<TimeProvider>());
         builder.AddVolatileJournalStorage().UseJsonJournalFormat(JournalingTestsJsonContext.Default);
-        builder.Services.AddStateMachine<IStateMachine, IStateMachine>(static (_, _) => Substitute.For<IStateMachine>());
+        builder.Services.AddStateMachine<IStateMachine, IStateMachine>(static (_, _) =>
+        {
+            var state = Substitute.For<IStateMachine>();
+            state.IsWritePrepared.Returns(true);
+            return state;
+        });
         return builder;
     }