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

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

@@ -62,13 +62,11 @@ Concurrent calls made while the same kind of write is queued can share that queu
 ## Safe-to-commit staging
 
 All interleaved callers share the manager's pending journal. Prepare fallible work, external acknowledgements,
-and proposed output in operation-local data. After establishing that an outcome is safe to commit, apply its
-mutations to durable state and initiate a write. Coordinate that transition with other interleaved operations
-which can affect the same decision. Any caller's write can include staged mutations from other calls.
-
-If an application error occurs after staging and makes those mutations unsafe to commit, end the activation's
-use of the manager and request deactivation. In-flight methods can retain local decisions and references
-across awaits; a fresh activation reconstructs both application and durable state together.
+and proposed output in operation-local data. After the final preparation await, check the relevant
+preconditions and apply the complete safe-to-commit update synchronously, then request an ordinary write.
+Orleans executes that synchronous block on a single activation thread. Another grain turn can run when
+the operation awaits, so keep shared state safe to commit at each await. Any caller's write can include
+staged mutations from other calls.
 
 ## Consistency and competing writers
 
@@ -102,6 +100,36 @@ 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.
 
+Owner shutdown which cancels initial recovery cancels all initialization waiters and leaves the manager
+stopped. Disposal waits for the owned read to finish before releasing journal resources. Cancelling an
+individual initialization caller's token ends only its wait; owned recovery continues for other callers.
+
+## Custom state lifecycle
+
+Custom <xref:Orleans.Journaling.IStateMachine> implementations share the manager's single logical execution thread.
+Supply command codecs as constructor dependencies. The registration factory selects codecs keyed by the
+same write-format key used to configure the journal owner. Activation-owned state factories resolve those
+dependencies from the activation's services; standalone callers supply codecs with the appropriate lifetime.
+The codec is available when the state is constructed, including for an empty journal.
+During replay, <xref:Orleans.Journaling.JournalReplayContext.GetRequiredCommandCodec*> selects the codec
+for each entry's stored format.
+
+States synchronously encode their pending changes through <xref:Orleans.Journaling.IStateMachine.WritePendingEntries*>
+or their current contents through <xref:Orleans.Journaling.IStateMachine.WriteSnapshot*>.
+After storage acknowledges captured bytes, <xref:Orleans.Journaling.IStateMachine.OnWriteCompleted*>
+performs durable-completion bookkeeping. A zero-byte write completes without this callback.
+
+The journal owner keeps feature operations quiescent through deletion's storage and reset outcome,
+including when a caller cancels its wait. Successful deletion calls <xref:Orleans.Journaling.IStateMachine.Reset*>
+before completing deletion waiters.
+
+The manager records the first capture or storage failure, fences further persistence, faults current
+and queued manager waiters, and requests grain deactivation. Features observe their write failures and
+complete their own operation waiters and resource cleanup through their operation and lifecycle ownership.
+Standalone callers own that cleanup explicitly. A previously captured write retains its actual storage
+outcome and acknowledgement bookkeeping. Owner-canceled initial recovery and idle shutdown complete
+through normal shutdown; admitted write/delete storage cancellation remains terminal.
+
 ## Compaction
 
 Each provider reports when its journal crosses a configured storage threshold. The next `WriteStateAsync`:

src/Orleans.Journaling/IJournaledStateManager.cs

@@ -19,7 +19,9 @@ public interface IJournaledStateManager : IAsyncDisposable
     /// Initializes the state manager by replaying its journal.
     /// </summary>
     /// <remarks>
-    /// A failed initialization permanently fences this instance. Recover by creating a new manager and new state instances.
+    /// A recovery failure permanently fences this instance. Recover by creating a new manager and new state instances.
+    /// Owner shutdown which cancels recovery cancels initialization and leaves this instance stopped.
+    /// Cancelling the caller's token ends only that caller's wait while owned recovery continues.
     /// </remarks>
     /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>A <see cref="ValueTask"/> which represents the operation.</returns>
@@ -58,7 +60,8 @@ public interface IJournaledStateManager : IAsyncDisposable
     /// Resets this instance, removing any persistent state.
     /// </summary>
     /// <remarks>
