feat(journaling): manage independently owned dictionary values

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 <xref:Orleans.Journaling.IDurableTaskCompletionSource`1> 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 <xref:Orleans.Journaling.IDurableDictionaryValueLifecycle`1> 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:

src/Orleans.Journaling/DurableDictionary.cs

@@ -16,30 +16,41 @@ public interface IDurableDictionary<TKey, TValue> : IDictionary<TKey, TValue> wh
 
 [DebuggerTypeProxy(typeof(IDurableDictionaryDebugView<,>))]
 [DebuggerDisplay("Count = {Count}")]
-internal class DurableDictionary<TKey, TValue> : IDurableDictionary<TKey, TValue>, IStateMachine, IDurableDictionaryCommandHandler<TKey, TValue> where TKey : notnull
+internal class DurableDictionary<TKey, TValue> : IDurableDictionary<TKey, TValue>, IStateMachine, IDurableDictionaryCommandHandler<TKey, TValue>, IDisposable where TKey : notnull
 {
     private readonly IDurableDictionaryCommandCodec<TKey, TValue> _codec;
+    private readonly IDurableDictionaryValueLifecycle<TValue>? _valueLifecycle;
     private readonly Dictionary<TKey, TValue> _items = [];
     private JournalStreamWriter _writer;
+    private bool _disposed;
 
-    protected DurableDictionary(IDurableDictionaryCommandCodec<TKey, TValue> codec)
+    protected DurableDictionary(
+        IDurableDictionaryCommandCodec<TKey, TValue> codec,
+        IDurableDictionaryValueLifecycle<TValue>? valueLifecycle = null)
     {
         ArgumentNullException.ThrowIfNull(codec);
         _codec = codec;
+        _valueLifecycle = valueLifecycle;
     }
 
     public DurableDictionary(
         [ServiceKey] string key,
         IJournaledStateManager manager,
         JournaledStateManagerShared shared,
         IServiceProvider serviceProvider)
-        : this(JournalFormatServices.GetRequiredCommandCodec<IDurableDictionaryCommandCodec<TKey, TValue>>(serviceProvider, shared.JournalFormatKey))
+        : this(
+            JournalFormatServices.GetRequiredCommandCodec<IDurableDictionaryCommandCodec<TKey, TValue>>(serviceProvider, shared.JournalFormatKey),
+            serviceProvider.GetService<IDurableDictionaryValueLifecycle<TValue>>())
     {
         ArgumentNullException.ThrowIfNullOrEmpty(key);
         manager.RegisterStateMachine(key, this);
     }
 
-    internal DurableDictionary(string key, IJournaledStateManager manager, IDurableDictionaryCommandCodec<TKey, TValue> codec) : this(codec)
+    internal DurableDictionary(
+        string key,
+        IJournaledStateManager manager,
+        IDurableDictionaryCommandCodec<TKey, TValue> codec,
+        IDurableDictionaryValueLifecycle<TValue>? 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<TKey, TValue>.ApplySet(TKey key, TValue value) => ApplySet(key, value);
     void IDurableDictionaryCommandHandler<TKey, TValue>.ApplyRemove(TKey key) => ApplyRemove(key);
     void IDurableDictionaryCommandHandler<TKey, TValue>.ApplyClear() => ApplyClear();
@@ -136,6 +209,7 @@ void IDurableDictionaryCommandHandler<TKey, TValue>.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<TKey, TValue> item)
         }
 
         WriteRemove(item.Key);
-        _ = ((ICollection<KeyValuePair<TKey, TValue>>)_items).Remove(item);
+        ApplyRemove(item.Key);
         return true;
     }
 
     public IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator() => ((IEnumerable<KeyValuePair<TKey, TValue>>)_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));
 }

src/Orleans.Journaling/Formats/Json/JsonCommandReader.cs

@@ -176,35 +176,48 @@ public bool ReadArrayItem(string operandName)
     public (TFirst? First, TSecond? Second) ReadCurrentPair<TFirst, TSecond>(
         string operandName,
         JsonTypeInfo<TFirst> firstTypeInfo,
-        JsonTypeInfo<TSecond> secondTypeInfo)
+        JsonTypeInfo<TSecond> 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,37 +226,29 @@ public bool ReadArrayItem(string operandName)
         JsonTypeInfo<TFirst> firstTypeInfo,
         JsonTypeInfo<TSecond> 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<TFirst, TSecond>(
         string operandName,
         JsonTypeInfo<TFirst> firstTypeInfo,
         JsonTypeInfo<TSecond> 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)

src/Orleans.Journaling/IDurableDictionaryValueLifecycle.cs

@@ -0,0 +1,34 @@
+namespace Orleans.Journaling;
+
+/// <summary>
+/// Defines independent resource ownership for values stored in durable dictionaries.
+/// </summary>
+/// <typeparam name="TValue">The dictionary value type.</typeparam>
+/// <remarks>
+/// 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.
+/// </remarks>
+public interface IDurableDictionaryValueLifecycle<TValue>
+{
+    /// <summary>
+    /// Acquires an independent owner of the value's resources.
+    /// </summary>
+    /// <param name="value">The borrowed value.</param>
+    /// <returns>The value with independently retained resources.</returns>
+    TValue Retain(TValue value);
+
+    /// <summary>
+    /// Releases the independently owned resources of a value.
+    /// </summary>
+    /// <param name="value">The owned value being retired.</param>
+    void Release(TValue value);
+}

src/Orleans.Journaling/README.md

@@ -169,6 +169,13 @@ For trimming and Native AOT, use `Configure<JsonJournalOptions>(...)` to configu
 
 ## Custom state and standalone ownership
 
+Register `IDurableDictionaryValueLifecycle<TValue>` 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<TState, TImplementation>()` on `IServiceCollection`. Both type arguments
 are reference types, and the implementation implements both `TState` and `IStateMachine`.

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

@@ -75,6 +75,12 @@ public partial interface IDurableDictionaryCommandHandler<TKey, TValue>
         void Reset(int capacityHint);
     }
 
+    public partial interface IDurableDictionaryValueLifecycle<TValue>
+    {
+        void Release(TValue value);
+        TValue Retain(TValue value);
+    }
+
     public partial interface IDurableDictionary<TKey, TValue> : System.Collections.Generic.IDictionary<TKey, TValue>, System.Collections.Generic.ICollection<System.Collections.Generic.KeyValuePair<TKey, TValue>>, System.Collections.Generic.IEnumerable<System.Collections.Generic.KeyValuePair<TKey, TValue>>, System.Collections.IEnumerable
     {
     }