STJ: Dispose enumerator on exception (#100194)

* STJ: Dispose enumerator on exception * Avoid code duplication * Rework fix * Remove useless Disposable field * Apply fix on all collection converters * Remove duplicate assignments * Skip fix for no-op Dispose implementation * Move IEnumerator disposal to WriteCore method. --------- Co-authored-by: Eirik Tsarpalis <eirik.tsarpalis@gmail.com>
dotnet · Jun 25, 2024 · a54d9e9 · a54d9e9
1 parent ddcbc8b
commit a54d9e9
Show file tree

Hide file tree

Showing 11 changed files with 68 additions and 37 deletions.
diff --git a/...on/src/System/Text/Json/Serialization/Converters/Collection/DictionaryDefaultConverter.cs b/...on/src/System/Text/Json/Serialization/Converters/Collection/DictionaryDefaultConverter.cs
@@ -28,6 +28,7 @@ protected internal override bool OnWriteResume(
             if (state.Current.CollectionEnumerator == null)
             {
                 enumerator = value.GetEnumerator();
+                state.Current.CollectionEnumerator = enumerator;
                 if (!enumerator.MoveNext())
                 {
                     enumerator.Dispose();
@@ -47,7 +48,6 @@ protected internal override bool OnWriteResume(
             {
                 if (ShouldFlush(ref state, writer))
                 {
-                    state.Current.CollectionEnumerator = enumerator;
                     return false;
                 }
 
@@ -61,7 +61,6 @@ protected internal override bool OnWriteResume(
                 TValue element = enumerator.Current.Value;
                 if (!_valueConverter.TryWrite(writer, element, options, ref state))
                 {
-                    state.Current.CollectionEnumerator = enumerator;
                     return false;
                 }
 

diff --git a/...ext.Json/src/System/Text/Json/Serialization/Converters/Collection/IDictionaryConverter.cs b/...ext.Json/src/System/Text/Json/Serialization/Converters/Collection/IDictionaryConverter.cs
@@ -45,6 +45,7 @@ protected internal override bool OnWriteResume(Utf8JsonWriter writer, TDictionar
             if (state.Current.CollectionEnumerator == null)
             {
                 enumerator = value.GetEnumerator();
+                state.Current.CollectionEnumerator = enumerator;
                 if (!enumerator.MoveNext())
                 {
                     return true;
@@ -62,7 +63,6 @@ protected internal override bool OnWriteResume(Utf8JsonWriter writer, TDictionar
             {
                 if (ShouldFlush(ref state, writer))
                 {
-                    state.Current.CollectionEnumerator = enumerator;
                     return false;
                 }
 
@@ -87,7 +87,6 @@ protected internal override bool OnWriteResume(Utf8JsonWriter writer, TDictionar
                 object? element = enumerator.Value;
                 if (!_valueConverter.TryWrite(writer, element, options, ref state))
                 {
-                    state.Current.CollectionEnumerator = enumerator;
                     return false;
                 }
 

diff --git a/...ext.Json/src/System/Text/Json/Serialization/Converters/Collection/IEnumerableConverter.cs b/...ext.Json/src/System/Text/Json/Serialization/Converters/Collection/IEnumerableConverter.cs
@@ -46,6 +46,7 @@ protected override bool OnWriteResume(
             if (state.Current.CollectionEnumerator == null)
             {
                 enumerator = value.GetEnumerator();
+                state.Current.CollectionEnumerator = enumerator;
                 if (!enumerator.MoveNext())
                 {
                     return true;
@@ -61,14 +62,12 @@ protected override bool OnWriteResume(
             {
                 if (ShouldFlush(ref state, writer))
                 {
-                    state.Current.CollectionEnumerator = enumerator;
                     return false;
                 }
 
                 object? element = enumerator.Current;
                 if (!converter.TryWrite(writer, element, options, ref state))
                 {
-                    state.Current.CollectionEnumerator = enumerator;
                     return false;
                 }
 

diff --git a/...n/src/System/Text/Json/Serialization/Converters/Collection/IEnumerableDefaultConverter.cs b/...n/src/System/Text/Json/Serialization/Converters/Collection/IEnumerableDefaultConverter.cs
@@ -22,6 +22,7 @@ protected override bool OnWriteResume(Utf8JsonWriter writer, TCollection value,
             if (state.Current.CollectionEnumerator == null)
             {
                 enumerator = value.GetEnumerator();
+                state.Current.CollectionEnumerator = enumerator;
                 if (!enumerator.MoveNext())
                 {
                     enumerator.Dispose();
@@ -39,14 +40,12 @@ protected override bool OnWriteResume(Utf8JsonWriter writer, TCollection value,
             {
                 if (ShouldFlush(ref state, writer))
                 {
-                    state.Current.CollectionEnumerator = enumerator;
                     return false;
                 }
 
                 TElement element = enumerator.Current;
                 if (!converter.TryWrite(writer, element, options, ref state))
                 {
-                    state.Current.CollectionEnumerator = enumerator;
                     return false;
                 }
 

diff --git a/...xt.Json/src/System/Text/Json/Serialization/Converters/Collection/StackOrQueueConverter.cs b/...xt.Json/src/System/Text/Json/Serialization/Converters/Collection/StackOrQueueConverter.cs
@@ -46,6 +46,7 @@ protected sealed override bool OnWriteResume(Utf8JsonWriter writer, TCollection
             if (state.Current.CollectionEnumerator == null)
             {
                 enumerator = value.GetEnumerator();
+                state.Current.CollectionEnumerator = enumerator;
                 if (!enumerator.MoveNext())
                 {
                     return true;
@@ -61,14 +62,12 @@ protected sealed override bool OnWriteResume(Utf8JsonWriter writer, TCollection
             {
                 if (ShouldFlush(ref state, writer))
                 {
-                    state.Current.CollectionEnumerator = enumerator;
                     return false;
                 }
 
                 object? element = enumerator.Current;
                 if (!converter.TryWrite(writer, element, options, ref state))
                 {
-                    state.Current.CollectionEnumerator = enumerator;
                     return false;
                 }
 

diff --git a/...braries/System.Text.Json/src/System/Text/Json/Serialization/JsonConverterOfT.WriteCore.cs b/...braries/System.Text.Json/src/System/Text/Json/Serialization/JsonConverterOfT.WriteCore.cs
@@ -15,32 +15,37 @@ internal bool WriteCore(
             {
                 return TryWrite(writer, value, options, ref state);
             }
-            catch (InvalidOperationException ex) when (ex.Source == ThrowHelper.ExceptionSourceValueToRethrowAsJsonException)
+            catch (Exception ex)
             {
-                ThrowHelper.ReThrowWithPath(ref state, ex);
-                throw;
-            }
-            catch (JsonException ex) when (ex.Path == null)
-            {
-                // JsonExceptions where the Path property is already set
-                // typically originate from nested calls to JsonSerializer;
-                // treat these cases as any other exception type and do not
-                // overwrite any exception information.
+                if (!state.SupportAsync)
+                {
+                    // Async serializers should dispose sync and
+                    // async disposables from the async root method.
+                    state.DisposePendingDisposablesOnException();
+                }
 
-                ThrowHelper.AddJsonExceptionInformation(ref state, ex);
-                throw;
-            }
-            catch (NotSupportedException ex)
-            {
-                // If the message already contains Path, just re-throw. This could occur in serializer re-entry cases.
-                // To get proper Path semantics in re-entry cases, APIs that take 'state' need to be used.
-                if (ex.Message.Contains(" Path: "))
+                switch (ex)
                 {
-                    throw;
+                    case InvalidOperationException when ex.Source == ThrowHelper.ExceptionSourceValueToRethrowAsJsonException:
+                        ThrowHelper.ReThrowWithPath(ref state, ex);
+                        break;
+
+                    case JsonException { Path: null } jsonException:
+                        // JsonExceptions where the Path property is already set
+                        // typically originate from nested calls to JsonSerializer;
+                        // treat these cases as any other exception type and do not
+                        // overwrite any exception information.
+                        ThrowHelper.AddJsonExceptionInformation(ref state, jsonException);
+                        break;
+
+                    case NotSupportedException when !ex.Message.Contains(" Path: "):
+                        // If the message already contains Path, just re-throw. This could occur in serializer re-entry cases.
+                        // To get proper Path semantics in re-entry cases, APIs that take 'state' need to be used.
+                        ThrowHelper.ThrowNotSupportedException(ref state, ex);
+                        break;
                 }
 
-                ThrowHelper.ThrowNotSupportedException(ref state, ex);
-                return default;
+                throw;
             }
         }
     }

diff --git a/...tem.Text.Json/src/System/Text/Json/Serialization/Metadata/JsonTypeInfoOfT.WriteHelpers.cs b/...tem.Text.Json/src/System/Text/Json/Serialization/Metadata/JsonTypeInfoOfT.WriteHelpers.cs
@@ -305,6 +305,7 @@ rootValue is not null &&
                 {
                     ThrowHelper.ThrowInvalidOperationException_PipeWriterDoesNotImplementUnflushedBytes(bufferWriter);
                 }
+
                 state.PipeWriter = bufferWriter;
                 state.FlushThreshold = (int)(bufferWriter.Capacity * JsonSerializer.FlushThreshold);
 

diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/WriteStack.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/WriteStack.cs
@@ -278,7 +278,7 @@ public void AddCompletedAsyncDisposable(IAsyncDisposable asyncDisposable)
             => (CompletedAsyncDisposables ??= new List<IAsyncDisposable>()).Add(asyncDisposable);
 
         // Asynchronously dispose of any AsyncDisposables that have been scheduled for disposal
-        public async ValueTask DisposeCompletedAsyncDisposables()
+        public readonly async ValueTask DisposeCompletedAsyncDisposables()
         {
             Debug.Assert(CompletedAsyncDisposables?.Count > 0);
             Exception? exception = null;
@@ -307,7 +307,7 @@ public async ValueTask DisposeCompletedAsyncDisposables()
         /// Walks the stack cleaning up any leftover IDisposables
         /// in the event of an exception on serialization
         /// </summary>
-        public void DisposePendingDisposablesOnException()
+        public readonly void DisposePendingDisposablesOnException()
         {
             Exception? exception = null;
 
@@ -346,7 +346,7 @@ static void DisposeFrame(IEnumerator? collectionEnumerator, ref Exception? excep
         /// Walks the stack cleaning up any leftover I(Async)Disposables
         /// in the event of an exception on async serialization
         /// </summary>
-        public async ValueTask DisposePendingDisposablesOnExceptionAsync()
+        public readonly async ValueTask DisposePendingDisposablesOnExceptionAsync()
         {
             Exception? exception = null;
 

diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/ThrowHelper.Serialization.cs b/src/libraries/System.Text.Json/src/System/Text/Json/ThrowHelper.Serialization.cs
@@ -586,9 +586,9 @@ public static void ThrowNotSupportedException(scoped ref ReadStack state, in Utf
         }
 
         [DoesNotReturn]
-        public static void ThrowNotSupportedException(ref WriteStack state, NotSupportedException ex)
+        public static void ThrowNotSupportedException(ref WriteStack state, Exception innerException)
         {
-            string message = ex.Message;
+            string message = innerException.Message;
 
             // The caller should check to ensure path is not already set.
             Debug.Assert(!message.Contains(" Path: "));
@@ -608,7 +608,7 @@ public static void ThrowNotSupportedException(ref WriteStack state, NotSupported
 
             message += $" Path: {state.PropertyPath()}.";
 
-            throw new NotSupportedException(message, ex);
+            throw new NotSupportedException(message, innerException);
         }
 
         [DoesNotReturn]

diff --git a/src/libraries/System.Text.Json/tests/Common/CollectionTests/CollectionTests.Generic.Write.cs b/src/libraries/System.Text.Json/tests/Common/CollectionTests/CollectionTests.Generic.Write.cs
@@ -955,6 +955,34 @@ public async Task WriteISetT_DisposesEnumerators()
             }
         }
 
+        [Fact]
+        public async Task WriteIEnumerableT_ElementSerializationThrows_DisposesEnumerators()
+        {
+            var items = new RefCountedList<IEnumerable<int>>(Enumerable.Repeat(ThrowingEnumerable(), 1));
+            await Assert.ThrowsAsync<DivideByZeroException>(() => Serializer.SerializeWrapper(items.AsEnumerable()));
+            Assert.Equal(0, items.RefCount);
+
+            static IEnumerable<int> ThrowingEnumerable()
+            {
+                yield return 42;
+                throw new DivideByZeroException();
+            }
+        }
+
+        [Fact]
+        public async Task WriteIDictionaryT_ElementSerializationThrows_DisposesEnumerators()
+        {
+            var items = new RefCountedDictionary<int, IEnumerable<int>>(Enumerable.Repeat(new KeyValuePair<int, IEnumerable<int>>(42, ThrowingEnumerable()), 1));
+            await Assert.ThrowsAsync<DivideByZeroException>(() => Serializer.SerializeWrapper((IDictionary<int, IEnumerable<int>>)items));
+            Assert.Equal(0, items.RefCount);
+
+            static IEnumerable<int> ThrowingEnumerable()
+            {
+                yield return 42;
+                throw new DivideByZeroException();
+            }
+        }
+
         public class SimpleClassWithKeyValuePairs
         {
             public KeyValuePair<string, string> KvpWStrVal { get; set; }

diff --git a/....Text.Json/tests/System.Text.Json.SourceGeneration.Tests/Serialization/CollectionTests.cs b/....Text.Json/tests/System.Text.Json.SourceGeneration.Tests/Serialization/CollectionTests.cs
@@ -130,6 +130,7 @@ public async Task DeserializeAsyncEnumerable()
         [JsonSerializable(typeof(JsonElement))]
         [JsonSerializable(typeof(string))]
         [JsonSerializable(typeof(IDictionary<int, int>))]
+        [JsonSerializable(typeof(IDictionary<int, IEnumerable<int>>))]
         [JsonSerializable(typeof(Dictionary<string, ClassWithInternalParameterlessConstructor>))]
         [JsonSerializable(typeof(Dictionary<string, ClassWithPrivateParameterlessConstructor>))]
         [JsonSerializable(typeof(Dictionary<string, Dictionary<string, CustomClass>>))]
@@ -556,6 +557,7 @@ public CollectionTests_Default()
         [JsonSerializable(typeof(JsonElement))]
         [JsonSerializable(typeof(string))]
         [JsonSerializable(typeof(IDictionary<int, int>))]
+        [JsonSerializable(typeof(IDictionary<int, IEnumerable<int>>))]
         [JsonSerializable(typeof(Dictionary<string, ClassWithInternalParameterlessConstructor>))]
         [JsonSerializable(typeof(Dictionary<string, ClassWithPrivateParameterlessConstructor>))]
         [JsonSerializable(typeof(Dictionary<string, Dictionary<string, CustomClass>>))]