-    /// Quiesce other operations before deleting state: deletion resets every registered state machine.
+    /// The caller keeps other operations quiescent through completion: deletion resets every registered state machine.
+    /// Cancellation ends the caller's wait; an already queued deletion continues to its storage and reset outcome.
     /// A failed deletion permanently fences the manager and requests deactivation of its owning grain.
     /// </remarks>
     /// <param name="cancellationToken">The cancellation token.</param>

src/Orleans.Journaling/JournaledStateManager.cs

@@ -191,6 +191,10 @@ private async Task WorkLoop()
         {
             await RecoverAsync(_shutdownCancellation.Token).ConfigureAwait(true);
         }
+        catch (OperationCanceledException) when (_shutdownCancellation.IsCancellationRequested)
+        {
+            return;
+        }
         catch (Exception exception)
         {
             Fence(exception);
@@ -548,6 +552,10 @@ private async Task WorkLoop()
                     }
                 }
             }
+            catch (OperationCanceledException) when (_shutdownCancellation.IsCancellationRequested)
+            {
+                return;
+            }
             catch (Exception exception)
             {
                 Fence(exception);
@@ -560,6 +568,11 @@ private void Fence(Exception exception)
     {
         lock (_lock)
         {
+            if (_state is ManagerState.Fenced)
+            {
+                return;
+            }
+
             _state = ManagerState.Fenced;
             _failure = exception;
         }

test/Orleans.Journaling.Tests/KeyedJournalingRegistrationTests.cs

@@ -382,6 +382,126 @@ public void DurableService_ResolvesCommandCodecFromJournalFormatKey()
             scope.ServiceProvider.GetRequiredService<IGrainContext>().ObservableLifecycle).Subscriptions);
     }
 
