diff --git a/docs/site/src/content/docs/grains/journaling/durable-state.md b/docs/site/src/content/docs/grains/journaling/durable-state.md index a08a08eae2b..d90dea49e92 100644 --- a/docs/site/src/content/docs/grains/journaling/durable-state.md +++ b/docs/site/src/content/docs/grains/journaling/durable-state.md @@ -51,6 +51,14 @@ Coordination with another grain or an external service requires an application p An changes status in memory when `TrySetResult`, `TrySetException`, or `TrySetCanceled` succeeds. Its `Task` completes after a write acknowledges that status or recovery replays it, allowing waiters to observe a durable completion. +## Own dictionary value resources + +Register for a dictionary value type whose resources have explicit owners. Each live add or assignment acquires an independent owner before encoding, so the caller can release its own value after the mutation. Reads borrow the stored value; acquire an independent owner when retaining a read beyond that entry's lifetime. Reusing a value under multiple keys and assigning an entry to itself each acquire a separate owner. + +Replay transfers decoded owners into the dictionary. Replacement, successful removal, clear, reset, and whole-journal deletion release the stored owners. A rejected mutation releases any owner it acquired and preserves existing entries. Snapshot encoding borrows the current entries. + +Lifecycle implementations preserve value contents and equality when retaining an owner, complete retention atomically, and release synchronously without throwing. Callbacks operate on the supplied value's resources, with dictionary mutation left to the caller. Command codecs borrow values during encoding and transfer decoded owners to the dictionary. Orleans disposes DI-created dictionaries with the activation scope, releasing their remaining owners even when recovery fails partway through a journal. Applications arrange disposal of manually constructed dictionaries and their dependencies. Value types with ordinary managed lifetimes retain the standard dictionary reference semantics. + ## Keep state names and schemas stable The name supplied to the manager or keyed service identifies a durable state component across activations and deployments. Apply these rules: diff --git a/src/Orleans.Journaling/DurableDictionary.cs b/src/Orleans.Journaling/DurableDictionary.cs index 8aeed395024..af36f826afd 100644 --- a/src/Orleans.Journaling/DurableDictionary.cs +++ b/src/Orleans.Journaling/DurableDictionary.cs @@ -16,16 +16,21 @@ public interface IDurableDictionary : IDictionary wh [DebuggerTypeProxy(typeof(IDurableDictionaryDebugView<,>))] [DebuggerDisplay("Count = {Count}")] -internal class DurableDictionary : IDurableDictionary, IStateMachine, IDurableDictionaryCommandHandler where TKey : notnull +internal class DurableDictionary : IDurableDictionary, IStateMachine, IDurableDictionaryCommandHandler, IDisposable where TKey : notnull { private readonly IDurableDictionaryCommandCodec _codec; + private readonly IDurableDictionaryValueLifecycle? _valueLifecycle; private readonly Dictionary _items = []; private JournalStreamWriter _writer; + private bool _disposed; - protected DurableDictionary(IDurableDictionaryCommandCodec codec) + protected DurableDictionary( + IDurableDictionaryCommandCodec codec, + IDurableDictionaryValueLifecycle? valueLifecycle = null) { ArgumentNullException.ThrowIfNull(codec); _codec = codec; + _valueLifecycle = valueLifecycle; } public DurableDictionary( @@ -33,13 +38,19 @@ public DurableDictionary( IJournaledStateManager manager, JournaledStateManagerShared shared, IServiceProvider serviceProvider) - : this(JournalFormatServices.GetRequiredCommandCodec>(serviceProvider, shared.JournalFormatKey)) + : this( + JournalFormatServices.GetRequiredCommandCodec>(serviceProvider, shared.JournalFormatKey), + serviceProvider.GetService>()) { ArgumentNullException.ThrowIfNullOrEmpty(key); manager.RegisterStateMachine(key, this); } - internal DurableDictionary(string key, IJournaledStateManager manager, IDurableDictionaryCommandCodec codec) : this(codec) + internal DurableDictionary( + string key, + IJournaledStateManager manager, + IDurableDictionaryCommandCodec codec, + IDurableDictionaryValueLifecycle? valueLifecycle = null) : this(codec, valueLifecycle) { ArgumentNullException.ThrowIfNullOrEmpty(key); manager.RegisterStateMachine(key, this); @@ -51,8 +62,7 @@ public TValue this[TKey key] set { - WriteSet(key, value); - ApplySet(key, value); + SetValue(key, value); } } @@ -69,7 +79,8 @@ void IStateMachine.ReplayEntry(JournalEntry entry, JournalReplayContext context) void IStateMachine.Reset(JournalStreamWriter writer) { - _items.Clear(); + ThrowIfDisposed(); + ApplyClear(); _writer = writer; } @@ -115,16 +126,78 @@ private void WriteSet(TKey key, TValue value) _codec.WriteSet(key, value, GetWriter()); } + private void SetValue(TKey key, TValue value) + { + ThrowIfDisposed(); + var owned = _valueLifecycle is { } lifecycle ? lifecycle.Retain(value) : value; + try + { + WriteSet(key, owned); + } + catch + { + _valueLifecycle?.Release(owned); + throw; + } + + ApplySet(key, owned); + } + protected virtual void OnSet(TKey key, TValue value) { } private void ApplySet(TKey key, TValue value) { - _items[key] = value; + if (_valueLifecycle is { } lifecycle) + { + bool replacing; + TValue? previous; + try + { + replacing = _items.TryGetValue(key, out previous); + _items[key] = value; + } + catch + { + lifecycle.Release(value); + throw; + } + + if (replacing) + { + lifecycle.Release(previous!); + } + } + else + { + _items[key] = value; + } + OnSet(key, value); } - internal bool ApplyRemove(TKey key) => _items.Remove(key); - private void ApplyClear() => _items.Clear(); + internal bool ApplyRemove(TKey key) + { + if (!_items.Remove(key, out var value)) + { + return false; + } + + _valueLifecycle?.Release(value); + return true; + } + + private void ApplyClear() + { + if (_valueLifecycle is { } lifecycle) + { + foreach (var value in _items.Values) + { + lifecycle.Release(value); + } + } + + _items.Clear(); + } void IDurableDictionaryCommandHandler.ApplySet(TKey key, TValue value) => ApplySet(key, value); void IDurableDictionaryCommandHandler.ApplyRemove(TKey key) => ApplyRemove(key); void IDurableDictionaryCommandHandler.ApplyClear() => ApplyClear(); @@ -136,6 +209,7 @@ void IDurableDictionaryCommandHandler.Reset(int capacityHint) protected virtual JournalStreamWriter GetWriter() { + ThrowIfDisposed(); Debug.Assert(_writer.IsInitialized); return _writer; } @@ -147,9 +221,7 @@ public void Add(TKey key, TValue value) ThrowDuplicateKey(key); } - WriteSet(key, value); - _items.Add(key, value); - OnSet(key, value); + SetValue(key, value); } public bool ContainsKey(TKey key) => _items.ContainsKey(key); @@ -165,12 +237,23 @@ public bool Remove(KeyValuePair item) } WriteRemove(item.Key); - _ = ((ICollection>)_items).Remove(item); + ApplyRemove(item.Key); return true; } public IEnumerator> GetEnumerator() => ((IEnumerable>)_items).GetEnumerator(); + public void Dispose() + { + if (!_disposed && _valueLifecycle is not null) + { + _disposed = true; + ApplyClear(); + } + } + + private void ThrowIfDisposed() => ObjectDisposedException.ThrowIf(_disposed, this); + [DoesNotReturn] private static void ThrowDuplicateKey(TKey key) => throw new ArgumentException($"An item with the same key has already been added. Key: {key}", nameof(key)); } diff --git a/src/Orleans.Journaling/Formats/Json/JsonCommandReader.cs b/src/Orleans.Journaling/Formats/Json/JsonCommandReader.cs index a82175f0bcc..e95ecea5cff 100644 --- a/src/Orleans.Journaling/Formats/Json/JsonCommandReader.cs +++ b/src/Orleans.Journaling/Formats/Json/JsonCommandReader.cs @@ -176,35 +176,48 @@ public bool ReadArrayItem(string operandName) public (TFirst? First, TSecond? Second) ReadCurrentPair( string operandName, JsonTypeInfo firstTypeInfo, - JsonTypeInfo secondTypeInfo) + JsonTypeInfo secondTypeInfo, + bool requireFirst = false) { if (_reader.TokenType is not JsonTokenType.StartArray) { throw new JsonException("JSON dictionary snapshot items must be [key,value] arrays."); } - if (!_reader.Read() || _reader.TokenType is JsonTokenType.EndArray) + // Validate the pair before deserialization can acquire resources for its values. + var validationReader = _reader; + if (!validationReader.Read() || validationReader.TokenType is JsonTokenType.EndArray) { throw new JsonException($"JSON journal command is missing operand '{JsonJournalEntryFields.Key}'."); } - var first = DeserializeCurrentAllowNull(JsonJournalEntryFields.Key, firstTypeInfo); - if (!_reader.Read() || _reader.TokenType is JsonTokenType.EndArray) + validationReader.Skip(); + if (!validationReader.Read() || validationReader.TokenType is JsonTokenType.EndArray) { throw new JsonException($"JSON journal command is missing operand '{JsonJournalEntryFields.Value}'."); } - var second = DeserializeCurrentAllowNull(JsonJournalEntryFields.Value, secondTypeInfo); - if (!_reader.Read()) + validationReader.Skip(); + if (!validationReader.Read()) { throw new JsonException($"JSON journal command operand '{operandName}' array is incomplete."); } - if (_reader.TokenType is not JsonTokenType.EndArray) + if (validationReader.TokenType is not JsonTokenType.EndArray) { throw new JsonException("JSON journal command contains unexpected extra elements."); } + _reader.Read(); + var first = DeserializeCurrentAllowNull(JsonJournalEntryFields.Key, firstTypeInfo); + if (requireFirst && first is null) + { + throw new JsonException($"JSON journal command operand '{JsonJournalEntryFields.Key}' must not be null."); + } + + _reader.Read(); + var second = DeserializeCurrentAllowNull(JsonJournalEntryFields.Value, secondTypeInfo); + _reader.Read(); return (first, second); } @@ -213,18 +226,14 @@ public bool ReadArrayItem(string operandName) JsonTypeInfo firstTypeInfo, JsonTypeInfo secondTypeInfo) { - var (first, second) = ReadCurrentPair(operandName, firstTypeInfo, secondTypeInfo); - if (first is null) - { - throw new JsonException($"JSON journal command operand '{JsonJournalEntryFields.Key}' must not be null."); - } + var (first, second) = ReadCurrentPair(operandName, firstTypeInfo, secondTypeInfo, requireFirst: true); if (second is null) { throw new JsonException($"JSON journal command operand '{JsonJournalEntryFields.Value}' must not be null."); } - return (first, second); + return (first!, second); } public (TFirst First, TSecond? Second) ReadCurrentPairRequiredFirst( @@ -232,18 +241,14 @@ public bool ReadArrayItem(string operandName) JsonTypeInfo firstTypeInfo, JsonTypeInfo secondTypeInfo) { - var (first, second) = ReadCurrentPair(operandName, firstTypeInfo, secondTypeInfo); - if (first is null) - { - throw new JsonException($"JSON journal command operand '{JsonJournalEntryFields.Key}' must not be null."); - } + var (first, second) = ReadCurrentPair(operandName, firstTypeInfo, secondTypeInfo, requireFirst: true); if (second is null && default(TSecond) is not null) { throw new JsonException($"JSON journal command operand '{JsonJournalEntryFields.Value}' must not be null."); } - return (first, second); + return (first!, second); } public void EnsureEnd(int nextIndex) diff --git a/src/Orleans.Journaling/IDurableDictionaryValueLifecycle.cs b/src/Orleans.Journaling/IDurableDictionaryValueLifecycle.cs new file mode 100644 index 00000000000..03d047530e1 --- /dev/null +++ b/src/Orleans.Journaling/IDurableDictionaryValueLifecycle.cs @@ -0,0 +1,34 @@ +namespace Orleans.Journaling; + +/// +/// Defines independent resource ownership for values stored in durable dictionaries. +/// +/// The dictionary value type. +/// +/// Register this service for values whose resources require explicit retention and release. +/// Live mutations acquire an owner before encoding. Replay transfers ownership of decoded values. +/// Replacement, removal, clear, reset, journal deletion, and dictionary disposal release stored owners. +/// Values returned by dictionary reads are borrowed from the dictionary. +/// Implementations must support independent owners of the same resources and run synchronously +/// on the dictionary's logical execution thread, with dictionary mutation left to the caller. +/// Retention preserves the value's contents and equality semantics. +/// Command codecs borrow values during encoding and transfer decoded owners to the command handler. +/// A failed retain leaves ownership with the caller; +/// release must complete without throwing. The activation scope disposes DI-created dictionaries. +/// Callers arrange disposal of manually constructed dictionaries and their dependencies. +/// +public interface IDurableDictionaryValueLifecycle +{ + /// + /// Acquires an independent owner of the value's resources. + /// + /// The borrowed value. + /// The value with independently retained resources. + TValue Retain(TValue value); + + /// + /// Releases the independently owned resources of a value. + /// + /// The owned value being retired. + void Release(TValue value); +} diff --git a/src/Orleans.Journaling/README.md b/src/Orleans.Journaling/README.md index bd3c40be155..c872f072267 100644 --- a/src/Orleans.Journaling/README.md +++ b/src/Orleans.Journaling/README.md @@ -169,6 +169,13 @@ For trimming and Native AOT, use `Configure(...)` to configu ## Custom state and standalone ownership +Register `IDurableDictionaryValueLifecycle` for dictionary values with explicitly owned +resources. Live mutations acquire independent owners before encoding; replay transfers decoded +owners. Reads borrow the stored value. Replacement, removal, clear, reset, journal deletion, and +dictionary disposal release stored owners. The activation scope disposes DI-created dictionaries; +callers arrange disposal of manually constructed components. Lifecycle implementations retain +atomically and release synchronously without throwing. + Register a grain-facing state component contract and its implementation with `services.AddStateMachine()` on `IServiceCollection`. Both type arguments are reference types, and the implementation implements both `TState` and `IStateMachine`. diff --git a/src/api/Orleans.Journaling/Orleans.Journaling.cs b/src/api/Orleans.Journaling/Orleans.Journaling.cs index 3fba3cb84e4..755cad3d8da 100644 --- a/src/api/Orleans.Journaling/Orleans.Journaling.cs +++ b/src/api/Orleans.Journaling/Orleans.Journaling.cs @@ -75,6 +75,12 @@ public partial interface IDurableDictionaryCommandHandler void Reset(int capacityHint); } + public partial interface IDurableDictionaryValueLifecycle + { + void Release(TValue value); + TValue Retain(TValue value); + } + public partial interface IDurableDictionary : System.Collections.Generic.IDictionary, System.Collections.Generic.ICollection>, System.Collections.Generic.IEnumerable>, System.Collections.IEnumerable { } diff --git a/test/Orleans.Journaling.Tests/OwnedDictionaryValueTests.cs b/test/Orleans.Journaling.Tests/OwnedDictionaryValueTests.cs new file mode 100644 index 00000000000..7a5d6f26c0a --- /dev/null +++ b/test/Orleans.Journaling.Tests/OwnedDictionaryValueTests.cs @@ -0,0 +1,815 @@ +using System.Buffers; +using System.Diagnostics.CodeAnalysis; +using System.Text; +using System.Text.Json; +using System.Text.Json.Serialization; +using System.Text.Json.Serialization.Metadata; +using Microsoft.Extensions.Configuration; +using Microsoft.Extensions.DependencyInjection; +using NSubstitute; +using Orleans.Hosting; +using Orleans.Journaling.Json; +using Orleans.Serialization; +using Orleans.Serialization.Buffers; +using Orleans.Serialization.Codecs; +using Orleans.Serialization.WireProtocol; +using Xunit; + +namespace Orleans.Journaling.Tests; + +[TestSuite("BVT")] +[TestProvider("None")] +[TestCategory("BVT")] +public sealed class OwnedDictionaryValueTests : JournalingTestBase +{ + [Theory] + [InlineData(false)] + [InlineData(true)] + public void AddAndSet_RetainBeforeEncoding_CallerReleasePreservesStoredValue(bool useIndexer) + { + using var fixture = CreateFixture(); + using var caller = new ResourceProbe(42); + if (useIndexer) fixture.Dictionary["key"] = caller; + else fixture.Dictionary.Add("key", caller); + + var stored = fixture.Dictionary["key"]; + Assert.NotSame(caller, stored); + Assert.Same(stored, Assert.Single(fixture.ValueCodec.Encoded)); + Assert.Equal(2, caller.Owners); + caller.Dispose(); + Assert.Equal(42, stored.Value); + Assert.Equal(1, stored.Owners); + fixture.Dictionary.Dispose(); + fixture.Dictionary.Dispose(); + Assert.Equal(1, stored.Disposes); + Assert.Equal(0, stored.Owners); + Assert.Equal(1, fixture.Lifecycle.Retains); + Assert.Equal([stored], fixture.Lifecycle.Released); + } + + [Fact] + public void MultipleKeysAndSelfReplacement_AcquireAndRetireIndependentOwners() + { + using var fixture = CreateFixture(); + using var caller = new ResourceProbe(17); + fixture.Dictionary.Add("first", caller); + caller.Dispose(); + var first = fixture.Dictionary["first"]; + fixture.Dictionary.Add("neighbor", first); + var neighbor = fixture.Dictionary["neighbor"]; + fixture.Dictionary["first"] = first; + var replacement = fixture.Dictionary["first"]; + + Assert.NotSame(first, neighbor); + Assert.NotSame(first, replacement); + Assert.Equal(1, first.Disposes); + Assert.Equal(2, replacement.Owners); + Assert.Equal(17, replacement.Value); + Assert.Equal(17, neighbor.Value); + Assert.True(fixture.Dictionary.Remove("first")); + Assert.False(fixture.Dictionary.Remove("first")); + Assert.Equal(1, replacement.Disposes); + Assert.Equal(17, neighbor.Value); + Assert.Equal(1, neighbor.Owners); + fixture.Dictionary.Clear(); + fixture.Dictionary.Dispose(); + Assert.Equal(1, neighbor.Disposes); + Assert.Equal(0, caller.Owners); + Assert.Equal(3, fixture.Lifecycle.Retains); + Assert.Equal([first, replacement, neighbor], fixture.Lifecycle.Released); + } + + [Fact] + public void DuplicateAddAndPairRemoval_ReleaseOnlyTheMatchingStoredOwner() + { + using var fixture = CreateFixture(); + using var caller = new ResourceProbe(23); + using var other = new ResourceProbe(99); + fixture.Dictionary.Add("key", caller); + var stored = fixture.Dictionary["key"]; + var before = fixture.JournalBytes(); + + Assert.Throws(() => fixture.Dictionary.Add("key", other)); + Assert.False(fixture.Dictionary.Remove("missing")); + Assert.False(fixture.Dictionary.Remove(new KeyValuePair("key", other))); + Assert.False(fixture.Dictionary.Remove(new KeyValuePair("missing", stored))); + Assert.Equal(before, fixture.JournalBytes()); + Assert.Equal(1, fixture.Lifecycle.Retains); + Assert.Empty(fixture.Lifecycle.Released); + Assert.Equal(23, stored.Value); + Assert.Equal(1, other.Owners); + + Assert.True(fixture.Dictionary.Remove(new KeyValuePair("key", stored))); + Assert.Empty(fixture.Dictionary); + Assert.Equal([stored], fixture.Lifecycle.Released); + Assert.Equal(1, stored.Disposes); + Assert.Equal(1, caller.Owners); + Assert.Equal(23, caller.Value); + } + + [Fact] + public void EqualPairRemoval_RetiresStoredOwner_AndPreservesBorrowedCaller() + { + using var fixture = CreateFixture(); + using var caller = new ResourceProbe(23); + fixture.Dictionary.Add("key", caller); + var stored = fixture.Dictionary["key"]; + Assert.NotSame(caller, stored); + Assert.Equal(caller, stored); + var pair = new KeyValuePair("key", caller); + + Assert.Contains(pair, fixture.Dictionary); + Assert.True(fixture.Dictionary.Remove(pair)); + + Assert.Same(stored, Assert.Single(fixture.Lifecycle.Released)); + Assert.Equal(1, stored.Disposes); + Assert.Equal(0, caller.Disposes); + Assert.Equal(1, caller.Owners); + Assert.Equal(23, caller.Value); + Assert.Empty(fixture.Dictionary); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public void EncodingFailure_ReleasesAcquiredOwner_PreservesJournalAndEntries(bool useIndexer) + { + using var fixture = CreateFixture(); + using var old = new ResourceProbe(10); + using var neighbor = new ResourceProbe(20); + using var proposed = new ResourceProbe(30); + fixture.Dictionary.Add("old", old); + fixture.Dictionary.Add("neighbor", neighbor); + var stored = fixture.Dictionary["old"]; + var storedNeighbor = fixture.Dictionary["neighbor"]; + var before = fixture.JournalBytes(); + var error = new InvalidOperationException("encoding failure"); + fixture.ValueCodec.Failure = error; + + Assert.Same(error, Record.Exception(() => + { + if (useIndexer) fixture.Dictionary["old"] = proposed; + else fixture.Dictionary.Add("new", proposed); + })); + + var rejected = fixture.ValueCodec.Encoded[^1]; + Assert.NotSame(proposed, rejected); + Assert.Equal([rejected], fixture.Lifecycle.Released); + Assert.Equal(1, rejected.Disposes); + Assert.Equal(1, proposed.Owners); + Assert.Equal(30, proposed.Value); + Assert.Same(stored, fixture.Dictionary["old"]); + Assert.Same(storedNeighbor, fixture.Dictionary["neighbor"]); + Assert.Equal(10, stored.Value); + Assert.Equal(20, storedNeighbor.Value); + Assert.Equal(2, fixture.Dictionary.Count); + Assert.False(fixture.Dictionary.ContainsKey("new")); + Assert.Equal(before, fixture.JournalBytes()); + } + + [Fact] + public void RetainFailure_PreservesCallerAndJournal() + { + using var fixture = CreateFixture(); + using var caller = new ResourceProbe(31); + var error = new InvalidOperationException("retain failure"); + fixture.Lifecycle.Failure = error; + + Assert.Same(error, Record.Exception(() => fixture.Dictionary.Add("key", caller))); + + Assert.Equal(1, fixture.Lifecycle.Retains); + Assert.Empty(fixture.ValueCodec.Encoded); + Assert.Empty(fixture.Lifecycle.Released); + Assert.Empty(fixture.JournalBytes()); + Assert.Empty(fixture.Dictionary); + Assert.Equal(1, caller.Owners); + Assert.Equal(31, caller.Value); + } + + [Fact] + public void ReplayInsertionFailure_ReleasesTransferredOwner() + { + using var fixture = CreateFixture(); + var decoded = new ResourceProbe(11); + var handler = (IDurableDictionaryCommandHandler)fixture.Dictionary; + + Assert.Throws(() => handler.ApplySet(null!, decoded)); + + Assert.Equal(1, decoded.Disposes); + Assert.Equal(0, decoded.Owners); + Assert.Equal([decoded], fixture.Lifecycle.Released); + Assert.Equal(0, fixture.Lifecycle.Retains); + Assert.Empty(fixture.Dictionary); + } + + [Theory] + [InlineData("remove")] + [InlineData("pair-remove")] + [InlineData("clear")] + public void RemoveAndClearEncodingFailure_PreserveStoredOwners(string operation) + { + using var fixture = CreateFixture(); + using var caller = new ResourceProbe(84); + fixture.Dictionary.Add("key", caller); + var stored = fixture.Dictionary["key"]; + var before = fixture.JournalBytes(); + var error = new InvalidOperationException("command encoding failure"); + fixture.Codec.Failure = error; + + Assert.Same(error, Record.Exception(() => + { + switch (operation) + { + case "remove": fixture.Dictionary.Remove("key"); break; + case "pair-remove": fixture.Dictionary.Remove(new KeyValuePair("key", stored)); break; + case "clear": fixture.Dictionary.Clear(); break; + default: throw new ArgumentException(operation); + } + })); + + Assert.Same(stored, fixture.Dictionary["key"]); + Assert.Equal(84, stored.Value); + Assert.Equal(2, caller.Owners); + Assert.Empty(fixture.Lifecycle.Released); + Assert.Equal(before, fixture.JournalBytes()); + } + + [Theory] + [InlineData("clear")] + [InlineData("state-reset")] + [InlineData("command-reset")] + [InlineData("replay-clear")] + [InlineData("dispose")] + public void ClearResetAndDispose_RetireEveryStoredOwnerOnce(string operation) + { + using var fixture = CreateFixture(); + using var first = new ResourceProbe(1); + using var second = new ResourceProbe(2); + fixture.Dictionary.Add("first", first); + fixture.Dictionary.Add("second", second); + var owners = fixture.Dictionary.Values.ToArray(); + void Act() + { + switch (operation) + { + case "clear": fixture.Dictionary.Clear(); break; + case "state-reset": ((IStateMachine)fixture.Dictionary).Reset(fixture.Writer); break; + case "command-reset": ((IDurableDictionaryCommandHandler)fixture.Dictionary).Reset(10); break; + case "replay-clear": ((IDurableDictionaryCommandHandler)fixture.Dictionary).ApplyClear(); break; + case "dispose": fixture.Dictionary.Dispose(); break; + default: throw new ArgumentException(operation); + } + } + + Act(); + Act(); + fixture.Dictionary.Dispose(); + Assert.Empty(fixture.Dictionary); + Assert.Equal(owners, fixture.Lifecycle.Released); + Assert.All(owners, owner => Assert.Equal(1, owner.Disposes)); + Assert.Equal(1, first.Owners); + Assert.Equal(1, second.Owners); + Assert.Equal(1, first.Value); + Assert.Equal(2, second.Value); + } + + [Fact] + public async Task JournalReplay_TransfersDecodedOwners_ReplacementRemovalAndDeletionRetireOnce() + { + var token = TestContext.Current.CancellationToken; + var storage = new VolatileJournalStorage(OrleansBinaryJournalFormat.JournalFormatKey); + var source = CreateTestSystem(storage); + await using var sourceManager = source.Manager; + var sourceCodec = CreateValueCodec(); + var sourceLifecycle = new ProbeLifecycle(); + using var sourceDictionary = new DurableDictionary( + "owned", source.Manager, CreateCodec(sourceCodec), sourceLifecycle); + await source.Manager.InitializeAsync(token); + using var first = new ResourceProbe(11); + using var second = new ResourceProbe(22); + using var neighbor = new ResourceProbe(33); + sourceDictionary.Add("key", first); + sourceDictionary.Add("neighbor", neighbor); + await source.Manager.WriteStateAsync(token); + sourceDictionary["key"] = second; + sourceDictionary.Add("removed", first); + sourceDictionary.Remove("removed"); + await source.Manager.WriteStateAsync(token); + Assert.NotEmpty(storage.Segments); + + var recovered = CreateTestSystem(storage); + await using var recoveredManager = recovered.Manager; + var replayCodec = CreateValueCodec(); + var replayLifecycle = new ProbeLifecycle(); + using var dictionary = new DurableDictionary( + "owned", recovered.Manager, CreateCodec(replayCodec), replayLifecycle); + await recovered.Manager.InitializeAsync(token); + + Assert.Equal(0, replayLifecycle.Retains); + Assert.Equal(4, replayCodec.Decoded.Count); + Assert.Equal([replayCodec.Decoded[0], replayCodec.Decoded[3]], replayLifecycle.Released); + Assert.Equal(22, dictionary["key"].Value); + Assert.Equal(33, dictionary["neighbor"].Value); + Assert.Same(replayCodec.Decoded[2], dictionary["key"]); + Assert.Same(replayCodec.Decoded[1], dictionary["neighbor"]); + Assert.Equal(1, dictionary["key"].Owners); + Assert.Equal(1, dictionary["neighbor"].Owners); + + await recovered.Manager.DeleteStateAsync(token); + Assert.Empty(dictionary); + Assert.Empty(storage.Segments); + dictionary.Dispose(); + Assert.Equal(4, replayLifecycle.Released.Count); + Assert.All(replayCodec.Decoded, value => + { + Assert.Equal(1, value.Disposes); + Assert.Equal(0, value.Owners); + }); + Assert.Equal(22, sourceDictionary["key"].Value); + Assert.Equal(33, sourceDictionary["neighbor"].Value); + } + + [Fact] + public void SnapshotReplay_TransfersDecodedOwners_ResetsPreviousContents() + { + using var source = CreateFixture(); + using var caller = new ResourceProbe(71); + source.Dictionary.Add("key", caller); + var stored = source.Dictionary["key"]; + byte[]? previous = null; + for (var i = 0; i < 3; i++) + { + var payload = CodecTestHelpers.WriteEntry(writer => ((IStateMachine)source.Dictionary).WriteSnapshot(writer)); + if (previous is not null) Assert.Equal(previous, payload); + previous = payload; + using var target = CreateFixture(); + using var prior = new ResourceProbe(99); + target.Dictionary.Add("prior", prior); + var oldOwner = target.Dictionary["prior"]; + target.Codec.Apply(CodecTestHelpers.ReadBuffer(payload), target.Dictionary); + var decoded = Assert.Single(target.ValueCodec.Decoded); + + Assert.Same(decoded, target.Dictionary["key"]); + Assert.Equal(71, decoded.Value); + Assert.Equal(["key"], target.Dictionary.Keys); + Assert.Equal([oldOwner], target.Lifecycle.Released); + Assert.Equal(1, target.Lifecycle.Retains); + Assert.Equal(1, decoded.Owners); + target.Dictionary.Dispose(); + Assert.Equal(1, decoded.Disposes); + Assert.Equal(0, decoded.Owners); + Assert.Equal(2, caller.Owners); + Assert.Equal(71, stored.Value); + Assert.Empty(source.Lifecycle.Released); + Assert.Equal(1, source.Lifecycle.Retains); + } + } + + [Theory] + [InlineData(false, false)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task JsonRecovery_ScopeDisposalRetiresTransferredOwners_AfterSuccessFailureOrRetry(bool failRecovery, bool retryRecovery) + { + var token = TestContext.Current.CancellationToken; + var storage = new VolatileJournalStorage(JsonLinesJournalFormat.JournalFormatKey); + var lifecycle = new ProbeLifecycle(); + var converter = new ProbeJsonConverter(); + var builder = new TestSiloBuilder(); + builder.Services.AddLogging(); + builder.Services.AddSerializer(); + builder.Services.AddKeyedSingleton(JournalingTimeProviderNames.Journaling, TimeProvider.System); + builder.AddJournaling(); + builder.Services.Configure(options => options.SerializerOptions = new() + { + TypeInfoResolver = new DefaultJsonTypeInfoResolver(), + Converters = { converter } + }); + builder.Services.AddSingleton>(lifecycle); + builder.Services.AddScoped(services => + new JournaledStateManager(services.GetRequiredService(), storage)); + await using var provider = builder.Services.BuildServiceProvider(); + using var first = new ResourceProbe(11); + using var second = new ResourceProbe(22); + using var neighbor = new ResourceProbe(33); + await using (var scope = provider.CreateAsyncScope()) + { + var dictionary = scope.ServiceProvider.GetRequiredKeyedService>("owned"); + var manager = scope.ServiceProvider.GetRequiredService(); + await manager.InitializeAsync(token); + dictionary.Add("key", first); + dictionary.Add("neighbor", neighbor); + dictionary["key"] = second; + await manager.WriteStateAsync(token); + Assert.Equal( + "[0,[\"set\",\"owned\",8]]\n" + + "[8,[\"set\",\"key\",11]]\n" + + "[8,[\"set\",\"neighbor\",33]]\n" + + "[8,[\"set\",\"key\",22]]\n", + Encoding.UTF8.GetString(Assert.Single(storage.Segments))); + Assert.Equal(1, converter.Encoded[0].Disposes); + Assert.Equal(0, converter.Encoded[1].Disposes); + Assert.Equal(0, converter.Encoded[2].Disposes); + } + + Assert.Equal(3, lifecycle.Released.Count); + Assert.All(converter.Encoded, value => Assert.Equal(1, value.Disposes)); + var validJournal = Assert.Single(storage.Segments); + if (failRecovery) + { + await storage.AppendAsync( + new ReadOnlySequence(Encoding.UTF8.GetBytes("[8,[\"set\",\"partial\",44,\"unexpected\"]]\n")), token); + } + + await using (var scope = provider.CreateAsyncScope()) + { + var dictionary = scope.ServiceProvider.GetRequiredKeyedService>("owned"); + var manager = scope.ServiceProvider.GetRequiredService(); + if (failRecovery) + { + var error = await Assert.ThrowsAsync(() => manager.InitializeAsync(token).AsTask()); + Assert.NotNull(error.InnerException); + var uninitialized = await Assert.ThrowsAsync(() => manager.WriteStateAsync(token).AsTask()); + Assert.Equal("The journaled state manager has not been initialized.", uninitialized.Message); + Assert.Same(converter.Decoded[3], dictionary["partial"]); + Assert.Equal(44, dictionary["partial"].Value); + } + else + { + await manager.InitializeAsync(token); + } + + Assert.Equal(3, lifecycle.Retains); + Assert.Equal(failRecovery ? 4 : 3, converter.Decoded.Count); + Assert.Same(converter.Decoded[2], dictionary["key"]); + Assert.Same(converter.Decoded[1], dictionary["neighbor"]); + Assert.Equal(22, dictionary["key"].Value); + Assert.Equal(33, dictionary["neighbor"].Value); + Assert.Equal(1, converter.Decoded[0].Disposes); + Assert.Equal(4, lifecycle.Released.Count); + ((IDurableDictionaryCommandHandler)dictionary).ApplyRemove("missing"); + Assert.Equal(4, lifecycle.Released.Count); + if (retryRecovery) + { + var previousOwners = converter.Decoded.ToArray(); + await storage.ReplaceAsync(new ReadOnlySequence(validJournal), token); + await manager.InitializeAsync(token); + + Assert.All(previousOwners, value => + { + Assert.Equal(1, value.Disposes); + Assert.Equal(0, value.Owners); + }); + Assert.Equal(7, converter.Decoded.Count); + Assert.Equal(8, lifecycle.Released.Count); + Assert.Equal(3, lifecycle.Retains); + Assert.Same(converter.Decoded[6], dictionary["key"]); + Assert.Same(converter.Decoded[5], dictionary["neighbor"]); + Assert.Equal(22, dictionary["key"].Value); + Assert.Equal(33, dictionary["neighbor"].Value); + Assert.False(dictionary.ContainsKey("partial")); + } + } + + Assert.Equal(3 + converter.Decoded.Count, lifecycle.Released.Count); + Assert.All(converter.Decoded, value => + { + Assert.Equal(1, value.Disposes); + Assert.Equal(0, value.Owners); + }); + Assert.Equal(1, first.Owners); + Assert.Equal(1, second.Owners); + Assert.Equal(1, neighbor.Owners); + Assert.Equal(11, first.Value); + Assert.Equal(22, second.Value); + Assert.Equal(33, neighbor.Value); + } + + [Theory] + [InlineData("[\"snapshot\",[[\"bad\",44,\"unexpected\"]]]")] + [InlineData("[\"snapshot\",[[null,44]]]")] + [InlineData("[\"snapshot\",[[\"bad\"]]]")] + public void JsonSnapshotValidation_RejectsInvalidPairBeforeMaterializingOwners(string payload) + { + using var fixture = CreateFixture(); + using var caller = new ResourceProbe(11); + fixture.Dictionary.Add("prior", caller); + var previousOwner = fixture.Dictionary["prior"]; + var converter = new ProbeJsonConverter(); + var codec = new JsonDurableDictionaryCommandCodec(new JsonSerializerOptions + { + TypeInfoResolver = new DefaultJsonTypeInfoResolver(), + Converters = { converter } + }); + + Assert.Throws(() => codec.Apply(CodecTestHelpers.ReadBuffer(Encoding.UTF8.GetBytes(payload)), fixture.Dictionary)); + + Assert.Empty(converter.Decoded); + Assert.Empty(fixture.Dictionary); + Assert.Same(previousOwner, Assert.Single(fixture.Lifecycle.Released)); + Assert.Equal(1, previousOwner.Disposes); + Assert.Equal(1, caller.Owners); + Assert.Equal(11, caller.Value); + } + + [Fact] + public void ScopedResolution_UsesRegisteredLifecycle_AndDisposesStoredOwners() + { + using var journal = new OrleansBinaryJournalBufferWriter(); + using var caller = new ResourceProbe(56); + var lifecycle = new ProbeLifecycle(); + var codec = CreateCodec(CreateValueCodec()); + var builder = new TestSiloBuilder(); + builder.Services.AddLogging(); + builder.Services.AddSerializer(); + builder.Services.AddKeyedSingleton(JournalingTimeProviderNames.Journaling, TimeProvider.System); + builder.Services.AddSingleton(CreateLocalManager(journal)); + builder.Services.Configure(options => options.JournalFormatKey = OrleansBinaryJournalFormat.JournalFormatKey); + builder.AddJournaling(); + builder.Services.AddSingleton>(lifecycle); + builder.Services.AddKeyedSingleton>(OrleansBinaryJournalFormat.JournalFormatKey, codec); + using var provider = builder.Services.BuildServiceProvider(); + ResourceProbe stored; + using (var scope = provider.CreateScope()) + { + var dictionary = scope.ServiceProvider.GetRequiredKeyedService>("owned"); + Assert.Same(dictionary, scope.ServiceProvider.GetRequiredKeyedService>("owned")); + dictionary.Add("key", caller); + stored = dictionary["key"]; + Assert.Equal(56, stored.Value); + Assert.Equal(2, caller.Owners); + Assert.Equal(1, lifecycle.Retains); + } + + Assert.Equal([stored], lifecycle.Released); + Assert.Equal(1, stored.Disposes); + Assert.Equal(1, caller.Owners); + Assert.Equal(56, caller.Value); + } + + [Fact] + public async Task StandaloneManagerDisposal_LeavesComponentOwnershipWithCaller() + { + var system = CreateTestSystem(); + await using var manager = system.Manager; + var lifecycle = new ProbeLifecycle(); + using var dictionary = new DurableDictionary( + "owned", manager, CreateCodec(CreateValueCodec()), lifecycle); + await manager.InitializeAsync(TestContext.Current.CancellationToken); + using var caller = new ResourceProbe(59); + dictionary.Add("key", caller); + var stored = dictionary["key"]; + + await manager.DisposeAsync(); + + Assert.Same(stored, dictionary["key"]); + Assert.Equal(59, stored.Value); + Assert.Empty(lifecycle.Released); + Assert.Equal(2, caller.Owners); + dictionary.Dispose(); + Assert.Equal([stored], lifecycle.Released); + Assert.Equal(1, stored.Disposes); + Assert.Equal(1, caller.Owners); + Assert.Equal(59, caller.Value); + } + + [Fact] + public void UnregisteredValues_PreserveReferenceAndDisposalBehavior() + { + using var fixture = CreateFixture(registerLifecycle: false); + using var first = new ResourceProbe(17); + using var second = new ResourceProbe(23); + fixture.Dictionary.Add("key", first); + Assert.Same(first, fixture.Dictionary["key"]); + fixture.Dictionary["key"] = second; + Assert.Same(second, fixture.Dictionary["key"]); + Assert.True(fixture.Dictionary.Remove("key")); + fixture.Dictionary.Add("key", first); + fixture.Dictionary.Clear(); + fixture.Dictionary.Add("key", second); + ((IStateMachine)fixture.Dictionary).Reset(fixture.Writer); + fixture.Dictionary.Add("survivor", first); + fixture.Dictionary.Dispose(); + Assert.Same(first, fixture.Dictionary["survivor"]); + fixture.Dictionary["survivor"] = second; + Assert.Same(second, fixture.Dictionary["survivor"]); + Assert.Equal(0, first.Disposes); + Assert.Equal(0, second.Disposes); + Assert.Equal(1, first.Owners); + Assert.Equal(1, second.Owners); + } + + [Fact] + public void DisposedDictionary_RejectsMutationsBeforeAcquiringOwners() + { + using var fixture = CreateFixture(); + using var caller = new ResourceProbe(26); + fixture.Dictionary.Dispose(); + Assert.Throws(() => fixture.Dictionary.Add("key", caller)); + Assert.Throws(() => fixture.Dictionary["key"] = caller); + Assert.Throws(fixture.Dictionary.Clear); + Assert.Throws(() => ((IStateMachine)fixture.Dictionary).Reset(fixture.Writer)); + Assert.Equal(0, fixture.Lifecycle.Retains); + Assert.Empty(fixture.Lifecycle.Released); + Assert.Empty(fixture.ValueCodec.Encoded); + Assert.Empty(fixture.JournalBytes()); + Assert.Equal(1, caller.Owners); + Assert.Equal(26, caller.Value); + } + + private ProbeCodec CreateValueCodec() => new(CodecProvider.GetCodec()); + + private CommandCodec CreateCodec(ProbeCodec valueCodec) => + new(new OrleansBinaryDurableDictionaryCommandCodec(CodecProvider.GetCodec(), valueCodec, SessionPool)); + + private Fixture CreateFixture(bool registerLifecycle = true) + { + var valueCodec = CreateValueCodec(); + return new(valueCodec, CreateCodec(valueCodec), registerLifecycle); + } + + private static IJournaledStateManager CreateLocalManager(OrleansBinaryJournalBufferWriter journal) + { + var manager = Substitute.For(); + manager.When(value => value.RegisterStateMachine(Arg.Any(), Arg.Any())) + .Do(call => call.Arg().Reset(journal.CreateJournalStreamWriter(new(1)))); + return manager; + } + + private sealed class Fixture : IDisposable + { + private readonly OrleansBinaryJournalBufferWriter _journal = new(); + + public Fixture(ProbeCodec valueCodec, CommandCodec codec, bool registerLifecycle) + { + ValueCodec = valueCodec; + Codec = codec; + Dictionary = new("owned", CreateLocalManager(_journal), codec, registerLifecycle ? Lifecycle : null); + } + + public ProbeLifecycle Lifecycle { get; } = new(); + public ProbeCodec ValueCodec { get; } + public CommandCodec Codec { get; } + public DurableDictionary Dictionary { get; } + public JournalStreamWriter Writer => _journal.CreateJournalStreamWriter(new(1)); + + public byte[] JournalBytes() + { + using var buffer = _journal.GetBuffer(); + return buffer.ToArray(); + } + + public void Dispose() + { + Dictionary.Dispose(); + _journal.Dispose(); + } + } + + private sealed class TestSiloBuilder : ISiloBuilder + { + public IServiceCollection Services { get; } = new ServiceCollection(); + public IConfiguration Configuration { get; } = new ConfigurationBuilder().Build(); + } + + private sealed class ResourceProbe : IDisposable, IEquatable + { + private readonly Resource _resource; + + public ResourceProbe(int value) : this(new Resource(value)) { } + + private ResourceProbe(Resource resource) + { + _resource = resource; + _resource.Owners++; + } + + public int Owners => _resource.Owners; + public int Disposes { get; private set; } + public int Value + { + get + { + ObjectDisposedException.ThrowIf(Disposes != 0, this); + return _resource.Value; + } + } + + public ResourceProbe Retain() + { + _ = Value; + return new(_resource); + } + + public bool Equals(ResourceProbe? other) => other is not null && ReferenceEquals(_resource, other._resource); + + public override bool Equals(object? obj) => obj is ResourceProbe other && Equals(other); + + public override int GetHashCode() => _resource.GetHashCode(); + + public void Dispose() + { + if (Disposes != 0) return; + Disposes++; + _resource.Owners--; + } + + private sealed class Resource(int value) + { + public int Value { get; } = value; + public int Owners { get; set; } + } + } + + private sealed class ProbeLifecycle : IDurableDictionaryValueLifecycle + { + public int Retains { get; private set; } + public List Released { get; } = []; + public Exception? Failure { get; set; } + + public ResourceProbe Retain(ResourceProbe value) + { + Retains++; + if (Failure is { } error) throw error; + return value.Retain(); + } + + public void Release(ResourceProbe value) + { + Assert.DoesNotContain(Released, item => ReferenceEquals(item, value)); + Assert.Equal(0, value.Disposes); + Released.Add(value); + value.Dispose(); + } + } + + private sealed class ProbeJsonConverter : JsonConverter + { + public List Encoded { get; } = []; + public List Decoded { get; } = []; + + public override ResourceProbe Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options) + { + var value = new ResourceProbe(reader.GetInt32()); + Decoded.Add(value); + return value; + } + + public override void Write(Utf8JsonWriter writer, ResourceProbe value, JsonSerializerOptions options) + { + Encoded.Add(value); + writer.WriteNumberValue(value.Value); + } + } + + private sealed class ProbeCodec(IFieldCodec inner) : IFieldCodec + { + public List Encoded { get; } = []; + public List Decoded { get; } = []; + public Exception? Failure { get; set; } + + public void WriteField(ref Writer writer, uint fieldIdDelta, [AllowNull] Type expectedType, [AllowNull] ResourceProbe value) + where TBufferWriter : IBufferWriter + { + Encoded.Add(value!); + inner.WriteField(ref writer, fieldIdDelta, typeof(int), value!.Value); + if (Failure is { } error) throw error; + } + + public ResourceProbe ReadValue(ref Reader reader, Field field) + { + var value = new ResourceProbe(inner.ReadValue(ref reader, field)); + Decoded.Add(value); + return value; + } + } + + private sealed class CommandCodec(IDurableDictionaryCommandCodec inner) : IDurableDictionaryCommandCodec + { + public Exception? Failure { get; set; } + + public void WriteSet(string key, ResourceProbe value, JournalStreamWriter writer) => inner.WriteSet(key, value, writer); + public void WriteSnapshot(IReadOnlyCollection> items, JournalStreamWriter writer) => inner.WriteSnapshot(items, writer); + public void Apply(JournalBufferReader input, IDurableDictionaryCommandHandler consumer) => inner.Apply(input, consumer); + + public void WriteRemove(string key, JournalStreamWriter writer) + { + ThrowEncodingFailure(writer); + inner.WriteRemove(key, writer); + } + + public void WriteClear(JournalStreamWriter writer) + { + ThrowEncodingFailure(writer); + inner.WriteClear(writer); + } + + private void ThrowEncodingFailure(JournalStreamWriter writer) + { + if (Failure is not { } error) return; + using var entry = writer.BeginEntry(); + entry.Writer.Write([0xff]); + throw error; + } + } +}