STJ: Dispose enumerator on exception (#100194)

manandre · eiriktsarpalis · web-flow · commit a54d9e9b991c · 2024-06-25T01:58:26.000+01:00
* 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 &lt;eirik.tsarpalis@gmail.com&gt;
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/DictionaryDefaultConverter.cs b/src/libraries/System.Text.Json/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/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/IDictionaryConverter.cs b/src/libraries/System.Text.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/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/IEnumerableConverter.cs b/src/libraries/System.Text.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/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/IEnumerableDefaultConverter.cs b/src/libraries/System.Text.Json/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/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Collection/StackOrQueueConverter.cs b/src/libraries/System.Text.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/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/JsonConverterOfT.WriteCore.cs b/src/libraries/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/src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Metadata/JsonTypeInfoOfT.WriteHelpers.cs b/src/libraries/System.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/src/libraries/System.Text.Json/tests/System.Text.Json.SourceGeneration.Tests/Serialization/CollectionTests.cs b/src/libraries/System.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>>))]

Original file line number	Diff line number	Diff line change
`@@ -28,6 +28,7 @@ protected internal override bool OnWriteResume(`
`28`	`28`	`if (state.Current.CollectionEnumerator == null)`
`29`	`29`	`{`
`30`	`30`	`enumerator = value.GetEnumerator();`
	`31`	`+ state.Current.CollectionEnumerator = enumerator;`
`31`	`32`	`if (!enumerator.MoveNext())`
`32`	`33`	`{`
`33`	`34`	`enumerator.Dispose();`
`@@ -47,7 +48,6 @@ protected internal override bool OnWriteResume(`
`47`	`48`	`{`
`48`	`49`	`if (ShouldFlush(ref state, writer))`
`49`	`50`	`{`
`50`		`- state.Current.CollectionEnumerator = enumerator;`
`51`	`51`	`return false;`
`52`	`52`	`}`
`53`	`53`
`@@ -61,7 +61,6 @@ protected internal override bool OnWriteResume(`
`61`	`61`	`TValue element = enumerator.Current.Value;`
`62`	`62`	`if (!_valueConverter.TryWrite(writer, element, options, ref state))`
`63`	`63`	`{`
`64`		`- state.Current.CollectionEnumerator = enumerator;`
`65`	`64`	`return false;`
`66`	`65`	`}`
`67`	`66`
Original file line number	Diff line number	Diff line change
`@@ -45,6 +45,7 @@ protected internal override bool OnWriteResume(Utf8JsonWriter writer, TDictionar`
`45`	`45`	`if (state.Current.CollectionEnumerator == null)`
`46`	`46`	`{`
`47`	`47`	`enumerator = value.GetEnumerator();`
	`48`	`+ state.Current.CollectionEnumerator = enumerator;`
`48`	`49`	`if (!enumerator.MoveNext())`
`49`	`50`	`{`
`50`	`51`	`return true;`
`@@ -62,7 +63,6 @@ protected internal override bool OnWriteResume(Utf8JsonWriter writer, TDictionar`
`62`	`63`	`{`
`63`	`64`	`if (ShouldFlush(ref state, writer))`
`64`	`65`	`{`
`65`		`- state.Current.CollectionEnumerator = enumerator;`
`66`	`66`	`return false;`
`67`	`67`	`}`
`68`	`68`
`@@ -87,7 +87,6 @@ protected internal override bool OnWriteResume(Utf8JsonWriter writer, TDictionar`
`87`	`87`	`object? element = enumerator.Value;`
`88`	`88`	`if (!_valueConverter.TryWrite(writer, element, options, ref state))`
`89`	`89`	`{`
`90`		`- state.Current.CollectionEnumerator = enumerator;`
`91`	`90`	`return false;`
`92`	`91`	`}`
`93`	`92`
Original file line number	Diff line number	Diff line change
`@@ -46,6 +46,7 @@ protected override bool OnWriteResume(`
`46`	`46`	`if (state.Current.CollectionEnumerator == null)`
`47`	`47`	`{`
`48`	`48`	`enumerator = value.GetEnumerator();`
	`49`	`+ state.Current.CollectionEnumerator = enumerator;`