+    [Fact]
+    public async Task StateConstruction_InjectsActivationScopedCodecBeforeRecovery()
+    {
+        var builder = CreateNamedProviderBuilder();
+        builder.AddVolatileJournalStorage();
+        builder.Services.Configure<JournaledStateManagerOptions>(options =>
+            options.JournalFormatKey = OrleansBinaryJournalFormat.JournalFormatKey);
+        builder.Services.AddScoped<IGrainContext>(services =>
+        {
+            var context = Substitute.For<IGrainContext>();
+            context.GrainId.Returns(GrainId.Create("codec-scope", Guid.NewGuid().ToString("N")));
+            context.ActivationServices.Returns(services);
+            context.ObservableLifecycle.Returns(new CompositionTestLifecycle());
+            return context;
+        });
+        builder.Services.AddKeyedScoped<IDurableValueCommandCodec<int>>(OrleansBinaryJournalFormat.JournalFormatKey,
+            static (services, _) => new OrleansBinaryDurableValueCommandCodec<int>(
+                services.GetRequiredService<ICodecProvider>().GetCodec<int>(),
+                services.GetRequiredService<SerializerSessionPool>()));
+        builder.Services.AddStateMachine<CodecState, CodecState>(static (services, _) =>
+            new CodecState(services.GetRequiredKeyedService<IDurableValueCommandCodec<int>>(OrleansBinaryJournalFormat.JournalFormatKey)));
+        await using var services = builder.Services.BuildServiceProvider(validateScopes: true);
+        await using var first = services.CreateAsyncScope();
+        await using var second = services.CreateAsyncScope();
+        var owner = first.ServiceProvider.GetRequiredService<IJournaledStateManager>();
+        var manager = first.ServiceProvider.GetRequiredService<IDurableStateManager>();
+        Assert.Same(owner, manager);
+        var state = manager.GetOrAddState<CodecState>("value");
+        var codec = state.Codec;
+        Assert.Same(first.ServiceProvider.GetRequiredKeyedService<IDurableValueCommandCodec<int>>(OrleansBinaryJournalFormat.JournalFormatKey), codec);
+        Assert.NotSame(codec, second.ServiceProvider.GetRequiredService<IDurableStateManager>()
+            .GetOrAddState<CodecState>("value").Codec);
+        Assert.Throws<InvalidOperationException>(() =>
+            services.GetRequiredKeyedService<IDurableValueCommandCodec<int>>(OrleansBinaryJournalFormat.JournalFormatKey));
+        Assert.Same(state, manager.GetOrAddState<CodecState>("value"));
+        Assert.Same(state, first.ServiceProvider.GetRequiredKeyedService<CodecState>("value"));
+        Assert.True(owner.TryGetStateMachine("value", out var registered));
+        Assert.Same(state, registered);
+        var lifecycle = Assert.IsType<CompositionTestLifecycle>(first.ServiceProvider.GetRequiredService<IGrainContext>().ObservableLifecycle);
+        Assert.Equal(1, lifecycle.Subscriptions);
+        await lifecycle.OnStart(TestContext.Current.CancellationToken);
+        Assert.Equal(0, state.Value);
+        state.Value = 42;
+        await manager.WriteStateAsync(TestContext.Current.CancellationToken);
+        Assert.Equal(0, owner.PendingWriteByteCount);
+        await lifecycle.OnStop(TestContext.Current.CancellationToken);
+    }
+
+    [Fact]
+    public async Task StateConstruction_InjectsNamedFormatCodecOnEmptyJournal()
+    {
+        var builder = CreateNamedProviderBuilder();
+        builder.AddVolatileJournalStorage();
+        var customStorage = new VolatileJournalStorage(CustomFormatKey);
+        builder.Services.AddKeyedSingleton<IJournalFormat>(CustomFormatKey, static (services, _) =>
+            new NamedBinaryJournalFormat(services.GetRequiredService<OrleansBinaryJournalFormat>()));
+        builder.Services.AddKeyedSingleton(typeof(IDurableDictionaryCommandCodec<,>), CustomFormatKey,
+            typeof(OrleansBinaryDurableDictionaryCommandCodec<,>));
+        builder.Services.AddKeyedScoped<IDurableValueCommandCodec<int>>(CustomFormatKey, static (services, _) =>
+            new OrleansBinaryDurableValueCommandCodec<int>(
+                services.GetRequiredService<ICodecProvider>().GetCodec<int>(),
+                services.GetRequiredService<SerializerSessionPool>()));
+        builder.Services.AddKeyedSingleton<IJournaledStateManagerFactory>("custom", (services, _) =>
+            new JournaledStateManagerFactory(
+                new JournaledStateManagerShared(
+                    services.GetRequiredService<ILogger<JournaledStateManager>>(),
+                    Options.Create(new JournaledStateManagerOptions { JournalFormatKey = CustomFormatKey }),
+                    TimeProvider.System,
+                    services),
+                new TestJournalStorageProvider(customStorage)));
+        await using var services = builder.Services.BuildServiceProvider(validateScopes: true);
+        await using var dependencies = services.CreateAsyncScope();
+        var factory = services.GetRequiredKeyedService<IJournaledStateManagerFactory>("custom");
+        await using var defaultManager = services.GetRequiredService<IJournaledStateManagerFactory>().CreateStandalone(new JournalId("default"));
+        await using var customManager = factory.CreateStandalone(new JournalId("custom"));
+        IJournaledStateManager delegating = new DelegatingStateManager(customManager);
+        var codec = dependencies.ServiceProvider.GetRequiredKeyedService<IDurableValueCommandCodec<int>>(CustomFormatKey);
+        Assert.Throws<InvalidOperationException>(() =>
+            services.GetRequiredKeyedService<IDurableValueCommandCodec<int>>(CustomFormatKey));
+        var defaultCodec = services.GetRequiredKeyedService<IDurableValueCommandCodec<int>>(JsonLinesJournalFormat.JournalFormatKey);
+        Assert.NotSame(defaultCodec, codec);
+        var state = new CodecState(codec);
+        Assert.Same(codec, state.Codec);
+        delegating.RegisterStateMachine("value", state);
+        var defaultState = new CodecState(defaultCodec);
+        defaultManager.RegisterStateMachine("value", defaultState);
+        await defaultManager.InitializeAsync(TestContext.Current.CancellationToken);
+        await delegating.InitializeAsync(TestContext.Current.CancellationToken);
+        Assert.Equal(0, state.Value);
+        Assert.Empty(customStorage.Segments);
+        state.Value = 42;
+        defaultState.Value = 7;
+        await delegating.WriteStateAsync(TestContext.Current.CancellationToken);
+        await defaultManager.WriteStateAsync(TestContext.Current.CancellationToken);
+        Assert.Single(customStorage.Segments);
+        await using var recovered = factory.CreateStandalone(new JournalId("custom"));
+        var recoveredState = new CodecState(codec);
+        recovered.RegisterStateMachine("value", recoveredState);
+        await recovered.InitializeAsync(TestContext.Current.CancellationToken);
+        Assert.Equal(42, recoveredState.Value);
+
+        await using var recoveredDefault = services.GetRequiredService<IJournaledStateManagerFactory>().CreateStandalone(new JournalId("default"));
+        var recoveredDefaultState = new CodecState(defaultCodec);
+        recoveredDefault.RegisterStateMachine("value", recoveredDefaultState);
+        await recoveredDefault.InitializeAsync(TestContext.Current.CancellationToken);
+        Assert.Equal(7, recoveredDefaultState.Value);
+    }
+
+    [Fact]
+    public void StateConstruction_MissingFormatCodecFails()
+    {
+        var builder = CreateNamedProviderBuilder();
+        builder.AddJournaling();
+        using var services = builder.Services.BuildServiceProvider();
+        Assert.NotNull(services.GetRequiredKeyedService<IDurableValueCommandCodec<int>>(JsonLinesJournalFormat.JournalFormatKey));
+        var exception = Assert.Throws<InvalidOperationException>(() =>
+            new CodecState(services.GetRequiredKeyedService<IDurableValueCommandCodec<int>>(CustomFormatKey)));
+        Assert.Contains(nameof(IDurableValueCommandCodec<int>), exception.Message);
+    }
+
     [Fact]
     public async Task StateManagerFactory_CreatesManagerForJournalId()
     {
@@ -421,6 +541,36 @@ private sealed class TestJournalStorageProvider(IJournalStorage storage) : IJour
         public IJournalStorage CreateStorage(JournalId journalId) => storage;
     }
 
