diff --git a/docs/site/src/content/docs/grains/journaling/runtime-behavior.md b/docs/site/src/content/docs/grains/journaling/runtime-behavior.md index e513f837fc9..90d8eb3e3b2 100644 --- a/docs/site/src/content/docs/grains/journaling/runtime-behavior.md +++ b/docs/site/src/content/docs/grains/journaling/runtime-behavior.md @@ -54,7 +54,12 @@ Durable collections encode their operation before applying it to the in-memory c - **Append** adds the encoded operation batch atomically. - **Snapshot replacement** writes the current state of every registered stream and atomically publishes it as the new journal generation. -Concurrent calls made while the same kind of write is queued can share that queued operation. Each caller observes its completion or failure. Calls made after a storage operation starts are processed by a later operation. +Concurrent calls made while a write is queued share that queued operation. A snapshot requirement +from any caller upgrades pending work to a snapshot. The manager consumes the requirement when +selecting a write and restores it if a prerequisite hook fails. Storage acknowledgement leaves +newer pending requirements untouched. +Each caller observes its completion or failure. Calls made after a storage operation starts are +processed by a later operation. > [!IMPORTANT] > In-memory mutation is visible before storage acknowledgement. Return success to a caller only after the required `WriteStateAsync` completes. Recovery reconstructs durable state in a new activation. @@ -68,6 +73,54 @@ Orleans executes that synchronous block on a single activation thread. Another g the operation awaits, so keep shared state safe to commit at each await. Any caller's write can include staged mutations from other calls. +## Journal operation hooks + +Features coordinate prerequisites and completion by registering + callbacks through + and +. +Registration completes synchronously before initialization, supporting constructors. +After initialization begins, idle changes apply synchronously; busy changes queue between +complete operations, after current after callbacks. +The hook set stays stable throughout each operation. Queued changes also separate coalesced writes, +so writes admitted after a change observe that change. +Cancellation ends only the caller's wait while the manager owns admitted changes. +Every owner implements the registration methods and callback boundaries; delegating owners forward +them to their inner owner. Synchronous state callbacks can request a queued change and observe its +completion outside the callback. + +Each actual append, snapshot, or deletion runs before callbacks in reverse list order, outside +the manager lock. Add the final prerequisite at index zero so its callback runs last before +synchronous capture or storage deletion. Prerequisites remain valid through capture or deletion, covering changes staged +during asynchronous preparation and while the work loop resumes. Preserve operation-local bookkeeping for +the captured batch separately from changes staged later. +When the final callback completes synchronously, capture or the start of deletion follows in the +same work-loop continuation. Other owner turns can stage changes while an incomplete callback is awaited, +so its prerequisite covers those changes through capture or deletion. + +Storage acknowledgement and registered-state acknowledgement or reset precede after callbacks. +All after callbacks run in list order, including for successful zero-byte writes. Coalesced callers +share callbacks for the actual operation. Features implement the before and after callbacks on +their own hook, retaining feature identity and operation-local bookkeeping there. + +A failed prerequisite reports with pending +state retained for an explicit persistence retry after the prerequisite is restored. A failed after +callback reports with persistence completed. +Remaining after callbacks run, multiple failures are aggregated, and the manager stays usable. +The feature's durable recovery protocol resumes interrupted post-persistence work. Storage and +state-processing failures retain the manager's fencing and fresh-recovery behavior. + +Hook callbacks receive the owner's shutdown token. Cancelling a caller's wait leaves the owned +operation running through its actual outcome. Disposal drains owned hooks and storage before releasing +journal resources, including when cancellation callbacks or cleanup fail. Concurrent disposal callers +share this completion. Hook implementations complete without calling initialization, persistence, +or disposal on their own owner. Awaiting an operation serialized behind the current callback creates +a circular dependency. Dependencies between journal owners also remain acyclic. +Shutdown closes work admission and cancels queued operations while the current +operation drains to its actual storage and hook outcome. For deletion, the feature owner stops +admission and drains feature operations before +queuing the whole-journal reset. + ## Consistency and competing writers Orleans grain placement normally supplies a single active writer for a grain identity. Journal storage providers also use optimistic concurrency to protect the journal when a stale or competing writer reaches storage. diff --git a/src/Orleans.Journaling/CompatibilitySuppressions.xml b/src/Orleans.Journaling/CompatibilitySuppressions.xml index eb687e7cf91..5e86005fc7b 100644 --- a/src/Orleans.Journaling/CompatibilitySuppressions.xml +++ b/src/Orleans.Journaling/CompatibilitySuppressions.xml @@ -239,6 +239,13 @@ lib/net8.0/Orleans.Journaling.dll true + + CP0006 + M:Orleans.Journaling.IJournaledStateManager.AddHookAsync(Orleans.Journaling.IJournaledStateHook,System.Nullable{System.Int32},System.Threading.CancellationToken) + lib/net10.0/Orleans.Journaling.dll + lib/net10.0/Orleans.Journaling.dll + true + CP0006 M:Orleans.Journaling.IJournaledStateManager.RegisterStateMachine(System.String,Orleans.Journaling.IStateMachine) @@ -246,6 +253,13 @@ lib/net10.0/Orleans.Journaling.dll true + + CP0006 + M:Orleans.Journaling.IJournaledStateManager.RemoveHookAsync(Orleans.Journaling.IJournaledStateHook,System.Threading.CancellationToken) + lib/net10.0/Orleans.Journaling.dll + lib/net10.0/Orleans.Journaling.dll + true + CP0006 M:Orleans.Journaling.IJournaledStateManager.TryGetStateMachine(System.String,Orleans.Journaling.IStateMachine@) @@ -274,6 +288,13 @@ lib/net10.0/Orleans.Journaling.dll true + + CP0006 + M:Orleans.Journaling.IJournaledStateManager.AddHookAsync(Orleans.Journaling.IJournaledStateHook,System.Nullable{System.Int32},System.Threading.CancellationToken) + lib/net8.0/Orleans.Journaling.dll + lib/net8.0/Orleans.Journaling.dll + true + CP0006 M:Orleans.Journaling.IJournaledStateManager.RegisterStateMachine(System.String,Orleans.Journaling.IStateMachine) @@ -281,6 +302,13 @@ lib/net8.0/Orleans.Journaling.dll true + + CP0006 + M:Orleans.Journaling.IJournaledStateManager.RemoveHookAsync(Orleans.Journaling.IJournaledStateHook,System.Threading.CancellationToken) + lib/net8.0/Orleans.Journaling.dll + lib/net8.0/Orleans.Journaling.dll + true + CP0006 M:Orleans.Journaling.IJournaledStateManager.TryGetStateMachine(System.String,Orleans.Journaling.IStateMachine@) diff --git a/src/Orleans.Journaling/IJournaledStateHook.cs b/src/Orleans.Journaling/IJournaledStateHook.cs new file mode 100644 index 00000000000..c2e187254ff --- /dev/null +++ b/src/Orleans.Journaling/IJournaledStateHook.cs @@ -0,0 +1,64 @@ +namespace Orleans.Journaling; + +/// +/// Identifies a journal persistence operation. +/// +public enum JournaledStateOperation +{ + /// Appends pending journal entries. + Write, + + /// Replaces storage with a snapshot. + Snapshot, + + /// Deletes storage and resets registered states. + Delete +} + +/// +/// Participates in the prerequisites and completion of actual journal operations. +/// +/// +/// Before hooks run in reverse list order on the owner's logical execution context, outside its lock. +/// The hook at index zero supplies the final prerequisite before capture or deletion. +/// When the final callback completes synchronously, capture or the start of deletion follows in the +/// same work-loop continuation. +/// Prerequisites remain valid until capture or deletion, covering changes staged during asynchronous +/// preparation and while the journal work loop resumes. After hooks run in forward list order after +/// storage acknowledgement and state acknowledgement or reset, including successful writes which +/// produce no storage bytes. Hooks retain operation-local data across these boundaries and keep later +/// pending changes separate. Hook implementations must complete without calling initialization, +/// persistence, or disposal on their own journal owner. Those operations are serialized behind the +/// current operation, so awaiting them from a hook would create a circular dependency. +/// Dependencies between journal owners must also remain acyclic. +/// Callbacks can request a hook registration change, observing its queued completion outside the +/// callback rather than awaiting work which follows the current operation. +/// +public interface IJournaledStateHook +{ + /// + /// Establishes prerequisites before capture or storage deletion. + /// + /// The operation about to execute. + /// The token for the owned operation's lifetime. + /// A completion representing the prerequisite work. + /// + /// Failure reports and retains pending changes + /// for an explicit retry. All staged changes remain safe to commit. Full deletion requires the + /// owner to stop admission and drain feature operations before queuing deletion. + /// + ValueTask BeforeOperationAsync(JournaledStateOperation operation, CancellationToken cancellationToken) => default; + + /// + /// Completes post-persistence work after the operation succeeds. + /// + /// The completed operation. + /// The token for the owned operation's lifetime. + /// A completion representing post-persistence work. + /// + /// Every after hook is invoked even when an earlier after hook fails. Failures are surfaced as + /// and leave the successfully persisted owner usable. + /// Durable feature state supplies recovery for interrupted post-persistence work. + /// + ValueTask AfterOperationAsync(JournaledStateOperation operation, CancellationToken cancellationToken) => default; +} diff --git a/src/Orleans.Journaling/IJournaledStateManager.cs b/src/Orleans.Journaling/IJournaledStateManager.cs index b6363a1b4cc..5c18162570f 100644 --- a/src/Orleans.Journaling/IJournaledStateManager.cs +++ b/src/Orleans.Journaling/IJournaledStateManager.cs @@ -8,10 +8,44 @@ namespace Orleans.Journaling; /// /// The owner registers state machines and initializes the journal before using recovered state. /// State machine instances and their dependencies retain the lifetime assigned by their caller. -/// Disposing this manager stops journal processing and releases its journal resources. +/// Disposal drains owned storage and hook operations before releasing resources. /// public interface IJournaledStateManager : IAsyncDisposable { + /// + /// Adds a hook before the next persistence operation which follows this registration. + /// + /// The hook to register. + /// The insertion index, or to append. + /// The token used to cancel the caller's wait. + /// A completion representing application of the registration. + /// + /// Registration completes synchronously before initialization begins, supporting constructor registration. + /// After initialization begins, idle registration applies synchronously; otherwise it is queued + /// between owned operations, after the current operation's after hooks finish. + /// Before hooks execute in reverse registration order; after hooks execute in forward order. + /// Register the final capture prerequisite at index zero so it runs last before capture or deletion. + /// Registration is independent of state-machine registration and persists through recovery and deletion. + /// Every implementation provides registration and invokes hooks at the documented operation boundaries. + /// Queued registration is a write-coalescing barrier. + /// Cancellation ends the caller's wait while admitted registration remains owned by the manager. + /// Synchronous state callbacks can request registration and observe its completion outside the callback. + /// Await registration before staging changes whose safety depends on the added hook. + /// + ValueTask AddHookAsync(IJournaledStateHook hook, int? index = null, CancellationToken cancellationToken = default); + + /// + /// Removes a hook between persistence operations. + /// + /// The registered hook to remove. + /// The token used to cancel the caller's wait. + /// A completion representing application of the removal. + /// The hook is not registered when the removal is applied. + /// + /// Removal has the same ownership and ordering guarantees as . + /// + ValueTask RemoveHookAsync(IJournaledStateHook hook, CancellationToken cancellationToken = default); + /// ValueTask IAsyncDisposable.DisposeAsync() => default; @@ -51,8 +85,11 @@ public interface IJournaledStateManager : IAsyncDisposable /// /// Stage mutations only after establishing that they are safe to commit. Pending changes are shared /// by all callers using this manager. Storage acknowledgement establishes durability. - /// A failed journal operation fences the manager; recovery requires a new manager and state machine instances. + /// A storage or state-processing failure fences the manager; recovery requires a new manager and state machine instances. /// Cancellation stops the caller's wait; an already queued write continues to its storage outcome. + /// Before-hook failure reports and retains pending + /// changes for an explicit retry. After-hook failure reports + /// after successful persistence. /// /// The token used to cancel the caller's wait. /// A task representing the write acknowledgement. @@ -64,7 +101,10 @@ public interface IJournaledStateManager : IAsyncDisposable /// /// 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. + /// A storage or state-reset failure permanently fences the manager and requests deactivation of its owning grain. + /// Before-hook failure reports with state retained. + /// After-hook failure reports after storage deletion + /// and state reset succeed. /// /// The cancellation token. /// A which represents the operation. diff --git a/src/Orleans.Journaling/JournaledStateManager.Hooks.cs b/src/Orleans.Journaling/JournaledStateManager.Hooks.cs new file mode 100644 index 00000000000..115729e6ac5 --- /dev/null +++ b/src/Orleans.Journaling/JournaledStateManager.Hooks.cs @@ -0,0 +1,105 @@ +using System.Runtime.ExceptionServices; + +namespace Orleans.Journaling; + +internal partial class JournaledStateManager +{ + private List? _hooks; + + public ValueTask AddHookAsync(IJournaledStateHook hook, int? index = null, CancellationToken cancellationToken = default) => + MutateHookAsync(hook, index, remove: false, cancellationToken); + + public ValueTask RemoveHookAsync(IJournaledStateHook hook, CancellationToken cancellationToken = default) => + MutateHookAsync(hook, index: null, remove: true, cancellationToken); + + private ValueTask MutateHookAsync(IJournaledStateHook hook, int? index, bool remove, CancellationToken cancellationToken) + { + ArgumentNullException.ThrowIfNull(hook); + cancellationToken.ThrowIfCancellationRequested(); + Task task; + lock (_lock) + { + ThrowIfFenced(); + ObjectDisposedException.ThrowIf(_disposeTask is not null, this); + _shutdownCancellation.Token.ThrowIfCancellationRequested(); + if (_workLoop is null || (!_workItemRunning && _workQueue.Count == 0)) + { + ApplyHookMutation(hook, index, remove); + return default; + } + + var workItem = new HookMutationWorkItem(hook, index, remove); + _workQueue.Enqueue(workItem); + task = workItem.Task; + } + + _workSignal.Signal(); + return new(task.WaitAsync(cancellationToken)); + } + + private void ApplyHookMutation(IJournaledStateHook hook, int? index, bool remove) + { + var hooks = _hooks ??= new(); + if (remove) + { + if (!hooks.Remove(hook)) + { + throw new InvalidOperationException("The journal operation hook is not registered."); + } + } + else + { + hooks.Insert(index ?? hooks.Count, hook); + } + } + + private ValueTask InvokeBeforeHooksAsync(ref int index, JournaledStateOperation operation, CancellationToken cancellationToken) + { + while (index >= 0) + { + cancellationToken.ThrowIfCancellationRequested(); + var pending = _hooks![index--].BeforeOperationAsync(operation, cancellationToken); + if (!pending.IsCompleted) + { + return pending; + } + + pending.GetAwaiter().GetResult(); + } + + return default; + } + + private async ValueTask InvokeAfterHooksAsync(JournaledStateOperation operation, CancellationToken cancellationToken) + { + List? failures = null; + for (var i = 0; i < _hooks!.Count; i++) + { + try + { + await _hooks[i].AfterOperationAsync(operation, cancellationToken).ConfigureAwait(true); + } + catch (Exception exception) + { + (failures ??= []).Add(exception); + } + } + + if (failures is { Count: 1 }) + { + ExceptionDispatchInfo.Capture(failures[0]).Throw(); + } + + if (failures is not null) + { + throw new AggregateException(failures); + } + } + + private sealed class HookMutationWorkItem(IJournaledStateHook hook, int? index, bool remove) : WorkItem + { + public IJournaledStateHook Hook { get; } = hook; + public int? Index { get; } = index; + public bool Remove { get; } = remove; + } +} diff --git a/src/Orleans.Journaling/JournaledStateManager.cs b/src/Orleans.Journaling/JournaledStateManager.cs index cbf6be752bc..bd0b1ae47bb 100644 --- a/src/Orleans.Journaling/JournaledStateManager.cs +++ b/src/Orleans.Journaling/JournaledStateManager.cs @@ -1,6 +1,7 @@ using System.Buffers; using System.Diagnostics; using System.Diagnostics.CodeAnalysis; +using System.Runtime.ExceptionServices; using Microsoft.Extensions.Logging; using Orleans.Diagnostics; using Orleans.Serialization.Buffers; @@ -28,11 +29,14 @@ internal partial class JournaledStateManager : IJournaledStateManager, IJournalS private readonly StateDirectory _journalStreamDirectory; private readonly RetiredStateTracker _retirementTracker; private Task? _workLoop; + // Under _lock, covers recovery and dequeued work through the last after callback. + private bool _workItemRunning; private ManagerState _state; private long _deletionGeneration; private Exception? _failure; - private bool _migrationSnapshotRequired; - private int _disposed; + // A snapshot requirement not yet claimed by a write. Consume under the lock before preparation awaits. + private bool _snapshotRequired; + private Task? _disposeTask; public JournaledStateManager(JournaledStateManagerShared shared, IJournalStorageProvider storageProvider, IGrainContext grainContext) : this(shared, CreateStorage(storageProvider, CreateJournalId(grainContext)), grainContext) @@ -98,6 +102,8 @@ internal static IJournalStorage CreateStorage(IJournalStorageProvider storagePro internal string WriteJournalFormatKey => _shared.JournalFormatKey; + private bool IsSnapshotRequired => _snapshotRequired || _storage.IsCompactionRequested; + internal IServiceProvider ServiceProvider => _grainContext is { } context ? context.ActivationServices : _shared.ServiceProvider; public bool TryGetStateMachine(string name, [NotNullWhen(true)] out IStateMachine? stateMachine) @@ -106,7 +112,7 @@ public bool TryGetStateMachine(string name, [NotNullWhen(true)] out IStateMachin lock (_lock) { ThrowIfFenced(); - ObjectDisposedException.ThrowIf(_disposed != 0, this); + ObjectDisposedException.ThrowIf(_disposeTask is not null, this); if (_states.TryGetValue(name, out var existing) && existing is not RetiredStateMachine) { stateMachine = existing; @@ -123,11 +129,11 @@ protected void EnsureRegistrationAllowed() lock (_lock) { ThrowIfFenced(); - ObjectDisposedException.ThrowIf(_disposed != 0, this); + ObjectDisposedException.ThrowIf(_disposeTask is not null, this); _shutdownCancellation.Token.ThrowIfCancellationRequested(); if (_workLoop is not null) { - throw new InvalidOperationException("New states cannot be registered after journaled state manager initialization has begun."); + throw new InvalidOperationException("Journal registrations cannot be changed after journaled state manager initialization has begun."); } } } @@ -164,10 +170,11 @@ public async ValueTask InitializeAsync(CancellationToken cancellationToken = def ThrowIfFenced(); if (_workLoop is null) { + _workItemRunning = true; _workLoop = Start(); } - task = EnqueueOrGetPendingWorkItem(out didEnqueue); + task = EnqueueOrGetPendingWorkItem(out didEnqueue).Task; } if (didEnqueue) @@ -189,6 +196,12 @@ private async Task WorkLoop() await Task.CompletedTask.ConfigureAwait(ConfigureAwaitOptions.ContinueOnCapturedContext | ConfigureAwaitOptions.ForceYielding); while (!_shutdownCancellation.IsCancellationRequested) { + lock (_lock) + { + _workItemRunning = true; + } + + Exception? recoveryFailure = null; try { await RecoverAsync(_shutdownCancellation.Token).ConfigureAwait(true); @@ -201,15 +214,17 @@ private async Task WorkLoop() } catch (Exception exception) { - try - { - LogErrorProcessingWorkItems(_shared.Logger, exception); - } - finally + recoveryFailure = exception; + LogErrorProcessingWorkItems(_shared.Logger, exception); + } + finally + { + lock (_lock) { - lock (_lock) + _workItemRunning = false; + if (recoveryFailure is not null) { - FaultQueuedWorkItemsUnderLock(exception); + FaultQueuedWorkItemsUnderLock(recoveryFailure); } } } @@ -245,6 +260,11 @@ private async Task WorkLoop() WorkItem workItem; lock (_lock) { + if (_shutdownCancellation.IsCancellationRequested) + { + return; + } + if (!_workQueue.TryDequeue(out var dequeuedWorkItem)) { // Wait for the queue to be signaled again. @@ -252,59 +272,77 @@ private async Task WorkLoop() } workItem = dequeuedWorkItem; + _workItemRunning = true; } var processingTimestamp = _shared.TimeProvider.GetTimestamp(); - var queueOperation = GetStorageQueueOperation(workItem); - var queueDuration = queueOperation is null - ? default - : _shared.TimeProvider.GetElapsedTime(workItem.EnqueuedTimestamp, processingTimestamp); + string? queueOperation = null; + TimeSpan queueDuration = default; var recordQueueDuration = workItem is DeleteStateWorkItem; - Activity? storageActivity = queueOperation is null - ? null - : ActivitySources.StorageGrainSource.StartActivity( - $"journal {queueOperation}", - ActivityKind.Internal, - workItem.TraceParent, - tags: null, - links: workItem.TraceLinks); - if (storageActivity is not null) - { - storageActivity.SetTag(ActivityTagKeys.JournalStorageOperation, queueOperation); - } + Activity? storageActivity = null; + JournaledStateOperation? operationType = null; + string? compactionReason = null; + var beforeHookRunning = false; + var afterHookRunning = false; try { - if (workItem is AppendJournalWorkItem or WriteSnapshotWorkItem - && workItem.DeletionGeneration != Volatile.Read(ref _deletionGeneration)) + lock (_lock) + { + if (workItem is WriteStateWorkItem + && workItem.DeletionGeneration != _deletionGeneration) + { + throw new InvalidOperationException( + "The journaled state operation was queued before deletion reset its state."); + } + + operationType = GetOperationType(workItem); + if (operationType == JournaledStateOperation.Snapshot) + { + compactionReason = _snapshotRequired + ? JournalingInstruments.CompactionReasonUserSnapshot + : JournalingInstruments.CompactionReasonStorageRequested; + _snapshotRequired = false; + } + } + + (queueOperation, queueDuration, storageActivity) = GetStorageQueueOperation(workItem, operationType, processingTimestamp); + + if (operationType is { } operation && _hooks is { Count: > 0 }) { - throw new InvalidOperationException( - "The journaled state operation was queued before deletion reset its state."); + beforeHookRunning = true; + // Keep synchronous final completion and capture in the same work-loop continuation. + for (var i = _hooks.Count - 1; i >= 0;) + { + await InvokeBeforeHooksAsync(ref i, operation, _shutdownCancellation.Token).ConfigureAwait(true); + } + + _shutdownCancellation.Token.ThrowIfCancellationRequested(); + beforeHookRunning = false; } // Note that the implementation of each command is inlined to avoid allocating unnecessary async states. // We are ok sacrificing some code organization for performance in the inner loop. switch (workItem) { - case AppendJournalWorkItem: - case WriteSnapshotWorkItem: + case HookMutationWorkItem mutation: + { + // The work loop owns the hook set while _workItemRunning is true. + ApplyHookMutation(mutation.Hook, mutation.Index, mutation.Remove); + break; + } + + case WriteStateWorkItem: { // 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 - || _migrationSnapshotRequired - || _storage.IsCompactionRequested; + var isSnapshot = operationType == JournaledStateOperation.Snapshot; var operationLabel = isSnapshot ? JournalingInstruments.OperationSnapshot : JournalingInstruments.OperationAppend; if (isSnapshot) { - var compactionReason = workItem is WriteSnapshotWorkItem - ? JournalingInstruments.CompactionReasonUserSnapshot - : _migrationSnapshotRequired - ? JournalingInstruments.CompactionReasonMigration - : JournalingInstruments.CompactionReasonStorageRequested; - _shared.Instruments.OnCompactionTriggered(compactionReason); + _shared.Instruments.OnCompactionTriggered(compactionReason!); } _shared.Instruments.OnWriteCoalesced(operationLabel, workItem.CallerCount); var gatherStartTimestamp = _shared.TimeProvider.GetTimestamp(); @@ -325,11 +363,6 @@ private async Task WorkLoop() RetireOrResurectStates(); } - if (_migrationSnapshotRequired) - { - ThrowIfMigrationBlockedByRetiredStates(); - } - // The map of state ids is itself stored as a durable state with the id 0. // This must be stored first, since it includes the identities of all other states, which are needed when replaying the journal. // If we removed retired states, this snapshot will persist that change. @@ -470,11 +503,6 @@ private async Task WorkLoop() { state.OnWriteCompleted(); } - - if (isSnapshot) - { - _migrationSnapshotRequired = false; - } } } @@ -540,13 +568,64 @@ private async Task WorkLoop() } } + if (operationType is { } completedOperation && _hooks is { Count: > 0 }) + { + afterHookRunning = true; + await InvokeAfterHooksAsync(completedOperation, _shutdownCancellation.Token).ConfigureAwait(true); + afterHookRunning = false; + } + if (recordQueueDuration && queueOperation is not null) { _shared.Instruments.OnStorageOperationQueued(queueOperation, queueDuration, succeeded: true); } storageActivity?.SetStatus(ActivityStatusCode.Ok); - workItem.SetResult(); + lock (_lock) + { + _workItemRunning = false; + workItem.SetResult(); + } + } + catch (Exception exception) when (beforeHookRunning || afterHookRunning) + { + if (beforeHookRunning && operationType == JournaledStateOperation.Snapshot) + { + lock (_lock) + { + _snapshotRequired = true; + } + } + + if (beforeHookRunning && exception is OperationCanceledException && _shutdownCancellation.IsCancellationRequested) + { + workItem.TrySetCanceled(_shutdownCancellation.Token); + continue; + } + + LogOperationHookFailed(_shared.Logger, exception, operationType!.Value, afterHookRunning); + storageActivity?.SetStatus(ActivityStatusCode.Error, "Journal operation hook failed."); + if (recordQueueDuration && queueOperation is not null) + { + _shared.Instruments.OnStorageOperationQueued(queueOperation, queueDuration, succeeded: afterHookRunning); + } + + lock (_lock) + { + _workItemRunning = false; + workItem.SetException(afterHookRunning + ? new JournaledStatePostCommitException(operationType.Value, exception) + : new JournaledStatePreCommitException(operationType.Value, exception)); + } + } + catch (Exception exception) when (workItem is HookMutationWorkItem) + { + LogHookMutationFailed(_shared.Logger, exception); + lock (_lock) + { + _workItemRunning = false; + workItem.SetException(exception); + } } catch (Exception exception) { @@ -580,6 +659,11 @@ private async Task WorkLoop() } finally { + lock (_lock) + { + _workItemRunning = false; + } + storageActivity?.Dispose(); } } @@ -596,6 +680,14 @@ private async Task WorkLoop() } } + private JournaledStateOperation? GetOperationType(WorkItem workItem) => workItem switch + { + WriteStateWorkItem when IsSnapshotRequired => JournaledStateOperation.Snapshot, + WriteStateWorkItem => JournaledStateOperation.Write, + DeleteStateWorkItem => JournaledStateOperation.Delete, + _ => null + }; + private void Fence(Exception exception) { lock (_lock) @@ -668,20 +760,6 @@ private void RetireOrResurectStates() } } - private void ThrowIfMigrationBlockedByRetiredStates() - { - foreach (var state in _statesMap.Values) - { - if (state is RetiredStateMachine { PreservedEntries.Count: > 0 } retiredState) - { - throw new InvalidOperationException( - $"Cannot migrate journal to format key '{_shared.JournalFormatKey}' because stream " + - $"{retiredState.StreamId.Value} belongs to a state which is not currently registered. Register the state so it can be decoded " + - "and snapshotted in the configured format, or keep the previous journal format configured until the state can be retired."); - } - } - } - private JournalStreamWriter CreateJournalStreamWriter(JournalStreamId streamId) => _journalWriter.CreateJournalStreamWriter(streamId); private static void AppendUpdatesOrSnapshotState(JournalBufferWriter journalWriter, bool isSnapshot, uint id, IStateMachine state) @@ -705,7 +783,7 @@ public async ValueTask DeleteStateAsync(CancellationToken cancellationToken = de lock (_lock) { ThrowIfStateOperationsUnavailable(); - task = EnqueueOrGetPendingWorkItem(out didEnqueue); + task = EnqueueOrGetPendingWorkItem(out didEnqueue).Task; } if (didEnqueue) @@ -780,7 +858,7 @@ private async Task RecoverAsync(CancellationToken cancellationToken) private void ResetForRecovery() { _journalWriter.Reset(); - _migrationSnapshotRequired = false; + _snapshotRequired = false; _statesMap.Clear(); _statesMap[StateDirectory.Id] = _journalStreamDirectory; _statesMap[RetiredStateTracker.Id] = _retirementTracker; @@ -851,7 +929,7 @@ private void ProcessRecoveryBuffer(JournalBufferReader buffer, IJournalMetadata? { if (!string.Equals(journalFormatKey, _shared.JournalFormatKey, StringComparison.Ordinal)) { - _migrationSnapshotRequired = true; + _snapshotRequired = true; } var journalFormat = GetJournalFormat(journalFormatKey); @@ -894,11 +972,9 @@ public async ValueTask WriteStateAsync(CancellationToken cancellationToken = def lock (_lock) { ThrowIfStateOperationsUnavailable(); - var isSnapshot = _migrationSnapshotRequired || _storage.IsCompactionRequested; - operation = isSnapshot ? JournalingInstruments.OperationSnapshot : JournalingInstruments.OperationAppend; - pendingWrite = isSnapshot - ? EnqueueOrGetPendingWorkItem(out didEnqueue) - : EnqueueOrGetPendingWorkItem(out didEnqueue); + _snapshotRequired = IsSnapshotRequired; + operation = _snapshotRequired ? JournalingInstruments.OperationSnapshot : JournalingInstruments.OperationAppend; + pendingWrite = EnqueueOrGetPendingWorkItem(out didEnqueue).Task; } if (didEnqueue) @@ -999,26 +1075,39 @@ private async ValueTask DeleteStorageAsync(CancellationToken cancellationToken) } } - private Task EnqueueOrGetPendingWorkItem(out bool didEnqueue) + private TWorkItem EnqueueOrGetPendingWorkItem(out bool didEnqueue) where TWorkItem : WorkItem, new() { + TWorkItem? pending = null; foreach (var workItem in _workQueue) { - if (workItem.GetType() != typeof(TWorkItem)) + if (workItem is HookMutationWorkItem) + { + // Later callers must not coalesce into work which precedes a registration change. + pending = null; + continue; + } + + if (pending is not null || workItem is not TWorkItem candidate) { continue; } - if (workItem is AppendJournalWorkItem or WriteSnapshotWorkItem + if (workItem is WriteStateWorkItem && workItem.DeletionGeneration != _deletionGeneration) { continue; } - workItem.RecordTraceContext(); - workItem.AddCaller(); + pending = candidate; + } + + if (pending is not null) + { + pending.RecordTraceContext(); + pending.AddCaller(); didEnqueue = false; - return workItem.Task; + return pending; } var newWorkItem = new TWorkItem(); @@ -1027,7 +1116,7 @@ private Task EnqueueOrGetPendingWorkItem(out bool didEnqueue) newWorkItem.RecordTraceContext(); _workQueue.Enqueue(newWorkItem); didEnqueue = true; - return newWorkItem.Task; + return newWorkItem; } private void BindState(string name, uint id) @@ -1062,9 +1151,22 @@ void ILifecycleParticipant.Participate(IGrainLifecycle observer private async Task StopAsync(CancellationToken cancellationToken) { + AggregateException? cancellationFailure = null; lock (_lock) { - _shutdownCancellation.Cancel(); + try + { + _shutdownCancellation.Cancel(); + } + catch (AggregateException exception) + { + cancellationFailure = exception; + } + } + + if (cancellationFailure is not null) + { + LogShutdownCancellationFailed(_shared.Logger, cancellationFailure); } _workSignal.Signal(); @@ -1079,6 +1181,11 @@ private async Task StopAsync(CancellationToken cancellationToken) { CancelQueuedWorkItems(_shutdownCancellation.Token); } + + if (cancellationFailure is not null) + { + ExceptionDispatchInfo.Capture(cancellationFailure).Throw(); + } } private void CancelQueuedWorkItems(CancellationToken cancellationToken) @@ -1097,21 +1204,41 @@ void IDisposable.Dispose() DisposeAsync().AsTask().GetAwaiter().GetResult(); } - public async ValueTask DisposeAsync() + public ValueTask DisposeAsync() { - if (Interlocked.Exchange(ref _disposed, 1) != 0) + lock (_lock) { - return; + if (_disposeTask is null) + { + var completion = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + // Publish before shutdown cancellation can reenter disposal. + _disposeTask = completion.Task; + _ = DisposeCoreAsync(completion); + } + + return new(_disposeTask); } + } + private async Task DisposeCoreAsync(TaskCompletionSource completion) + { try { - await StopAsync(CancellationToken.None).ConfigureAwait(false); + try + { + await StopAsync(CancellationToken.None).ConfigureAwait(false); + } + finally + { + _shutdownCancellation.Dispose(); + _journalWriter.Dispose(); + } + + completion.SetResult(); } - finally + catch (Exception exception) { - _shutdownCancellation.Dispose(); - _journalWriter.Dispose(); + completion.SetException(exception); } } @@ -1193,20 +1320,37 @@ public void RecordTraceContext() } } - private static string? GetStorageQueueOperation(WorkItem workItem) => - workItem switch + private (string? Operation, TimeSpan QueueDuration, Activity? Activity) GetStorageQueueOperation( + WorkItem workItem, + JournaledStateOperation? operationType, + long processingTimestamp) + { + var operation = operationType switch { - AppendJournalWorkItem => JournalingInstruments.OperationAppend, - WriteSnapshotWorkItem => JournalingInstruments.OperationSnapshot, - DeleteStateWorkItem => JournalingInstruments.OperationDelete, + JournaledStateOperation.Snapshot => JournalingInstruments.OperationSnapshot, + JournaledStateOperation.Write => JournalingInstruments.OperationAppend, + JournaledStateOperation.Delete => JournalingInstruments.OperationDelete, _ => null }; + if (operation is null) + { + return default; + } - private sealed class InitializeWorkItem : WorkItem; + var queueDuration = _shared.TimeProvider.GetElapsedTime(workItem.EnqueuedTimestamp, processingTimestamp); + var activity = ActivitySources.StorageGrainSource.StartActivity( + $"journal {operation}", + ActivityKind.Internal, + workItem.TraceParent, + tags: null, + links: workItem.TraceLinks); + activity?.SetTag(ActivityTagKeys.JournalStorageOperation, operation); + return (operation, queueDuration, activity); + } - private sealed class AppendJournalWorkItem : WorkItem; + private sealed class InitializeWorkItem : WorkItem; - private sealed class WriteSnapshotWorkItem : WorkItem; + private sealed class WriteStateWorkItem : WorkItem; private sealed class DeleteStateWorkItem : WorkItem; @@ -1330,8 +1474,6 @@ private sealed class RetiredStateMachine(JournalStreamId streamId) : IStateMachi public JournalStreamId StreamId { get; } = streamId; - public IReadOnlyList PreservedEntries => _preservedEntries; - void IStateMachine.ReplayEntry(JournalEntry entry, JournalReplayContext context) => _preservedEntries.Add(new PreservedJournalEntry(entry.FormatKey, entry.Reader)); @@ -1352,6 +1494,21 @@ void IStateMachine.WritePendingEntries(JournalStreamWriter writer) { } Message = "Error processing work items.")] private static partial void LogErrorProcessingWorkItems(ILogger logger, Exception exception); + [LoggerMessage( + Level = LogLevel.Error, + Message = "Journal {Operation} hook failed. Persistence completed: {Committed}.")] + private static partial void LogOperationHookFailed(ILogger logger, Exception exception, JournaledStateOperation operation, bool committed); + + [LoggerMessage( + Level = LogLevel.Error, + Message = "Journal operation hook registration change failed.")] + private static partial void LogHookMutationFailed(ILogger logger, Exception exception); + + [LoggerMessage( + Level = LogLevel.Error, + Message = "Journal shutdown cancellation callback failed; owned operations are drained before resources are released.")] + private static partial void LogShutdownCancellationFailed(ILogger logger, Exception exception); + [LoggerMessage( Level = LogLevel.Information, Message = "State \"{Name}\" was not found. I have substituted a placeholder for graceful time-based retirement.")] diff --git a/src/Orleans.Journaling/JournaledStatePostCommitException.cs b/src/Orleans.Journaling/JournaledStatePostCommitException.cs new file mode 100644 index 00000000000..8274c8726d6 --- /dev/null +++ b/src/Orleans.Journaling/JournaledStatePostCommitException.cs @@ -0,0 +1,22 @@ +namespace Orleans.Journaling; + +/// +/// Reports a hook failure after the journal operation and its state acknowledgement or reset succeeded. +/// +/// +/// The journal owner remains usable. The caller handles the failed post-persistence work using +/// the feature's durable recovery protocol, preserving the completed business operation. +/// +[GenerateSerializer] +public sealed class JournaledStatePostCommitException : Exception +{ + /// + /// Initializes a post-persistence hook failure. + /// + /// The successfully completed journal operation. + /// The failure or aggregate of failures from after hooks. + public JournaledStatePostCommitException(JournaledStateOperation operation, Exception innerException) + : base($"Journal operation '{operation}' completed, but a post-persistence hook failed.", innerException) + { + } +} diff --git a/src/Orleans.Journaling/JournaledStatePreCommitException.cs b/src/Orleans.Journaling/JournaledStatePreCommitException.cs new file mode 100644 index 00000000000..fe8d2eb79d4 --- /dev/null +++ b/src/Orleans.Journaling/JournaledStatePreCommitException.cs @@ -0,0 +1,22 @@ +namespace Orleans.Journaling; + +/// +/// Reports a prerequisite hook failure which prevented the journal storage operation. +/// +/// +/// Pending changes remain staged and safe to commit. The caller can restore the prerequisite +/// and explicitly retry persistence, or retire its owner and recover from durable state. +/// +[GenerateSerializer] +public sealed class JournaledStatePreCommitException : Exception +{ + /// + /// Initializes a prerequisite hook failure. + /// + /// The journal operation prevented by the failed prerequisite. + /// The original prerequisite failure. + public JournaledStatePreCommitException(JournaledStateOperation operation, Exception innerException) + : base($"Journal operation '{operation}' was prevented by a prerequisite hook failure.", innerException) + { + } +} diff --git a/src/Orleans.Journaling/README.md b/src/Orleans.Journaling/README.md index 2f5046a30be..d50c3b34275 100644 --- a/src/Orleans.Journaling/README.md +++ b/src/Orleans.Journaling/README.md @@ -264,6 +264,56 @@ the operation's token before use, or deliberately supply them through a registra lifecycle ownership. Creation through the explicit-journal factory keeps failure handling independent of the ambient grain context, including when the caller subsequently enrolls the manager in a lifecycle. +## Journal operation hooks + +Use `IJournaledStateManager.AddHookAsync(hook, index)` and `RemoveHookAsync(hook)` to register +`IJournaledStateHook` callbacks. Registration completes synchronously before initialization, +supporting constructors. After initialization begins, changes apply synchronously when idle, +or queue between complete journal operations while busy. +Each operation keeps its hook set through preparation, capture, storage ACK and after callbacks. +Queued changes are write-coalescing barriers: writes admitted afterward observe the changed set. +Caller cancellation ends only its wait for an admitted change; the manager retains ownership. +Await registration before staging changes whose safety depends on the new hook. +Every journal owner implementation provides registration and invokes callbacks at the operation +boundaries described below. Delegating owners forward registration methods to their inner owner. +Synchronous state callbacks can request a queued registration and observe completion outside the callback. + +For each actual write, snapshot, or deletion, before hooks run in reverse list order outside +the manager lock. Add the final prerequisite with `index: 0` so it runs last before the work loop +captures the registered states or starts deletion. Prerequisites remain valid through capture or +deletion, covering changes staged during asynchronous preparation and while the work loop resumes. +When the final callback completes synchronously, capture or the start of deletion follows in the +same work-loop continuation. Other owner turns can stage changes while an incomplete callback is awaited, +so its prerequisite covers those changes through capture or deletion. +Keep operation-local bookkeeping for the captured batch separate from later pending changes. + +After hooks run in list order after storage acknowledgement and state acknowledgement or reset. +They also run for a successful zero-byte write. Coalesced callers share the hooks for their actual +operation. Features implement the before and after callbacks on their own hook, retaining their +identity and operation-local bookkeeping there. +Queued writes coalesce into one operation. A snapshot requirement from any caller upgrades that +pending work to a snapshot, even if the storage compaction signal clears before it runs. +The manager consumes the pending requirement when selecting a write and restores it when a before +hook fails. Storage acknowledgement leaves newer snapshot requests untouched. + +| Outcome | Owner and caller behavior | +| --- | --- | +| Before hook fails | `JournaledStatePreCommitException` retains pending state for an explicit retry after restoring the prerequisite. | +| Storage or state processing fails | The original failure fences the manager; create a fresh owner and recover the durable outcome. | +| After hook fails | `JournaledStatePostCommitException` reports successful persistence. Every remaining after hook runs, failures are aggregated, and the manager stays usable. The feature's durable recovery protocol resumes interrupted post-persistence work. | + +Hooks receive the owner's shutdown token. Caller cancellation ends the caller's wait while +the owned prerequisite, capture, storage, and completion phases continue. Disposal cancels the +owner token and drains owned work before releasing journal resources, including when cancellation +callbacks or after-hook cleanup fail. Concurrent disposal callers share that drain and its outcome. +Shutdown closes work admission and cancels queued operations while an already running operation +drains to its actual storage and hook outcome. +Hook implementations complete without calling initialization, persistence, or disposal on their +own owner: awaiting an operation serialized behind the current callback creates a circular dependency. +Dependencies between journal owners also remain acyclic. +Before whole-journal deletion, the feature owner stops admission and drains its own operations; +deletion completion follows storage deletion and registered-state reset. + ## State identity and retirement Preserve state names across activations and deployments. A stream absent from the setup declarations diff --git a/src/api/Orleans.Journaling/Orleans.Journaling.cs b/src/api/Orleans.Journaling/Orleans.Journaling.cs index 396a1692cad..cf2ecdf72e7 100644 --- a/src/api/Orleans.Journaling/Orleans.Journaling.cs +++ b/src/api/Orleans.Journaling/Orleans.Journaling.cs @@ -219,13 +219,21 @@ public partial interface IDurableValue T? Value { get; set; } } + public partial interface IJournaledStateHook + { + System.Threading.Tasks.ValueTask AfterOperationAsync(JournaledStateOperation operation, System.Threading.CancellationToken cancellationToken); + System.Threading.Tasks.ValueTask BeforeOperationAsync(JournaledStateOperation operation, System.Threading.CancellationToken cancellationToken); + } + public partial interface IJournaledStateManager : System.IAsyncDisposable { long PendingWriteByteCount { get; } + System.Threading.Tasks.ValueTask AddHookAsync(IJournaledStateHook hook, int? index = null, System.Threading.CancellationToken cancellationToken = default); System.Threading.Tasks.ValueTask DeleteStateAsync(System.Threading.CancellationToken cancellationToken = default); System.Threading.Tasks.ValueTask InitializeAsync(System.Threading.CancellationToken cancellationToken = default); void RegisterStateMachine(string name, IStateMachine stateMachine); + System.Threading.Tasks.ValueTask RemoveHookAsync(IJournaledStateHook hook, System.Threading.CancellationToken cancellationToken = default); System.Threading.Tasks.ValueTask System.IAsyncDisposable.DisposeAsync(); bool TryGetStateMachine(string name, out IStateMachine? stateMachine); System.Threading.Tasks.ValueTask WriteStateAsync(System.Threading.CancellationToken cancellationToken = default); @@ -424,6 +432,25 @@ public sealed partial class JournaledStateManagerOptions public System.TimeSpan RetirementGracePeriod { get { throw null; } set { } } } + public enum JournaledStateOperation + { + Write = 0, + Snapshot = 1, + Delete = 2 + } + + [GenerateSerializer] + public sealed partial class JournaledStatePostCommitException : System.Exception + { + public JournaledStatePostCommitException(JournaledStateOperation operation, System.Exception innerException) { } + } + + [GenerateSerializer] + public sealed partial class JournaledStatePreCommitException : System.Exception + { + public JournaledStatePreCommitException(JournaledStateOperation operation, System.Exception innerException) { } + } + public readonly ref partial struct JournalEntry { private readonly object _dummy; @@ -800,10 +827,62 @@ public void WriteField(ref global::Orleans.Serialization.Buffers. where TBufferWriter : System.Buffers.IBufferWriter { } } + [System.CodeDom.Compiler.GeneratedCode("OrleansCodeGen", "10.0.0.0")] + [System.ComponentModel.EditorBrowsable(System.ComponentModel.EditorBrowsableState.Never)] + [System.Diagnostics.CodeAnalysis.ExcludeFromCodeCoverage] + public sealed partial class Codec_JournaledStatePostCommitException : global::Orleans.Serialization.Codecs.IFieldCodec, global::Orleans.Serialization.Codecs.IFieldCodec + { + public Codec_JournaledStatePostCommitException(global::Orleans.Serialization.Serializers.IBaseCodec _baseTypeSerializer, global::Orleans.Serialization.Activators.IActivator _activator) { } + + public void Deserialize(ref global::Orleans.Serialization.Buffers.Reader reader, global::Orleans.Journaling.JournaledStatePostCommitException instance) { } + + public global::Orleans.Journaling.JournaledStatePostCommitException ReadValue(ref global::Orleans.Serialization.Buffers.Reader reader, global::Orleans.Serialization.WireProtocol.Field field) { throw null; } + + public void Serialize(ref global::Orleans.Serialization.Buffers.Writer writer, global::Orleans.Journaling.JournaledStatePostCommitException instance) + where TBufferWriter : System.Buffers.IBufferWriter { } + + public void WriteField(ref global::Orleans.Serialization.Buffers.Writer writer, uint fieldIdDelta, System.Type expectedType, global::Orleans.Journaling.JournaledStatePostCommitException value) + where TBufferWriter : System.Buffers.IBufferWriter { } + } + + [System.CodeDom.Compiler.GeneratedCode("OrleansCodeGen", "10.0.0.0")] + [System.ComponentModel.EditorBrowsable(System.ComponentModel.EditorBrowsableState.Never)] + [System.Diagnostics.CodeAnalysis.ExcludeFromCodeCoverage] + public sealed partial class Codec_JournaledStatePreCommitException : global::Orleans.Serialization.Codecs.IFieldCodec, global::Orleans.Serialization.Codecs.IFieldCodec + { + public Codec_JournaledStatePreCommitException(global::Orleans.Serialization.Serializers.IBaseCodec _baseTypeSerializer, global::Orleans.Serialization.Activators.IActivator _activator) { } + + public void Deserialize(ref global::Orleans.Serialization.Buffers.Reader reader, global::Orleans.Journaling.JournaledStatePreCommitException instance) { } + + public global::Orleans.Journaling.JournaledStatePreCommitException ReadValue(ref global::Orleans.Serialization.Buffers.Reader reader, global::Orleans.Serialization.WireProtocol.Field field) { throw null; } + + public void Serialize(ref global::Orleans.Serialization.Buffers.Writer writer, global::Orleans.Journaling.JournaledStatePreCommitException instance) + where TBufferWriter : System.Buffers.IBufferWriter { } + + public void WriteField(ref global::Orleans.Serialization.Buffers.Writer writer, uint fieldIdDelta, System.Type expectedType, global::Orleans.Journaling.JournaledStatePreCommitException value) + where TBufferWriter : System.Buffers.IBufferWriter { } + } + [System.CodeDom.Compiler.GeneratedCode("OrleansCodeGen", "10.0.0.0")] [System.ComponentModel.EditorBrowsable(System.ComponentModel.EditorBrowsableState.Never)] [System.Diagnostics.CodeAnalysis.ExcludeFromCodeCoverage] public sealed partial class Copier_DurableTaskCompletionSourceState : global::Orleans.Serialization.Cloning.ShallowCopier> { } + + [System.CodeDom.Compiler.GeneratedCode("OrleansCodeGen", "10.0.0.0")] + [System.ComponentModel.EditorBrowsable(System.ComponentModel.EditorBrowsableState.Never)] + [System.Diagnostics.CodeAnalysis.ExcludeFromCodeCoverage] + public sealed partial class Copier_JournaledStatePostCommitException : global::Orleans.Serialization.GeneratedCodeHelpers.OrleansGeneratedCodeHelper.ExceptionCopier + { + public Copier_JournaledStatePostCommitException(global::Orleans.Serialization.Serializers.ICodecProvider codecProvider) : base(default(Serialization.Serializers.ICodecProvider)!) { } + } + + [System.CodeDom.Compiler.GeneratedCode("OrleansCodeGen", "10.0.0.0")] + [System.ComponentModel.EditorBrowsable(System.ComponentModel.EditorBrowsableState.Never)] + [System.Diagnostics.CodeAnalysis.ExcludeFromCodeCoverage] + public sealed partial class Copier_JournaledStatePreCommitException : global::Orleans.Serialization.GeneratedCodeHelpers.OrleansGeneratedCodeHelper.ExceptionCopier + { + public Copier_JournaledStatePreCommitException(global::Orleans.Serialization.Serializers.ICodecProvider codecProvider) : base(default(Serialization.Serializers.ICodecProvider)!) { } + } } \ No newline at end of file diff --git a/test/Benchmarks/Journaling/DurableListJournalBenchmarks.cs b/test/Benchmarks/Journaling/DurableListJournalBenchmarks.cs index c94bd6045e2..5cb36fbffcf 100644 --- a/test/Benchmarks/Journaling/DurableListJournalBenchmarks.cs +++ b/test/Benchmarks/Journaling/DurableListJournalBenchmarks.cs @@ -126,6 +126,23 @@ private void ReplayEncodedJournalData() private sealed class BenchmarkJournalManager(OrleansBinaryJournalBufferWriter buffer, JournalStreamId streamId) : IJournaledStateManager { + private readonly List _hooks = []; + + public ValueTask AddHookAsync(IJournaledStateHook hook, int? index = null, CancellationToken cancellationToken = default) + { + _hooks.Insert(index ?? _hooks.Count, hook); + return default; + } + + public ValueTask RemoveHookAsync(IJournaledStateHook hook, CancellationToken cancellationToken = default) + { + if (!_hooks.Remove(hook)) + { + throw new InvalidOperationException("The journal operation hook is not registered."); + } + return default; + } + public ValueTask InitializeAsync(CancellationToken cancellationToken) => default; public void RegisterStateMachine(string name, IStateMachine state) => state.Reset(buffer.CreateJournalStreamWriter(streamId)); diff --git a/test/Orleans.DurableJobs.Tests/DurableJobs/JournaledJobShardDiscoveryTests.cs b/test/Orleans.DurableJobs.Tests/DurableJobs/JournaledJobShardDiscoveryTests.cs index 7a9862c5cab..39e4213c5c1 100644 --- a/test/Orleans.DurableJobs.Tests/DurableJobs/JournaledJobShardDiscoveryTests.cs +++ b/test/Orleans.DurableJobs.Tests/DurableJobs/JournaledJobShardDiscoveryTests.cs @@ -727,6 +727,10 @@ public IJournaledStateManager CreateStandalone(JournalId journalId) private sealed class TrackingJournaledStateManager(IJournaledStateManager inner) : IJournaledStateManager { + public ValueTask AddHookAsync(IJournaledStateHook hook, int? index = null, CancellationToken cancellationToken = default) => + inner.AddHookAsync(hook, index, cancellationToken); + public ValueTask RemoveHookAsync(IJournaledStateHook hook, CancellationToken cancellationToken = default) => + inner.RemoveHookAsync(hook, cancellationToken); public int DisposeCalls { get; private set; } public long PendingWriteByteCount => inner.PendingWriteByteCount; public ValueTask InitializeAsync(CancellationToken cancellationToken) => inner.InitializeAsync(cancellationToken); diff --git a/test/Orleans.Journaling.Json.Tests/CodecRecoveryTests.cs b/test/Orleans.Journaling.Json.Tests/CodecRecoveryTests.cs index f1cd56ccf37..5217532cad1 100644 --- a/test/Orleans.Journaling.Json.Tests/CodecRecoveryTests.cs +++ b/test/Orleans.Journaling.Json.Tests/CodecRecoveryTests.cs @@ -293,27 +293,31 @@ public async Task Recovery_EmptyJournalWithStaleMetadata_WritesConfiguredFormat( Assert.Contains("""[8,["set","alpha",1]]""", Encoding.UTF8.GetString(storage.Segments.Single()), StringComparison.Ordinal); } - [Fact] - public async Task Migration_WithUnregisteredRetiredState_ThrowsClearError() + [Theory] + [InlineData(OrleansBinaryJournalFormat.JournalFormatKey, JsonLinesJournalFormat.JournalFormatKey)] + [InlineData(JsonLinesJournalFormat.JournalFormatKey, OrleansBinaryJournalFormat.JournalFormatKey)] + public async Task Migration_WithUnregisteredRetiredState_PreservesStorage(string sourceFormat, string targetFormat) { - var storage = new VolatileJournalStorage(OrleansBinaryJournalFormat.JournalFormatKey); - using var first = CreateFormatAwareTestSystem(storage, OrleansBinaryJournalFormat.JournalFormatKey); - var dict = CreateFormatAwareDictionary(first, OrleansBinaryJournalFormat.JournalFormatKey, "dict"); + var storage = new VolatileJournalStorage(sourceFormat); + using var first = CreateFormatAwareTestSystem(storage, sourceFormat); + var dict = CreateFormatAwareDictionary(first, sourceFormat, "dict"); await first.Lifecycle.OnStart(TestContext.Current.CancellationToken); dict.Add("alpha", 1); await first.Manager.WriteStateAsync(CancellationToken.None); - storage.SetConfiguredJournalFormatKey(JsonLinesJournalFormat.JournalFormatKey); - using var recovered = CreateFormatAwareTestSystem(storage, JsonLinesJournalFormat.JournalFormatKey); - var other = CreateFormatAwareDictionary(recovered, JsonLinesJournalFormat.JournalFormatKey, "other"); + var originalBytes = Assert.Single(storage.Segments).ToArray(); + storage.SetConfiguredJournalFormatKey(targetFormat); + using var recovered = CreateFormatAwareTestSystem(storage, targetFormat); + var other = CreateFormatAwareDictionary(recovered, targetFormat, "other"); await recovered.Lifecycle.OnStart(TestContext.Current.CancellationToken); other.Add("beta", 2); var exception = await Assert.ThrowsAsync( () => recovered.Manager.WriteStateAsync(CancellationToken.None).AsTask()); - Assert.Contains("Cannot migrate journal", exception.Message, StringComparison.Ordinal); - Assert.Contains("not currently registered", exception.Message, StringComparison.Ordinal); + Assert.Contains("cannot append preserved entry", exception.Message, StringComparison.Ordinal); + Assert.Equal(sourceFormat, storage.StoredJournalFormatKey); + Assert.Equal(originalBytes, Assert.Single(storage.Segments)); } [Theory] diff --git a/test/Orleans.Journaling.Tests/DurableCollectionDirectWriteTests.cs b/test/Orleans.Journaling.Tests/DurableCollectionDirectWriteTests.cs index 26cf16aac46..ede9edb2f56 100644 --- a/test/Orleans.Journaling.Tests/DurableCollectionDirectWriteTests.cs +++ b/test/Orleans.Journaling.Tests/DurableCollectionDirectWriteTests.cs @@ -193,6 +193,23 @@ public void Collections_UseDirectEntryWriter() private sealed class TestJournalManager(TestJournalStreamWriter writer) : IJournaledStateManager { + private readonly List _hooks = []; + + public ValueTask AddHookAsync(IJournaledStateHook hook, int? index = null, CancellationToken cancellationToken = default) + { + _hooks.Insert(index ?? _hooks.Count, hook); + return default; + } + + public ValueTask RemoveHookAsync(IJournaledStateHook hook, CancellationToken cancellationToken = default) + { + if (!_hooks.Remove(hook)) + { + throw new InvalidOperationException("The journal operation hook is not registered."); + } + return default; + } + public ValueTask InitializeAsync(CancellationToken cancellationToken) => default; public void RegisterStateMachine(string name, IStateMachine state) => state.Reset(writer.CreateWriter()); diff --git a/test/Orleans.Journaling.Tests/DurableListDirectWriteTests.cs b/test/Orleans.Journaling.Tests/DurableListDirectWriteTests.cs index e1cc3708672..245561f8c50 100644 --- a/test/Orleans.Journaling.Tests/DurableListDirectWriteTests.cs +++ b/test/Orleans.Journaling.Tests/DurableListDirectWriteTests.cs @@ -92,6 +92,23 @@ public void Add_UsesDirectEntryWriter() private sealed class TestJournalManager(TestJournalStreamWriter writer) : IJournaledStateManager { + private readonly List _hooks = []; + + public ValueTask AddHookAsync(IJournaledStateHook hook, int? index = null, CancellationToken cancellationToken = default) + { + _hooks.Insert(index ?? _hooks.Count, hook); + return default; + } + + public ValueTask RemoveHookAsync(IJournaledStateHook hook, CancellationToken cancellationToken = default) + { + if (!_hooks.Remove(hook)) + { + throw new InvalidOperationException("The journal operation hook is not registered."); + } + return default; + } + public ValueTask InitializeAsync(CancellationToken cancellationToken) => default; public void RegisterStateMachine(string name, IStateMachine state) => state.Reset(writer.CreateWriter()); diff --git a/test/Orleans.Journaling.Tests/DurableStateManagerTests.cs b/test/Orleans.Journaling.Tests/DurableStateManagerTests.cs index d1b0b4b797f..e85899b2bdb 100644 --- a/test/Orleans.Journaling.Tests/DurableStateManagerTests.cs +++ b/test/Orleans.Journaling.Tests/DurableStateManagerTests.cs @@ -592,7 +592,7 @@ private static void AssertIncompatible( private static void AssertAdmissionRejected(Action action) { var error = Assert.Throws(action); - Assert.Equal("New states cannot be registered after journaled state manager initialization has begun.", error.Message); + Assert.Equal("Journal registrations cannot be changed after journaled state manager initialization has begun.", error.Message); } private static void AssertFenced(InvalidOperationException error, IOException cause) diff --git a/test/Orleans.Journaling.Tests/KeyedJournalingRegistrationTests.cs b/test/Orleans.Journaling.Tests/KeyedJournalingRegistrationTests.cs index e9822dd9557..dd7945ec7cb 100644 --- a/test/Orleans.Journaling.Tests/KeyedJournalingRegistrationTests.cs +++ b/test/Orleans.Journaling.Tests/KeyedJournalingRegistrationTests.cs @@ -746,6 +746,10 @@ public void ReplayEntry(JournalEntry entry, JournalReplayContext context) => private sealed class DelegatingStateManager(IJournaledStateManager inner) : IJournaledStateManager { + public ValueTask AddHookAsync(IJournaledStateHook hook, int? index = null, CancellationToken cancellationToken = default) => + inner.AddHookAsync(hook, index, cancellationToken); + public ValueTask RemoveHookAsync(IJournaledStateHook hook, CancellationToken cancellationToken = default) => + inner.RemoveHookAsync(hook, cancellationToken); 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) diff --git a/test/Orleans.Journaling.Tests/StateManagerHookTests.cs b/test/Orleans.Journaling.Tests/StateManagerHookTests.cs new file mode 100644 index 00000000000..8833f87cb85 --- /dev/null +++ b/test/Orleans.Journaling.Tests/StateManagerHookTests.cs @@ -0,0 +1,1566 @@ +using System.Buffers; +using System.Collections.Concurrent; +using System.Diagnostics; +using System.Reflection; +using System.Threading.Channels; +using Microsoft.Extensions.DependencyInjection; +using NSubstitute; +using Orleans.Diagnostics; +using Orleans.Serialization; +using Xunit; + +namespace Orleans.Journaling.Tests; + +public partial class StateManagerTests +{ + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task Hooks_QueueActivityUsesSelectedOperationAndCompletesWithStorageOutcome(bool snapshot, bool storageFails) + { + using var parent = new Activity("journal queue test").SetIdFormat(ActivityIdFormat.W3C).Start(); + var stopped = new ConcurrentBag(); + using var listener = new ActivityListener + { + ShouldListenTo = source => source.Name == ActivitySources.StorageActivitySourceName, + Sample = (ref ActivityCreationOptions _) => ActivitySamplingResult.AllDataAndRecorded, + ActivityStopped = activity => + { + if (activity.ParentSpanId == parent.SpanId) + { + stopped.Add(activity); + } + } + }; + ActivitySource.AddActivityListener(listener); + var failure = new IOException("Storage failed."); + var storage = new CapturingStorage + { + IsCompactionRequested = snapshot, + NextAppendException = storageFails && !snapshot ? failure : null, + NextReplaceException = storageFails && snapshot ? failure : null + }; + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableValue("state", manager, CreateValueCodec()); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state.Value = 1; + if (storageFails) + { + Assert.Same(failure, await Assert.ThrowsAsync(() => + manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask())); + } + else + { + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + } + + await manager.DisposeAsync(); + var activity = Assert.Single(stopped); + var operation = snapshot ? "snapshot" : "append"; + Assert.Equal($"journal {operation}", activity.OperationName); + Assert.Equal(operation, activity.GetTagItem(ActivityTagKeys.JournalStorageOperation)); + Assert.Equal(storageFails ? ActivityStatusCode.Error : ActivityStatusCode.Ok, activity.Status); + Assert.Equal(storageFails ? failure.Message : null, activity.StatusDescription); + } + + [Fact] + public async Task Hooks_IdleRegistrationAppliesSynchronously() + { + await using var manager = CreateTestSystem().Manager; + _ = new DurableValue("state", manager, CreateValueCodec()); + var field = typeof(JournaledStateManager).GetField("_hooks", BindingFlags.Instance | BindingFlags.NonPublic)!; + Assert.Null(field.GetValue(manager)); + + var before = 0; + var hook = new DelegateTestHook { BeforeOperation = (_, _) => before++ }; + Assert.True(manager.AddHookAsync(hook, cancellationToken: TestContext.Current.CancellationToken).IsCompletedSuccessfully); + Assert.True(manager.RemoveHookAsync(hook, cancellationToken: TestContext.Current.CancellationToken).IsCompletedSuccessfully); + await Assert.ThrowsAsync(() => manager.AddHookAsync(null!, cancellationToken: TestContext.Current.CancellationToken).AsTask()); + await Assert.ThrowsAsync(() => manager.RemoveHookAsync(hook, cancellationToken: TestContext.Current.CancellationToken).AsTask()); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Equal(0, before); + Assert.True(manager.AddHookAsync(hook, index: 0, cancellationToken: TestContext.Current.CancellationToken).IsCompletedSuccessfully); + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Equal(1, before); + Assert.True(manager.RemoveHookAsync(hook, cancellationToken: TestContext.Current.CancellationToken).IsCompletedSuccessfully); + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Equal(1, before); + + await using var withoutHooks = CreateTestSystem().Manager; + await withoutHooks.InitializeAsync(TestContext.Current.CancellationToken); + await withoutHooks.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Null(field.GetValue(withoutHooks)); + } + + [Fact] + public async Task Hooks_ShutdownRejectsRegistrationAndRemoval() + { + var manager = CreateTestSystem().Manager; + var hook = new DelegateTestHook(); + await manager.AddHookAsync(hook, cancellationToken: TestContext.Current.CancellationToken); + + await manager.DisposeAsync(); + + await Assert.ThrowsAsync(() => manager.AddHookAsync(hook, cancellationToken: TestContext.Current.CancellationToken).AsTask()); + await Assert.ThrowsAsync(() => manager.RemoveHookAsync(hook, cancellationToken: TestContext.Current.CancellationToken).AsTask()); + } + + [Fact] + public async Task Hooks_RegistrationDuringRecoveryAppliesBeforeLaterPersistence() + { + var entered = NewSignal(); + var release = NewSignal(); + var storage = Substitute.For(); + storage.ReadAsync(Arg.Any(), Arg.Any()) + .Returns(call => ReadAsync(call.Arg(), call.Arg())); + await using var manager = CreateTestSystem(storage).Manager; + var before = 0; + var after = 0; + var hook = new DelegateTestHook + { + BeforeOperation = (_, _) => before++, + AfterOperation = (_, _) => after++ + }; + var initialization = manager.InitializeAsync(TestContext.Current.CancellationToken).AsTask(); + var registration = manager.AddHookAsync(hook, cancellationToken: TestContext.Current.CancellationToken).AsTask(); + try + { + await WaitFor(entered.Task); + Assert.False(initialization.IsCompleted); + Assert.False(registration.IsCompleted); + release.TrySetResult(); + await WaitFor(Task.WhenAll(initialization, registration)); + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Equal(1, before); + Assert.Equal(1, after); + } + finally + { + release.TrySetResult(); + } + + async ValueTask ReadAsync(IJournalStorageConsumer consumer, CancellationToken token) + { + entered.TrySetResult(); + await release.Task.WaitAsync(token); + consumer.Complete(metadata: null); + } + } + + [Fact] + public async Task Hooks_QueuedRegistrationReportsRecoveryFailureAndIdleRegistrationDoesNotRetryRecovery() + { + var entered = NewSignal(); + var release = NewSignal(); + var storage = Substitute.For(); + var failure = new IOException("Recovery failed."); + var fail = true; + var reads = 0; + storage.ReadAsync(Arg.Any(), Arg.Any()) + .Returns(call => ReadAsync(call.Arg(), call.Arg())); + await using var manager = CreateTestSystem(storage).Manager; + var initialization = manager.InitializeAsync(TestContext.Current.CancellationToken).AsTask(); + var failedBefore = 0; + var failedRegistration = manager.AddHookAsync(new DelegateTestHook + { + BeforeOperation = (_, _) => failedBefore++ + }, cancellationToken: TestContext.Current.CancellationToken).AsTask(); + try + { + await WaitFor(entered.Task); + release.TrySetResult(); + Assert.Same(failure, await Assert.ThrowsAsync(() => WaitFor(initialization))); + Assert.Same(failure, await Assert.ThrowsAsync(() => WaitFor(failedRegistration))); + var before = 0; + var hook = new DelegateTestHook { BeforeOperation = (_, _) => before++ }; + Assert.True(manager.AddHookAsync(hook, cancellationToken: TestContext.Current.CancellationToken).IsCompletedSuccessfully); + Assert.Equal(1, reads); + fail = false; + await manager.InitializeAsync(TestContext.Current.CancellationToken); + Assert.Equal(2, reads); + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Equal(0, failedBefore); + Assert.Equal(1, before); + } + finally + { + release.TrySetResult(); + } + + async ValueTask ReadAsync(IJournalStorageConsumer consumer, CancellationToken token) + { + reads++; + entered.TrySetResult(); + await release.Task.WaitAsync(token); + if (fail) + { + throw failure; + } + consumer.Complete(metadata: null); + } + } + + [Theory] + [InlineData("before")] + [InlineData("storage")] + [InlineData("after")] + public async Task Hooks_BusyMutationWaitsForCompleteOperationAndSurvivesCallerCancellation(string phase) + { + var storage = new CapturingStorage { BlockNextAppend = phase == "storage" }; + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableValue("state", manager, CreateValueCodec()); + var entered = NewSignal(); + var release = NewSignal(); + var oldBefore = 0; + var oldAfter = 0; + var newBefore = 0; + var newAfter = 0; + var original = new DelegateTestHook + { + BeforeOperationAsync = async (_, _) => + { + oldBefore++; + if (phase == "before") + { + entered.TrySetResult(); + await release.Task; + } + }, + AfterOperationAsync = async (_, _) => + { + oldAfter++; + if (phase == "after") + { + entered.TrySetResult(); + await release.Task; + } + } + }; + var replacement = new DelegateTestHook + { + BeforeOperation = (_, _) => newBefore++, + AfterOperation = (_, _) => newAfter++ + }; + await manager.AddHookAsync(original, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state.Value = 1; + var first = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await WaitFor(phase == "storage" ? storage.BlockedAppendStarted.Task : entered.Task); + using var caller = new CancellationTokenSource(); + var added = manager.AddHookAsync(replacement, index: 0, caller.Token).AsTask(); + var removed = manager.RemoveHookAsync(original, cancellationToken: TestContext.Current.CancellationToken).AsTask(); + state.Value = 2; + var next = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + Assert.False(added.IsCompleted); + Assert.False(removed.IsCompleted); + Assert.Equal(0, newBefore); + Assert.Equal(0, newAfter); + caller.Cancel(); + await Assert.ThrowsAnyAsync(() => added); + + release.TrySetResult(); + storage.ReleaseAppend.TrySetResult(); + await WaitFor(Task.WhenAll(first, removed, next)); + Assert.Equal(1, oldBefore); + Assert.Equal(1, oldAfter); + Assert.Equal(1, newBefore); + Assert.Equal(1, newAfter); + Assert.Equal(2, state.Value); + } + finally + { + release.TrySetResult(); + storage.ReleaseAppend.TrySetResult(); + } + } + + [Fact] + public async Task Hooks_QueuedMutationSeparatesWriteCoalescingGroups() + { + var storage = new CapturingStorage { BlockNextAppend = true }; + await using var manager = CreateTestSystem(storage).Manager; + var state = new LifecycleState(); + manager.RegisterStateMachine("state", state); + var oldBefore = 0; + var oldAfter = 0; + var newBefore = 0; + var newAfter = 0; + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperation = (_, _) => oldBefore++, + AfterOperation = (_, _) => oldAfter++ + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + var first = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await WaitFor(storage.BlockedAppendStarted.Task); + var beforeChange = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + var change = manager.AddHookAsync(new DelegateTestHook + { + BeforeOperation = (_, _) => newBefore++, + AfterOperation = (_, _) => newAfter++ + }, cancellationToken: TestContext.Current.CancellationToken).AsTask(); + var afterChange = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + var coalesced = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + storage.ReleaseAppend.TrySetResult(); + await WaitFor(Task.WhenAll(first, beforeChange, change, afterChange, coalesced)); + Assert.Equal(3, state.CaptureCount); + Assert.Equal(3, storage.Appends.Count); + Assert.Equal(3, oldBefore); + Assert.Equal(3, oldAfter); + Assert.Equal(1, newBefore); + Assert.Equal(1, newAfter); + } + finally + { + storage.ReleaseAppend.TrySetResult(); + } + } + + [Fact] + public async Task Hooks_ResetCanQueueRegistrationWithoutChangingCurrentAfterHooks() + { + var storage = new CapturingStorage(); + await using var manager = CreateTestSystem(storage).Manager; + var state = new LifecycleState(); + Task? registration = null; + var newBefore = 0; + var newAfter = 0; + state.ResetAction = () => + { + if (state.ResetCount == 2) + { + registration = manager.AddHookAsync(new DelegateTestHook + { + BeforeOperation = (_, _) => newBefore++, + AfterOperation = (_, _) => newAfter++ + }, cancellationToken: TestContext.Current.CancellationToken).AsTask(); + } + }; + manager.RegisterStateMachine("state", state); + await manager.AddHookAsync(new DelegateTestHook + { + AfterOperation = (operation, _) => + { + if (operation == JournaledStateOperation.Delete) + { + Assert.NotNull(registration); + Assert.False(registration.IsCompleted); + Assert.Equal(0, newBefore); + Assert.Equal(0, newAfter); + } + } + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + await manager.DeleteStateAsync(TestContext.Current.CancellationToken); + Assert.NotNull(registration); + await WaitFor(registration); + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Equal(1, newBefore); + Assert.Equal(1, newAfter); + } + + [Fact] + public async Task Hooks_FailedQueuedMutationReportsErrorAndLeavesOwnerUsable() + { + var storage = new CapturingStorage { BlockNextAppend = true }; + await using var manager = CreateTestSystem(storage).Manager; + manager.RegisterStateMachine("state", new LifecycleState()); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + var first = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await WaitFor(storage.BlockedAppendStarted.Task); + var invalid = manager.AddHookAsync(new DelegateTestHook(), index: 1, cancellationToken: TestContext.Current.CancellationToken).AsTask(); + var before = 0; + var valid = new DelegateTestHook { BeforeOperation = (_, _) => before++ }; + var accepted = manager.AddHookAsync(valid, cancellationToken: TestContext.Current.CancellationToken).AsTask(); + storage.ReleaseAppend.TrySetResult(); + await Assert.ThrowsAsync(() => WaitFor(invalid)); + await WaitFor(Task.WhenAll(first, accepted)); + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Equal(1, before); + } + finally + { + storage.ReleaseAppend.TrySetResult(); + } + } + + [Fact] + public async Task Hooks_StorageFailureFaultsQueuedRegistrationWithOriginalCause() + { + var failure = new IOException("Storage failed."); + var storage = new CapturingStorage { BlockNextAppend = true, NextAppendException = failure }; + await using var manager = CreateTestSystem(storage).Manager; + manager.RegisterStateMachine("state", new LifecycleState()); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + var write = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await WaitFor(storage.BlockedAppendStarted.Task); + var registration = manager.AddHookAsync(new DelegateTestHook(), cancellationToken: TestContext.Current.CancellationToken).AsTask(); + storage.ReleaseAppend.TrySetResult(); + Assert.Same(failure, await Assert.ThrowsAsync(() => WaitFor(write))); + Assert.Same(failure, await Assert.ThrowsAsync(() => WaitFor(registration))); + var fenced = await Assert.ThrowsAsync(() => + manager.AddHookAsync(new DelegateTestHook(), cancellationToken: TestContext.Current.CancellationToken).AsTask()); + Assert.Same(failure, fenced.InnerException); + } + finally + { + storage.ReleaseAppend.TrySetResult(); + } + } + + [Fact] + public async Task Hooks_ShutdownCancelsQueuedMutationWithoutChangingCommittedOperation() + { + var storage = new CapturingStorage(); + await using var manager = CreateTestSystem(storage).Manager; + manager.RegisterStateMachine("state", new LifecycleState()); + var entered = NewSignal(); + var release = NewSignal(); + var original = new DelegateTestHook + { + AfterOperationAsync = async (_, _) => + { + entered.TrySetResult(); + await release.Task; + } + }; + await manager.AddHookAsync(original, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + var write = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await WaitFor(entered.Task); + var added = manager.AddHookAsync(new DelegateTestHook(), cancellationToken: TestContext.Current.CancellationToken).AsTask(); + var removed = manager.RemoveHookAsync(original, cancellationToken: TestContext.Current.CancellationToken).AsTask(); + var shutdown = manager.DisposeAsync().AsTask(); + Assert.False(shutdown.IsCompleted); + release.TrySetResult(); + await WaitFor(Task.WhenAll(write, shutdown)); + await Assert.ThrowsAnyAsync(() => WaitFor(added)); + await Assert.ThrowsAnyAsync(() => WaitFor(removed)); + Assert.Single(storage.Appends); + } + finally + { + release.TrySetResult(); + } + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Hooks_DelegateOrderSurroundsCaptureAndActualAcknowledgement(bool snapshot) + { + var storage = new CapturingStorage { IsCompactionRequested = snapshot, BlockNextAppend = !snapshot, BlockNextReplace = snapshot }; + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableDictionary("state", manager, CreateDictionaryCodec()); + List events = []; + manager.RegisterStateMachine("ack", new HookAcknowledgementState(events)); + var expected = snapshot ? JournaledStateOperation.Snapshot : JournaledStateOperation.Write; + CancellationToken ownedToken = default; + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperation = (operation, token) => + { + Assert.Equal(expected, operation); + Assert.True(token.CanBeCanceled); + ownedToken = token; + events.Add("before-sync"); + state["hook"] = 2; + }, + BeforeOperationAsync = async (operation, token) => + { + Assert.Equal(expected, operation); + Assert.Equal(ownedToken, token); + await Task.Yield(); + events.Add("before-async"); + }, + AfterOperation = (operation, token) => + { + Assert.Equal(expected, operation); + Assert.Equal(ownedToken, token); + Assert.Single(snapshot ? storage.Replaces : storage.Appends); + events.Add("after-sync"); + }, + AfterOperationAsync = async (_, token) => + { + Assert.Equal(ownedToken, token); + await Task.Yield(); + events.Add("after-async"); + } + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state["business"] = 1; + var write = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await (snapshot ? storage.ReplaceEntered : storage.BlockedAppendStarted).Task + .WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.Equal(new[] { "before-sync", "before-async", "capture" }, events); + Assert.False(write.IsCompleted); + (snapshot ? storage.ReleaseReplace : storage.ReleaseAppend).TrySetResult(); + await write.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.Equal(new[] { "before-sync", "before-async", "capture", "ack", "after-sync", "after-async" }, events); + } + finally + { + (snapshot ? storage.ReleaseReplace : storage.ReleaseAppend).TrySetResult(); + } + + await using var replayManager = CreateTestSystem(storage).Manager; + var recovered = new DurableDictionary("state", replayManager, CreateDictionaryCodec()); + await replayManager.InitializeAsync(TestContext.Current.CancellationToken); + Assert.Equal(2, recovered.Count); + Assert.Equal(1, recovered["business"]); + Assert.Equal(2, recovered["hook"]); + } + + [Theory] + [InlineData(JournaledStateOperation.Write)] + [InlineData(JournaledStateOperation.Snapshot)] + [InlineData(JournaledStateOperation.Delete)] + public async Task Hooks_BeforeFailurePreservesStateForExplicitRetry(JournaledStateOperation operation) + { + var storage = new CapturingStorage { IsCompactionRequested = operation == JournaledStateOperation.Snapshot }; + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableDictionary("state", manager, CreateDictionaryCodec()); + var capture = new AlwaysWritingState(); + manager.RegisterStateMachine("capture", capture); + var failure = new IOException("Prerequisite failed."); + var fail = true; + var after = 0; + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperation = (_, _) => + { + if (fail) + { + throw failure; + } + }, + AfterOperation = (_, _) => after++ + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state["business"] = 1; + var pendingBytes = manager.PendingWriteByteCount; + var caught = await Assert.ThrowsAsync(() => InvokeAsync()); + Assert.Same(failure, caught.InnerException); + Assert.Contains($"'{operation}'", caught.Message, StringComparison.Ordinal); + Assert.Equal(1, state["business"]); + Assert.Equal(pendingBytes, manager.PendingWriteByteCount); + Assert.Equal(0, capture.AppendEntriesCount); + Assert.Equal(0, capture.WriteCompletedCount); + Assert.Empty(storage.Appends); + Assert.Empty(storage.Replaces); + Assert.Equal(0, storage.DeleteCount); + Assert.Equal(0, after); + storage.IsCompactionRequested = false; + fail = false; + await InvokeAsync(); + Assert.Equal(1, after); + Assert.Equal(operation == JournaledStateOperation.Delete ? 0 : 1, state.Count); + Assert.Equal(operation == JournaledStateOperation.Delete ? 1 : 0, storage.DeleteCount); + Assert.Equal(operation == JournaledStateOperation.Write ? 1 : 0, storage.Appends.Count); + Assert.Equal(operation == JournaledStateOperation.Snapshot ? 1 : 0, storage.Replaces.Count); + + Task InvokeAsync() => (operation == JournaledStateOperation.Delete + ? manager.DeleteStateAsync(TestContext.Current.CancellationToken) + : manager.WriteStateAsync(TestContext.Current.CancellationToken)).AsTask(); + } + + [Theory] + [InlineData(JournaledStateOperation.Write)] + [InlineData(JournaledStateOperation.Snapshot)] + [InlineData(JournaledStateOperation.Delete)] + public async Task Hooks_AfterFailureReportsCompletedOperationAndRunsRemainingHooks(JournaledStateOperation operation) + { + var storage = new CapturingStorage { IsCompactionRequested = operation == JournaledStateOperation.Snapshot }; + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableDictionary("state", manager, CreateDictionaryCodec()); + List events = []; + manager.RegisterStateMachine("ack", new HookAcknowledgementState(events)); + var first = new IOException("First cleanup failed."); + var second = new IOException("Second cleanup failed."); + var fail = true; + await manager.AddHookAsync(new DelegateTestHook + { + AfterOperation = (_, _) => + { + if (fail) + { + throw first; + } + } + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.AddHookAsync(new DelegateTestHook + { + AfterOperationAsync = async (_, _) => + { + await Task.Yield(); + events.Add("second"); + if (fail) + { + throw second; + } + } + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.AddHookAsync(new DelegateTestHook { AfterOperation = (_, _) => events.Add("last") }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state["business"] = 1; + var caught = await Assert.ThrowsAsync(() => (operation == JournaledStateOperation.Delete + ? manager.DeleteStateAsync(TestContext.Current.CancellationToken) + : manager.WriteStateAsync(TestContext.Current.CancellationToken)).AsTask()); + Assert.Contains($"'{operation}'", caught.Message, StringComparison.Ordinal); + Assert.Equal(new Exception[] { first, second }, Assert.IsType(caught.InnerException).InnerExceptions); + Assert.Equal(operation == JournaledStateOperation.Delete + ? new[] { "second", "last" } + : new[] { "capture", "ack", "second", "last" }, events); + Assert.Equal(operation == JournaledStateOperation.Delete ? 0 : 1, state.Count); + Assert.Equal(operation == JournaledStateOperation.Delete ? 1 : 0, storage.DeleteCount); + fail = false; + + await using var replayManager = CreateTestSystem(storage).Manager; + var recovered = new DurableDictionary("state", replayManager, CreateDictionaryCodec()); + await replayManager.InitializeAsync(TestContext.Current.CancellationToken); + Assert.Equal(operation == JournaledStateOperation.Delete ? 0 : 1, recovered.Count); + if (operation != JournaledStateOperation.Delete) + { + Assert.Equal(1, recovered["business"]); + } + + state["later"] = 2; + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Equal(2, state["later"]); + } + + [Fact] + public async Task Hooks_StorageFailureSkipsAfterHooksAndPreservesTerminalCause() + { + var failure = new IOException("Storage failed."); + var storage = new CapturingStorage { NextAppendException = failure }; + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableDictionary("state", manager, CreateDictionaryCodec()); + var after = 0; + await manager.AddHookAsync(new DelegateTestHook { AfterOperation = (_, _) => after++ }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state["business"] = 1; + Assert.Same(failure, await Assert.ThrowsAsync(() => manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask())); + Assert.Equal(0, after); + var fenced = await Assert.ThrowsAsync(() => manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask()); + Assert.Same(failure, fenced.InnerException); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Hooks_FinalPrerequisiteCoversChangesArrivingDuringLaterOrdinaryHook(bool snapshot) + { + var storage = new CapturingStorage { IsCompactionRequested = snapshot }; + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableDictionary("state", manager, CreateDictionaryCodec()); + List events = []; + manager.RegisterStateMachine("ack", new HookAcknowledgementState(events)); + var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + await manager.AddHookAsync(new DelegateTestHook( + (_, _) => + { + Assert.Equal(3, state["late"]); + state["owner"] = 2; + events.Add("capture-prerequisite"); + return default; + }, + (_, _) => events.Add("capture-after")), cancellationToken: TestContext.Current.CancellationToken); + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperationAsync = async (_, token) => + { + events.Add("ordinary-before"); + entered.TrySetResult(); + await release.Task.WaitAsync(token); + events.Add("ordinary-complete"); + }, + AfterOperation = (_, _) => events.Add("ordinary-after") + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state["business"] = 1; + var write = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await entered.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.Equal(new[] { "ordinary-before" }, events); + Assert.Empty(storage.Appends); + Assert.Empty(storage.Replaces); + state["late"] = 3; + release.TrySetResult(); + await write.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.Equal(new[] { "ordinary-before", "ordinary-complete", "capture-prerequisite", "capture", "ack", "capture-after", "ordinary-after" }, events); + } + finally + { + release.TrySetResult(); + } + + await using var recoveredManager = CreateTestSystem(storage).Manager; + var recovered = new DurableDictionary("state", recoveredManager, CreateDictionaryCodec()); + await recoveredManager.InitializeAsync(TestContext.Current.CancellationToken); + Assert.Equal(3, recovered.Count); + Assert.Equal(1, recovered["business"]); + Assert.Equal(2, recovered["owner"]); + Assert.Equal(3, recovered["late"]); + } + + [Fact] + public async Task Hooks_CanceledCallerRetainsFinalPreparationAndPostCompletion() + { + var storage = new CapturingStorage(); + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableDictionary("state", manager, CreateDictionaryCodec()); + var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var completed = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + CancellationToken ownedToken = default; + Func before = async (_, token) => + { + ownedToken = token; + entered.TrySetResult(); + await release.Task.WaitAsync(token); + state["scheduled"] = 2; + }; + Action after = (_, token) => + { + Assert.Equal(ownedToken, token); + Assert.False(token.IsCancellationRequested); + Assert.Single(storage.Appends); + completed.TrySetResult(); + }; + await manager.AddHookAsync(new DelegateTestHook(before, after), cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + using var caller = new CancellationTokenSource(); + state["business"] = 1; + var write = manager.WriteStateAsync(caller.Token).AsTask(); + try + { + await entered.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.NotEqual(caller.Token, ownedToken); + state["late"] = 3; + caller.Cancel(); + await Assert.ThrowsAnyAsync(() => write); + Assert.False(ownedToken.IsCancellationRequested); + Assert.Empty(storage.Appends); + Assert.False(completed.Task.IsCompleted); + release.TrySetResult(); + await completed.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.Equal(2, state["scheduled"]); + await using var recoveredManager = CreateTestSystem(storage).Manager; + var recovered = new DurableDictionary("state", recoveredManager, CreateDictionaryCodec()); + await recoveredManager.InitializeAsync(TestContext.Current.CancellationToken); + Assert.Equal(3, recovered.Count); + Assert.Equal(1, recovered["business"]); + Assert.Equal(2, recovered["scheduled"]); + Assert.Equal(3, recovered["late"]); + } + finally + { + release.TrySetResult(); + } + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Hooks_SynchronousFinalPrerequisiteAndCaptureShareContinuationAfterAsyncEarlierHook(bool snapshot) + { + var scheduler = new QueuedHookTaskScheduler(); + await scheduler.Run(async () => + { + var storage = new CapturingStorage { IsCompactionRequested = snapshot }; + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableDictionary("state", manager, CreateDictionaryCodec()); + List events = []; + List prepared = []; + manager.RegisterStateMachine("ack", new HookAcknowledgementState(events)); + var queueMutation = true; + Task? mutation = null; + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperation = (_, _) => + { + events.Add("final-before"); + prepared.AddRange(state.Keys); + if (queueMutation) + { + Assert.Empty(state); + queueMutation = false; + mutation = scheduler.Queue(() => + { + state["intent"] = 1; + events.Add("queued-mutation"); + }); + } + } + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperationAsync = async (_, _) => + { + events.Add("ordinary-before"); + await Task.Yield(); + events.Add("ordinary-complete"); + } + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + + Assert.NotNull(mutation); + await mutation; + Assert.Empty(prepared); + Assert.Equal(1, state["intent"]); + await using (var recoveredManager = CreateTestSystem(storage).Manager) + { + var recovered = new DurableDictionary("state", recoveredManager, CreateDictionaryCodec()); + await recoveredManager.InitializeAsync(TestContext.Current.CancellationToken); + Assert.Empty(recovered); + } + + Assert.Equal( + ["ordinary-before", "ordinary-complete", "final-before", "capture", "ack", "queued-mutation"], + events); + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Equal(["intent"], prepared); + await using var replayManager = CreateTestSystem(storage).Manager; + var replayed = new DurableDictionary("state", replayManager, CreateDictionaryCodec()); + await replayManager.InitializeAsync(TestContext.Current.CancellationToken); + var intent = Assert.Single(replayed); + Assert.Equal("intent", intent.Key); + Assert.Equal(1, intent.Value); + }); + } + + [Fact] + public async Task Hooks_BeforeCallbacksRunInReverseAndAfterCallbacksRunInListOrder() + { + await using var manager = CreateTestSystem().Manager; + List events = []; + manager.RegisterStateMachine("ack", new HookAcknowledgementState(events)); + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperation = (_, _) => events.Add("second-before"), + AfterOperation = (_, _) => events.Add("second-after") + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperation = (_, _) => events.Add("third-before"), + AfterOperation = (_, _) => events.Add("third-after") + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperation = (_, _) => events.Add("first-before"), + AfterOperation = (_, _) => events.Add("first-after") + }, index: 0, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Equal( + ["third-before", "second-before", "first-before", "capture", "ack", "first-after", "second-after", "third-after"], + events); + } + + [Fact] + public async Task Hooks_DeleteWaitsForPrerequisitesAndCompletesAfterReset() + { + var storage = new CapturingStorage(); + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableValue("state", manager, CreateValueCodec()); + var deleting = false; + var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var before = 0; + var after = 0; + await manager.AddHookAsync(new DelegateTestHook( + async (operation, token) => + { + before++; + if (!deleting) + { + Assert.Equal(JournaledStateOperation.Write, operation); + return; + } + + Assert.Equal(JournaledStateOperation.Delete, operation); + Assert.Equal(1, state.Value); + entered.TrySetResult(); + await release.Task.WaitAsync(token); + }, + (operation, _) => + { + if (!deleting) + { + return; + } + + Assert.Equal(JournaledStateOperation.Delete, operation); + Assert.Equal(0, state.Value); + Assert.Equal(1, storage.DeleteCount); + after++; + }), cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state.Value = 1; + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + deleting = true; + var deletion = manager.DeleteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await entered.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.Equal(0, storage.DeleteCount); + Assert.Equal(0, after); + Assert.False(deletion.IsCompleted); + release.TrySetResult(); + await deletion.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.Equal(1, after); + deleting = false; + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Equal(3, before); + } + finally + { + release.TrySetResult(); + } + } + + [Fact] + public async Task Hooks_ShutdownCallbackFailureStillDrainsOwnedPreparation() + { + var storage = new CapturingStorage(); + var manager = CreateTestSystem(storage).Manager; + var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var drained = false; + var failure = new IOException("Cancellation callback failed."); + var hook = new DelegateTestHook + { + BeforeOperationAsync = async (_, token) => + { + using var registration = token.Register(() => throw failure); + entered.TrySetResult(); + await release.Task; + drained = true; + } + }; + await manager.AddHookAsync(hook, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + var write = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await entered.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + var queued = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + var shutdown = manager.DisposeAsync().AsTask(); + var concurrentShutdown = manager.DisposeAsync().AsTask(); + Assert.Same(shutdown, concurrentShutdown); + Assert.False(shutdown.IsCompleted); + Assert.False(drained); + release.TrySetResult(); + var caught = await Assert.ThrowsAsync(() => + shutdown.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken)); + Assert.Same(failure, Assert.Single(caught.InnerExceptions)); + await Assert.ThrowsAnyAsync(() => write); + await Assert.ThrowsAnyAsync(() => queued); + Assert.True(drained); + Assert.Empty(storage.Appends); + await Assert.ThrowsAsync(() => + manager.RemoveHookAsync(hook, cancellationToken: TestContext.Current.CancellationToken).AsTask()); + } + finally + { + release.TrySetResult(); + await Assert.ThrowsAsync(() => manager.DisposeAsync().AsTask()); + } + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Hooks_ShutdownCancelsQueuedInitializationAfterOwnedCallbackDrains(bool after) + { + var storage = new CapturingStorage(); + var manager = CreateTestSystem(storage).Manager; + var state = new LifecycleState(); + manager.RegisterStateMachine("state", state); + var entered = NewSignal(); + var release = NewSignal(); + Func callback = async (_, _) => + { + entered.TrySetResult(); + await release.Task; + }; + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperationAsync = after ? null : callback, + AfterOperationAsync = after ? callback : null + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + var write = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await WaitFor(entered.Task); + var initializing = manager.InitializeAsync(TestContext.Current.CancellationToken).AsTask(); + var deleting = manager.DeleteStateAsync(TestContext.Current.CancellationToken).AsTask(); + var shutdown = manager.DisposeAsync().AsTask(); + Assert.False(shutdown.IsCompleted); + Assert.False(initializing.IsCompleted); + Assert.False(deleting.IsCompleted); + release.TrySetResult(); + await WaitFor(shutdown); + if (after) + { + await WaitFor(write); + } + else + { + await Assert.ThrowsAnyAsync(() => WaitFor(write)); + } + + await Assert.ThrowsAnyAsync(() => WaitFor(initializing)); + await Assert.ThrowsAnyAsync(() => WaitFor(deleting)); + Assert.True(initializing.IsCanceled); + Assert.True(deleting.IsCanceled); + Assert.Equal(after ? 1 : 0, state.CaptureCount); + Assert.Equal(after ? 1 : 0, state.WriteCompletedCount); + Assert.Equal(after ? 1 : 0, storage.Appends.Count); + Assert.Equal(0, storage.DeleteCount); + } + finally + { + release.TrySetResult(); + await manager.DisposeAsync(); + } + } + + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task Hooks_ShutdownCallbackFailureDrainsActualStorageAndPreservesItsOutcome(bool snapshot, bool storageFails) + { + var entered = NewSignal(); + var release = NewSignal(); + var cancellationFailure = new IOException("Cancellation callback failed."); + var storageFailure = new IOException("Actual storage failure."); + var completed = false; + var storage = Substitute.For(); + storage.IsCompactionRequested.Returns(snapshot); + storage.ReadAsync(Arg.Any(), Arg.Any()).Returns(call => + { + call.Arg().Complete(metadata: null); + return ValueTask.CompletedTask; + }); + storage.AppendAsync(Arg.Any>(), Arg.Any()) + .Returns(call => WriteAsync(call.Arg>(), call.Arg())); + storage.ReplaceAsync(Arg.Any>(), Arg.Any()) + .Returns(call => WriteAsync(call.Arg>(), call.Arg())); + var manager = CreateTestSystem(storage).Manager; + var state = new LifecycleState(); + manager.RegisterStateMachine("state", state); + CancellationTokenRegistration registration = default; + var after = 0; + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperation = (_, token) => registration = token.Register(() => throw cancellationFailure), + AfterOperation = (_, _) => after++ + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + var write = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await WaitFor(entered.Task); + var queued = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + var shutdown = manager.DisposeAsync().AsTask(); + Assert.Same(shutdown, manager.DisposeAsync().AsTask()); + Assert.False(shutdown.IsCompleted); + Assert.False(completed); + Assert.Equal(0, state.WriteCompletedCount); + Assert.Equal(0, after); + release.TrySetResult(); + var caught = await Assert.ThrowsAsync(() => WaitFor(shutdown)); + Assert.Same(cancellationFailure, Assert.Single(caught.InnerExceptions)); + Assert.True(completed); + if (storageFails) + { + Assert.Same(storageFailure, await Assert.ThrowsAsync(() => WaitFor(write))); + Assert.Same(storageFailure, await Assert.ThrowsAsync(() => WaitFor(queued))); + } + else + { + await WaitFor(write); + await Assert.ThrowsAnyAsync(() => WaitFor(queued)); + } + + Assert.Equal(storageFails ? 0 : 1, state.WriteCompletedCount); + Assert.Equal(storageFails ? 0 : 1, after); + } + finally + { + release.TrySetResult(); + registration.Dispose(); + await Assert.ThrowsAsync(() => manager.DisposeAsync().AsTask()); + } + + async ValueTask WriteAsync(ReadOnlySequence bytes, CancellationToken token) + { + var original = bytes.ToArray(); + entered.TrySetResult(); + await release.Task; + Assert.True(token.IsCancellationRequested); + Assert.Equal(original, bytes.ToArray()); + completed = true; + if (storageFails) + { + throw storageFailure; + } + } + } + + [Fact] + public async Task Hooks_ShutdownCallbackReentrySharesDisposalCompletion() + { + var storage = new CapturingStorage(); + var manager = CreateTestSystem(storage).Manager; + var state = new DurableValue("state", manager, CreateValueCodec()); + var entered = NewSignal(); + var release = NewSignal(); + Task? reenteredDisposal = null; + var callbackCount = 0; + await manager.AddHookAsync(new DelegateTestHook + { + AfterOperationAsync = async (_, token) => + { + using var registration = token.Register(() => + { + callbackCount++; + reenteredDisposal = manager.DisposeAsync().AsTask(); + }); + entered.TrySetResult(); + await release.Task; + } + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state.Value = 1; + var write = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await WaitFor(entered.Task); + var disposal = manager.DisposeAsync().AsTask(); + Assert.Equal(1, callbackCount); + Assert.Same(disposal, reenteredDisposal); + Assert.False(disposal.IsCompleted); + release.TrySetResult(); + await WaitFor(Task.WhenAll(write, disposal)); + Assert.Same(disposal, manager.DisposeAsync().AsTask()); + Assert.Single(storage.Appends); + } + finally + { + release.TrySetResult(); + await manager.DisposeAsync(); + } + } + + [Fact] + public async Task Hooks_ShutdownDrainsEveryAfterHookDespiteCleanupFailures() + { + var storage = new CapturingStorage(); + var manager = CreateTestSystem(storage).Manager; + var state = new DurableValue("state", manager, CreateValueCodec()); + var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var first = new IOException("First cleanup failed."); + var second = new IOException("Second cleanup failed."); + List events = []; + await manager.AddHookAsync(new DelegateTestHook + { + AfterOperationAsync = async (_, token) => + { + Assert.Single(storage.Appends); + entered.TrySetResult(); + await release.Task; + Assert.True(token.IsCancellationRequested); + events.Add("first"); + throw first; + } + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.AddHookAsync(new DelegateTestHook { AfterOperation = (_, _) => throw second }, cancellationToken: TestContext.Current.CancellationToken); + await manager.AddHookAsync(new DelegateTestHook { AfterOperation = (_, _) => events.Add("last") }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state.Value = 1; + var write = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await entered.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + var shutdown = manager.DisposeAsync().AsTask(); + var concurrentShutdown = manager.DisposeAsync().AsTask(); + Assert.Same(shutdown, concurrentShutdown); + Assert.False(shutdown.IsCompleted); + release.TrySetResult(); + var caught = await Assert.ThrowsAsync(() => + write.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken)); + Assert.Equal(new Exception[] { first, second }, Assert.IsType(caught.InnerException).InnerExceptions); + await shutdown.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.Equal(new[] { "first", "last" }, events); + Assert.Single(storage.Appends); + } + finally + { + release.TrySetResult(); + await manager.DisposeAsync(); + } + } + + [Fact] + public async Task Hooks_QueuedWritesCoalesceAndRetainSnapshotRequirement() + { + var context = new QueuedSynchronizationContext(); + await context.Run(async () => + { + var storage = new CapturingStorage(); + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableValue("state", manager, CreateValueCodec()); + List before = []; + List after = []; + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperation = (operation, _) => before.Add(operation), + AfterOperation = (operation, _) => after.Add(operation) + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state.Value = 1; + var append = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + storage.IsCompactionRequested = true; + state.Value = 2; + var snapshot = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + storage.IsCompactionRequested = false; + state.Value = 3; + var anotherAppend = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + + Assert.Empty(storage.Appends); + Assert.Empty(storage.Replaces); + Assert.Empty(before); + await Task.WhenAll(append, snapshot, anotherAppend); + Assert.Equal([JournaledStateOperation.Snapshot], before); + Assert.Equal([JournaledStateOperation.Snapshot], after); + Assert.Empty(storage.Appends); + Assert.Single(storage.Replaces); + + await using var recoveredManager = CreateTestSystem(storage).Manager; + var recovered = new DurableValue("state", recoveredManager, CreateValueCodec()); + await recoveredManager.InitializeAsync(TestContext.Current.CancellationToken); + Assert.Equal(3, recovered.Value); + }); + } + + [Fact] + public async Task Hooks_BeforeFailureRestoresSnapshotForAlreadyQueuedWrite() + { + var storage = new CapturingStorage { IsCompactionRequested = true }; + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableValue("state", manager, CreateValueCodec()); + var entered = NewSignal(); + var release = NewSignal(); + var failure = new IOException("Snapshot prerequisite failed."); + List before = []; + List after = []; + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperationAsync = async (operation, _) => + { + before.Add(operation); + if (before.Count == 1) + { + entered.TrySetResult(); + await release.Task; + throw failure; + } + }, + AfterOperation = (operation, _) => after.Add(operation) + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state.Value = 1; + var first = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await WaitFor(entered.Task); + storage.IsCompactionRequested = false; + state.Value = 2; + var retry = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + Assert.Empty(storage.Appends); + Assert.Empty(storage.Replaces); + release.TrySetResult(); + var caught = await Assert.ThrowsAsync(() => WaitFor(first)); + Assert.Same(failure, caught.InnerException); + await WaitFor(retry); + Assert.Equal([JournaledStateOperation.Snapshot, JournaledStateOperation.Snapshot], before); + Assert.Equal([JournaledStateOperation.Snapshot], after); + Assert.Empty(storage.Appends); + Assert.Single(storage.Replaces); + await using var recoveredManager = CreateTestSystem(storage).Manager; + var recovered = new DurableValue("state", recoveredManager, CreateValueCodec()); + await recoveredManager.InitializeAsync(TestContext.Current.CancellationToken); + Assert.Equal(2, recovered.Value); + } + finally + { + release.TrySetResult(); + } + } + + [Fact] + public async Task Hooks_SnapshotAcknowledgementPreservesLaterQueuedSnapshot() + { + var storage = new CapturingStorage { IsCompactionRequested = true, BlockNextReplace = true }; + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableValue("state", manager, CreateValueCodec()); + List before = []; + List after = []; + await manager.AddHookAsync(new DelegateTestHook + { + BeforeOperation = (operation, _) => before.Add(operation), + AfterOperation = (operation, _) => after.Add(operation) + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state.Value = 1; + var first = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await WaitFor(storage.ReplaceEntered.Task); + state.Value = 2; + var second = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + storage.IsCompactionRequested = false; + Assert.Empty(storage.Replaces); + Assert.Equal([JournaledStateOperation.Snapshot], before); + Assert.Empty(after); + storage.ReleaseReplace.TrySetResult(); + await WaitFor(Task.WhenAll(first, second)); + Assert.Equal([JournaledStateOperation.Snapshot, JournaledStateOperation.Snapshot], before); + Assert.Equal([JournaledStateOperation.Snapshot, JournaledStateOperation.Snapshot], after); + Assert.Empty(storage.Appends); + Assert.Equal(2, storage.Replaces.Count); + await using var recoveredManager = CreateTestSystem(storage).Manager; + var recovered = new DurableValue("state", recoveredManager, CreateValueCodec()); + await recoveredManager.InitializeAsync(TestContext.Current.CancellationToken); + Assert.Equal(2, recovered.Value); + } + finally + { + storage.ReleaseReplace.TrySetResult(); + } + } + + [Fact] + public async Task Hooks_CoalescedCallersShareActualOperationCallbacks() + { + var storage = new CapturingStorage { BlockNextAppend = true }; + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableDictionary("state", manager, CreateDictionaryCodec()); + var before = 0; + var after = 0; + await manager.AddHookAsync(new DelegateTestHook { BeforeOperation = (_, _) => before++, AfterOperation = (_, _) => after++ }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state["first"] = 1; + var first = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await storage.BlockedAppendStarted.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.Equal(1, before); + Assert.Equal(0, after); + state["second"] = 2; + var second = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + state["third"] = 3; + var third = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + storage.ReleaseAppend.TrySetResult(); + await Task.WhenAll(first, second, third).WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.Equal(2, before); + Assert.Equal(2, after); + Assert.Equal(2, storage.Appends.Count); + } + finally + { + storage.ReleaseAppend.TrySetResult(); + } + } + + [Fact] + public async Task Hooks_StateChangesDuringAfterCallbackPersistInNextOperation() + { + var storage = new CapturingStorage(); + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableDictionary("state", manager, CreateDictionaryCodec()); + var entered = NewSignal(); + var release = NewSignal(); + var after = 0; + await manager.AddHookAsync(new DelegateTestHook + { + AfterOperationAsync = async (_, _) => + { + if (++after == 1) + { + entered.TrySetResult(); + await release.Task; + } + } + }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state["first"] = 1; + var first = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + try + { + await WaitFor(entered.Task); + Assert.Single(storage.Appends); + Assert.False(first.IsCompleted); + state["later"] = 2; + var next = manager.WriteStateAsync(TestContext.Current.CancellationToken).AsTask(); + Assert.False(next.IsCompleted); + + await using (var beforeNext = CreateTestSystem(storage).Manager) + { + var recovered = new DurableDictionary("state", beforeNext, CreateDictionaryCodec()); + await beforeNext.InitializeAsync(TestContext.Current.CancellationToken); + Assert.Equal(1, Assert.Single(recovered).Value); + } + + release.TrySetResult(); + await WaitFor(Task.WhenAll(first, next)); + Assert.Equal(2, after); + Assert.Equal(2, storage.Appends.Count); + await using var afterNext = CreateTestSystem(storage).Manager; + var persisted = new DurableDictionary("state", afterNext, CreateDictionaryCodec()); + await afterNext.InitializeAsync(TestContext.Current.CancellationToken); + Assert.Equal(2, persisted.Count); + Assert.Equal(1, persisted["first"]); + Assert.Equal(2, persisted["later"]); + } + finally + { + release.TrySetResult(); + } + } + + [Fact] + public async Task Hooks_NoByteWriteStillCompletesLogicalOperation() + { + var storage = new CapturingStorage(); + await using var manager = CreateTestSystem(storage).Manager; + var state = new DurableValue("state", manager, CreateValueCodec()); + var before = 0; + var after = 0; + await manager.AddHookAsync(new DelegateTestHook { BeforeOperation = (_, _) => before++, AfterOperation = (_, _) => after++ }, cancellationToken: TestContext.Current.CancellationToken); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + state.Value = 1; + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Equal(1, before); + Assert.Equal(1, after); + before = 0; + after = 0; + await manager.WriteStateAsync(TestContext.Current.CancellationToken); + Assert.Equal(1, before); + Assert.Equal(1, after); + Assert.Single(storage.Appends); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public void Hooks_FailureOutcomesRoundTripWithTypeAndOriginalCause(bool after) + { + Exception original = after + ? new JournaledStatePostCommitException(JournaledStateOperation.Delete, new IOException("Cleanup failed.")) + : new JournaledStatePreCommitException(JournaledStateOperation.Snapshot, new IOException("Scheduling failed.")); + var serializer = ServiceProvider.GetRequiredService(); + var copy = serializer.Deserialize(serializer.SerializeToArray(original)); + if (after) + { + Assert.IsType(copy); + } + else + { + Assert.IsType(copy); + } + + Assert.Equal(original.InnerException!.Message, Assert.IsType(copy.InnerException).Message); + Assert.Equal(original.Message, copy.Message); + } + + private sealed class QueuedHookTaskScheduler : TaskScheduler + { + private readonly Channel _queue = Channel.CreateUnbounded(); + + public override int MaximumConcurrencyLevel => 1; + protected override IEnumerable? GetScheduledTasks() => null; + protected override void QueueTask(Task task) => Assert.True(_queue.Writer.TryWrite(task)); + protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued) => false; + + public Task Queue(Action action) => + Task.Factory.StartNew(action, TestContext.Current.CancellationToken, TaskCreationOptions.DenyChildAttach, this); + + public async Task Run(Func action) + { + var task = Task.Factory.StartNew(action, TestContext.Current.CancellationToken, TaskCreationOptions.DenyChildAttach, this).Unwrap(); + while (!task.IsCompleted) + { + var next = await _queue.Reader.ReadAsync(TestContext.Current.CancellationToken).AsTask().WaitAsync(TimeSpan.FromSeconds(10)); + var previous = SynchronizationContext.Current; + SynchronizationContext.SetSynchronizationContext(null); + try + { + Assert.True(TryExecuteTask(next)); + } + finally + { + SynchronizationContext.SetSynchronizationContext(previous); + } + } + + await task; + } + } + + private sealed class DelegateTestHook( + Func? before = null, + Action? after = null) : IJournaledStateHook + { + public Action? BeforeOperation { get; init; } + public Func? BeforeOperationAsync { get; init; } = before; + public Action? AfterOperation { get; init; } = after; + public Func? AfterOperationAsync { get; init; } + + ValueTask IJournaledStateHook.BeforeOperationAsync(JournaledStateOperation operation, CancellationToken cancellationToken) + { + BeforeOperation?.Invoke(operation, cancellationToken); + return BeforeOperationAsync?.Invoke(operation, cancellationToken) ?? default; + } + + ValueTask IJournaledStateHook.AfterOperationAsync(JournaledStateOperation operation, CancellationToken cancellationToken) + { + AfterOperation?.Invoke(operation, cancellationToken); + return AfterOperationAsync?.Invoke(operation, cancellationToken) ?? default; + } + } + + private sealed class HookAcknowledgementState(List events) : IStateMachine + { + public void ReplayEntry(JournalEntry entry, JournalReplayContext context) { } + public void Reset(JournalStreamWriter writer) { } + public void WritePendingEntries(JournalStreamWriter writer) => events.Add("capture"); + public void WriteSnapshot(JournalStreamWriter writer) => events.Add("capture"); + public void OnWriteCompleted() => events.Add("ack"); + } +}