`49`	`50`	`if (!enumerator.MoveNext())`
`50`	`51`	`{`
`51`	`52`	`return true;`
`@@ -61,14 +62,12 @@ protected override bool OnWriteResume(`
`61`	`62`	`{`
`62`	`63`	`if (ShouldFlush(ref state, writer))`
`63`	`64`	`{`
`64`		`- state.Current.CollectionEnumerator = enumerator;`
`65`	`65`	`return false;`
`66`	`66`	`}`
`67`	`67`
`68`	`68`	`object? element = enumerator.Current;`
`69`	`69`	`if (!converter.TryWrite(writer, element, options, ref state))`
`70`	`70`	`{`
`71`		`- state.Current.CollectionEnumerator = enumerator;`
`72`	`71`	`return false;`
`73`	`72`	`}`
`74`	`73`
Original file line number	Diff line number	Diff line change
`@@ -22,6 +22,7 @@ protected override bool OnWriteResume(Utf8JsonWriter writer, TCollection value,`
`22`	`22`	`if (state.Current.CollectionEnumerator == null)`
`23`	`23`	`{`
`24`	`24`	`enumerator = value.GetEnumerator();`
	`25`	`+ state.Current.CollectionEnumerator = enumerator;`
`25`	`26`	`if (!enumerator.MoveNext())`
`26`	`27`	`{`
`27`	`28`	`enumerator.Dispose();`
`@@ -39,14 +40,12 @@ protected override bool OnWriteResume(Utf8JsonWriter writer, TCollection value,`
`39`	`40`	`{`
`40`	`41`	`if (ShouldFlush(ref state, writer))`
`41`	`42`	`{`
`42`		`- state.Current.CollectionEnumerator = enumerator;`
`43`	`43`	`return false;`
`44`	`44`	`}`
`45`	`45`
`46`	`46`	`TElement element = enumerator.Current;`
`47`	`47`	`if (!converter.TryWrite(writer, element, options, ref state))`
`48`	`48`	`{`
`49`		`- state.Current.CollectionEnumerator = enumerator;`
`50`	`49`	`return false;`
`51`	`50`	`}`
`52`	`51`
Original file line number	Diff line number	Diff line change
`@@ -46,6 +46,7 @@ protected sealed override bool OnWriteResume(Utf8JsonWriter writer, TCollection`
`46`	`46`	`if (state.Current.CollectionEnumerator == null)`
`47`	`47`	`{`
`48`	`48`	`enumerator = value.GetEnumerator();`
	`49`	`+ state.Current.CollectionEnumerator = enumerator;`
`49`	`50`	`if (!enumerator.MoveNext())`
`50`	`51`	`{`
`51`	`52`	`return true;`
`@@ -61,14 +62,12 @@ protected sealed override bool OnWriteResume(Utf8JsonWriter writer, TCollection`
`61`	`62`	`{`
`62`	`63`	`if (ShouldFlush(ref state, writer))`
`63`	`64`	`{`
`64`		`- state.Current.CollectionEnumerator = enumerator;`
`65`	`65`	`return false;`
`66`	`66`	`}`
`67`	`67`
`68`	`68`	`object? element = enumerator.Current;`
`69`	`69`	`if (!converter.TryWrite(writer, element, options, ref state))`
`70`	`70`	`{`
`71`		`- state.Current.CollectionEnumerator = enumerator;`
`72`	`71`	`return false;`
`73`	`72`	`}`
`74`	`73`
Original file line number	Diff line number	Diff line change
`@@ -305,6 +305,7 @@ rootValue is not null &&`
`305`	`305`	`{`
`306`	`306`	`ThrowHelper.ThrowInvalidOperationException_PipeWriterDoesNotImplementUnflushedBytes(bufferWriter);`
`307`	`307`	`}`
	`308`	`+`
`308`	`309`	`state.PipeWriter = bufferWriter;`
`309`	`310`	`state.FlushThreshold = (int)(bufferWriter.Capacity * JsonSerializer.FlushThreshold);`
`310`	`311`
Original file line number	Diff line number	Diff line change
`@@ -586,9 +586,9 @@ public static void ThrowNotSupportedException(scoped ref ReadStack state, in Utf`
`586`	`586`	`}`
`587`	`587`
`588`	`588`	`[DoesNotReturn]`
`589`		`- public static void ThrowNotSupportedException(ref WriteStack state, NotSupportedException ex)`
	`589`	`+ public static void ThrowNotSupportedException(ref WriteStack state, Exception innerException)`
`590`	`590`	`{`
`591`		`- string message = ex.Message;`
	`591`	`+ string message = innerException.Message;`
`592`	`592`
`593`	`593`	`// The caller should check to ensure path is not already set.`
`594`	`594`	`Debug.Assert(!message.Contains(" Path: "));`
`@@ -608,7 +608,7 @@ public static void ThrowNotSupportedException(ref WriteStack state, NotSupported`
`608`	`608`
`609`	`609`	`message += $" Path: {state.PropertyPath()}.";`
`610`	`610`
`611`		`- throw new NotSupportedException(message, ex);`
	`611`	`+ throw new NotSupportedException(message, innerException);`
`612`	`612`	`}`
`613`	`613`
`614`	`614`	`[DoesNotReturn]`