+    private sealed class NamedBinaryJournalFormat(IJournalFormat inner) : IJournalFormat
+    {
+        public string FormatKey => CustomFormatKey;
+        public string? MimeType => inner.MimeType;
+        public JournalBufferWriter CreateWriter() => inner.CreateWriter();
+        public void Replay(JournalBufferReader input, JournalReplayContext context) => inner.Replay(input, context);
+    }
+
+    private sealed class CodecState(IDurableValueCommandCodec<int> codec) : IStateMachine, IDurableValueCommandHandler<int>
+    {
+        public IDurableValueCommandCodec<int> Codec { get; } = codec;
+        public int Value { get; set; }
+        public void Reset(JournalStreamWriter writer) => Value = 0;
+        public void WritePendingEntries(JournalStreamWriter writer) => Codec.WriteSet(Value, writer);
+        public void WriteSnapshot(JournalStreamWriter writer) => Codec.WriteSet(Value, writer);
+        public void ReplayEntry(JournalEntry entry, JournalReplayContext context) =>
+            context.GetRequiredCommandCodec(entry.FormatKey, Codec).Apply(entry.Reader, this);
+        public void ApplySet(int value) => Value = value;
+    }
+
+    private sealed class DelegatingStateManager(IJournaledStateManager inner) : IJournaledStateManager
+    {
+        public ValueTask InitializeAsync(CancellationToken cancellationToken) => inner.InitializeAsync(cancellationToken);
+        public void RegisterStateMachine(string name, IStateMachine state) => inner.RegisterStateMachine(name, state);
+        public bool TryGetStateMachine(string name, [System.Diagnostics.CodeAnalysis.NotNullWhen(true)] out IStateMachine? state)
+            => inner.TryGetStateMachine(name, out state);
+        public ValueTask WriteStateAsync(CancellationToken cancellationToken) => inner.WriteStateAsync(cancellationToken);
+        public ValueTask DeleteStateAsync(CancellationToken cancellationToken) => inner.DeleteStateAsync(cancellationToken);
+    }
+
     private sealed class LifecycleJournalStorageProvider : IJournalStorageProvider, IJournalStorageCatalog, ILifecycleParticipant<ISiloLifecycle>
     {
         private readonly VolatileJournalStorageProvider _storage = new();