Fix enumerator disposal when cancellation outside of serialization occurs.

src/libraries/System.Text.Json/src/System/Text/Json/Serialization/WriteStack.cs

@@ -183,6 +183,8 @@ public readonly JsonTypeInfo PeekNestedJsonTypeInfo()
 
         public void Push()
         {
+            Debug.Assert(_continuationCount == 0 || _count < _continuationCount);
+
             if (_continuationCount == 0)
             {
                 Debug.Assert(Current.PolymorphicSerializationState != PolymorphicSerializationState.PolymorphicReEntrySuspended);
@@ -234,6 +236,7 @@ public void Push()
         public void Pop(bool success)
         {
             Debug.Assert(_count > 0);
+            Debug.Assert(_continuationCount == 0 || _count < _continuationCount);
 
             if (!success)
             {
@@ -314,11 +317,23 @@ public readonly void DisposePendingDisposablesOnException()
             Debug.Assert(Current.AsyncDisposable is null);
             DisposeFrame(Current.CollectionEnumerator, ref exception);
 
-            int stackSize = Math.Max(_count, _continuationCount);
-            for (int i = 0; i < stackSize - 1; i++)
+            if (_stack is not null)
             {
-                Debug.Assert(_stack[i].AsyncDisposable is null);
-                DisposeFrame(_stack[i].CollectionEnumerator, ref exception);
+                int currentIndex = _count - _indexOffset;
+                int stackSize = Math.Max(currentIndex, _continuationCount);
+                for (int i = 0; i < stackSize; i++)
+                {
+                    Debug.Assert(_stack[i].AsyncDisposable is null);
+
+                    if (i == currentIndex)
+                    {
+                        // Matches the entry in Current, skip to avoid double disposal.
+                        Debug.Assert(_stack[i].CollectionEnumerator is null || ReferenceEquals(Current.CollectionEnumerator, _stack[i].CollectionEnumerator));
+                        continue;
+                    }
+
+                    DisposeFrame(_stack[i].CollectionEnumerator, ref exception);
+                }
             }
 
             if (exception is not null)
@@ -352,10 +367,23 @@ public readonly async ValueTask DisposePendingDisposablesOnExceptionAsync()
 
             exception = await DisposeFrame(Current.CollectionEnumerator, Current.AsyncDisposable, exception).ConfigureAwait(false);
 
-            int stackSize = Math.Max(_count, _continuationCount);
-            for (int i = 0; i < stackSize - 1; i++)
+            if (_stack is not null)
             {
-                exception = await DisposeFrame(_stack[i].CollectionEnumerator, _stack[i].AsyncDisposable, exception).ConfigureAwait(false);
+                Debug.Assert(_continuationCount == 0 || _count < _continuationCount);
+                int currentIndex = _count - _indexOffset;
+                int stackSize = Math.Max(currentIndex, _continuationCount);
+                for (int i = 0; i < stackSize; i++)
+                {
+                    if (i == currentIndex)
+                    {
+                        // Matches the entry in Current, skip to avoid double disposal.
+                        Debug.Assert(_stack[i].CollectionEnumerator is null || ReferenceEquals(Current.CollectionEnumerator, _stack[i].CollectionEnumerator));
+                        Debug.Assert(_stack[i].AsyncDisposable is null || ReferenceEquals(Current.AsyncDisposable, _stack[i].AsyncDisposable));
+                        continue;
+                    }
+
+                    exception = await DisposeFrame(_stack[i].CollectionEnumerator, _stack[i].AsyncDisposable, exception).ConfigureAwait(false);
+                }
             }
 
             if (exception is not null)

src/libraries/System.Text.Json/tests/Common/AsyncEnumerableTests.cs

@@ -387,6 +387,32 @@ public static IEnumerable<object[]> GetAsyncEnumerableSources()
             static object[] WrapArgs<TSource>(IEnumerable<TSource> source, int bufferSize, DeserializeAsyncEnumerableOverload overload) => new object[] { source, bufferSize, overload };
         }
 
+        [Theory]
+        [InlineData(0)]
+        [InlineData(1)]
+        [InlineData(2)]
+        [InlineData(3)]
+        [InlineData(5)]
+        [InlineData(43)]
+        public async Task SerializeAsyncEnumerable_Cancellation_DisposesEnumerators(int depth)
+        {
+            // Regression test for https://github.com/dotnet/runtime/issues/120010
+
+            using SelfCancellingAsyncEnumerable enumerable = new();
+            using MemoryStream stream = new MemoryStream();
+
+            object wrappingValue = enumerable;
+            while (depth-- > 0)
+            {
+                // Use a LINQ enumerable instead of array/list
+                // to force use of enumerators in every layer.
+                wrappingValue = Enumerable.Repeat(wrappingValue, 1);
+            }
+
+            await Assert.ThrowsAsync<TaskCanceledException>(() => StreamingSerializer.SerializeWrapper(stream, wrappingValue, cancellationToken: enumerable.CancellationToken));
+            Assert.True(enumerable.IsEnumeratorDisposed);
+        }
+
         public enum DeserializeAsyncEnumerableOverload { JsonSerializerOptions, JsonTypeInfo };
 
         private IAsyncEnumerable<T> DeserializeAsyncEnumerableWrapper<T>(Stream stream, JsonSerializerOptions options = null, CancellationToken cancellationToken = default, DeserializeAsyncEnumerableOverload overload = DeserializeAsyncEnumerableOverload.JsonSerializerOptions)
@@ -457,5 +483,35 @@ public override string Read(ref Utf8JsonReader reader, Type typeToConvert, JsonS
 
             public override void Write(Utf8JsonWriter writer, string value, JsonSerializerOptions options) => throw new NotImplementedException();
         }
+
+        sealed class SelfCancellingAsyncEnumerable : IAsyncEnumerable<int>, IDisposable
+        {
+            private readonly CancellationTokenSource _cts = new();
+            public CancellationToken CancellationToken => _cts.Token;
+            public bool IsEnumeratorDisposed { get; private set; }
+            public IAsyncEnumerator<int> GetAsyncEnumerator(CancellationToken _) => new Enumerator(this);
+            private sealed class Enumerator(SelfCancellingAsyncEnumerable parent) : IAsyncEnumerator<int>
+            {
+                public int Current { get; private set; }
+                public async ValueTask<bool> MoveNextAsync()
+                {
+                    await Task.Yield();
+                    if (++Current == 10)
+                    {
+                        parent._cts.Cancel();
+                    }
+
+                    return true;
+                }
+
+                public ValueTask DisposeAsync()
+                {
+                    parent.IsEnumeratorDisposed = true;
+                    return default;
+                }
+            }
+
+            public void Dispose() => _cts.Dispose();
+        }
     }
 }

src/libraries/System.Text.Json/tests/Common/CollectionTests/CollectionTests.Generic.Write.cs

@@ -1,8 +1,10 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Collections;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
 
@@ -983,6 +985,66 @@ static IEnumerable<int> ThrowingEnumerable()
             }
         }
 
+        [Theory]
+        [InlineData(0)]
+        [InlineData(1)]
+        [InlineData(2)]
+        [InlineData(3)]
+        [InlineData(5)]
+        [InlineData(43)]
+        public async Task WriteIEnumerableOfT_Cancellation_DisposesEnumerators(int depth)
+        {
+            // Regression test for https://github.com/dotnet/runtime/issues/120010
+
+            if (StreamingSerializer?.IsAsyncSerializer is not true)
+            {
+                return; // require serializers with cancellation support.
+            }
+
+            JsonSerializerOptions options = Serializer.CreateOptions(opts => opts.DefaultBufferSize = 1); // Force early async writes
+            using SelfCancellingEnumerable enumerable = new();
+            using Utf8MemoryStream stream = new();
+
+            object wrappingValue = enumerable;
+            while (depth-- > 0)
+            {
+                // Use a LINQ enumerable instead of array/list
+                // to force use of enumerators in every layer.
+                wrappingValue = Enumerable.Repeat(wrappingValue, 1);
+            }
+
+            await Assert.ThrowsAsync<TaskCanceledException>(() => StreamingSerializer.SerializeWrapper(stream, wrappingValue, options, enumerable.CancellationToken));
+            Assert.True(enumerable.IsEnumeratorDisposed);
+        }
+
+        public sealed class SelfCancellingEnumerable : IEnumerable<int>, IDisposable
+        {
+            private readonly CancellationTokenSource _cts = new();
+            public bool IsEnumeratorDisposed { get; private set; }
+            public CancellationToken CancellationToken => _cts.Token;
+            public IEnumerator<int> GetEnumerator() => new Enumerator(this);
+            IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+            public void Dispose() => _cts.Dispose();
+
+            private sealed class Enumerator(SelfCancellingEnumerable parent) : IEnumerator<int>
+            {
+                public int Current { get; private set; }
+                object IEnumerator.Current => Current;
+                public bool MoveNext()
+                {
+                    if (++Current == 10)
+                    {
+                        parent._cts.Cancel();
+                    }
+
+                    return true;
+                }
+
+                public void Dispose() => parent.IsEnumeratorDisposed = true;
+                public void Reset() { }
+            }
+        }
+
         public class SimpleClassWithKeyValuePairs
         {
             public KeyValuePair<string, string> KvpWStrVal { get; set; }

src/libraries/System.Text.Json/tests/Common/PipeJsonSerializerWrapper.cs

@@ -4,6 +4,7 @@
 using System.Buffers;
 using System.IO.Pipelines;
 using System.Text.Json.Serialization.Metadata;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace System.Text.Json.Serialization.Tests
@@ -13,16 +14,16 @@ namespace System.Text.Json.Serialization.Tests
     /// </summary>
     public abstract partial class PipeJsonSerializerWrapper : StreamingJsonSerializerWrapper
     {
-        public abstract Task SerializeWrapper(PipeWriter stream, object value, Type inputType, JsonSerializerOptions? options = null);
-        public abstract Task SerializeWrapper<T>(PipeWriter stream, T value, JsonSerializerOptions? options = null);
-        public abstract Task SerializeWrapper(PipeWriter stream, object value, Type inputType, JsonSerializerContext context);
-        public abstract Task SerializeWrapper<T>(PipeWriter stream, T value, JsonTypeInfo<T> jsonTypeInfo);
-        public abstract Task SerializeWrapper(PipeWriter stream, object value, JsonTypeInfo jsonTypeInfo);
-        public abstract Task<object> DeserializeWrapper(PipeReader utf8Json, Type returnType, JsonSerializerOptions? options = null);
-        public abstract Task<T> DeserializeWrapper<T>(PipeReader utf8Json, JsonSerializerOptions? options = null);
-        public abstract Task<object> DeserializeWrapper(PipeReader utf8Json, Type returnType, JsonSerializerContext context);
-        public abstract Task<T> DeserializeWrapper<T>(PipeReader utf8Json, JsonTypeInfo<T> jsonTypeInfo);
-        public abstract Task<object> DeserializeWrapper(PipeReader utf8Json, JsonTypeInfo jsonTypeInfo);
+        public abstract Task SerializeWrapper(PipeWriter stream, object value, Type inputType, JsonSerializerOptions? options = null, CancellationToken cancellationToken = default);
+        public abstract Task SerializeWrapper<T>(PipeWriter stream, T value, JsonSerializerOptions? options = null, CancellationToken cancellationToken = default);
+        public abstract Task SerializeWrapper(PipeWriter stream, object value, Type inputType, JsonSerializerContext context, CancellationToken cancellationToken = default);
+        public abstract Task SerializeWrapper<T>(PipeWriter stream, T value, JsonTypeInfo<T> jsonTypeInfo, CancellationToken cancellationToken = default);
+        public abstract Task SerializeWrapper(PipeWriter stream, object value, JsonTypeInfo jsonTypeInfo, CancellationToken cancellationToken = default);
+        public abstract Task<object> DeserializeWrapper(PipeReader utf8Json, Type returnType, JsonSerializerOptions? options = null, CancellationToken cancellationToken = default);
+        public abstract Task<T> DeserializeWrapper<T>(PipeReader utf8Json, JsonSerializerOptions? options = null, CancellationToken cancellationToken = default);
+        public abstract Task<object> DeserializeWrapper(PipeReader utf8Json, Type returnType, JsonSerializerContext context, CancellationToken cancellationToken = default);
+        public abstract Task<T> DeserializeWrapper<T>(PipeReader utf8Json, JsonTypeInfo<T> jsonTypeInfo, CancellationToken cancellationToken = default);
+        public abstract Task<object> DeserializeWrapper(PipeReader utf8Json, JsonTypeInfo jsonTypeInfo, CancellationToken cancellationToken = default);
 
         public override async Task<string> SerializeWrapper(object value, Type inputType, JsonSerializerOptions options = null)
         {

src/libraries/System.Text.Json/tests/Common/StreamingJsonSerializerWrapper.cs

@@ -3,6 +3,7 @@
 
 using System.IO;
 using System.Text.Json.Serialization.Metadata;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace System.Text.Json.Serialization.Tests
@@ -18,16 +19,16 @@ public abstract partial class StreamingJsonSerializerWrapper : JsonSerializerWra
         public abstract bool IsAsyncSerializer { get; }
         public virtual bool ForceSmallBufferInOptions { get; } = false;
 
-        public abstract Task SerializeWrapper(Stream stream, object value, Type inputType, JsonSerializerOptions? options = null);
-        public abstract Task SerializeWrapper<T>(Stream stream, T value, JsonSerializerOptions? options = null);
-        public abstract Task SerializeWrapper(Stream stream, object value, Type inputType, JsonSerializerContext context);
-        public abstract Task SerializeWrapper<T>(Stream stream, T value, JsonTypeInfo<T> jsonTypeInfo);
-        public abstract Task SerializeWrapper(Stream stream, object value, JsonTypeInfo jsonTypeInfo);
-        public abstract Task<object> DeserializeWrapper(Stream utf8Json, Type returnType, JsonSerializerOptions? options = null);
-        public abstract Task<T> DeserializeWrapper<T>(Stream utf8Json, JsonSerializerOptions? options = null);
-        public abstract Task<object> DeserializeWrapper(Stream utf8Json, Type returnType, JsonSerializerContext context);
-        public abstract Task<T> DeserializeWrapper<T>(Stream utf8Json, JsonTypeInfo<T> jsonTypeInfo);
-        public abstract Task<object> DeserializeWrapper(Stream utf8Json, JsonTypeInfo jsonTypeInfo);
+        public abstract Task SerializeWrapper(Stream stream, object value, Type inputType, JsonSerializerOptions? options = null, CancellationToken cancellationToken = default);
+        public abstract Task SerializeWrapper<T>(Stream stream, T value, JsonSerializerOptions? options = null, CancellationToken cancellationToken = default);
+        public abstract Task SerializeWrapper(Stream stream, object value, Type inputType, JsonSerializerContext context, CancellationToken cancellationToken = default);
+        public abstract Task SerializeWrapper<T>(Stream stream, T value, JsonTypeInfo<T> jsonTypeInfo, CancellationToken cancellationToken = default);
+        public abstract Task SerializeWrapper(Stream stream, object value, JsonTypeInfo jsonTypeInfo, CancellationToken cancellationToken = default);
+        public abstract Task<object> DeserializeWrapper(Stream utf8Json, Type returnType, JsonSerializerOptions? options = null, CancellationToken cancellationToken = default);
+        public abstract Task<T> DeserializeWrapper<T>(Stream utf8Json, JsonSerializerOptions? options = null, CancellationToken cancellationToken = default);
+        public abstract Task<object> DeserializeWrapper(Stream utf8Json, Type returnType, JsonSerializerContext context, CancellationToken cancellationToken = default);
+        public abstract Task<T> DeserializeWrapper<T>(Stream utf8Json, JsonTypeInfo<T> jsonTypeInfo, CancellationToken cancellationToken = default);
+        public abstract Task<object> DeserializeWrapper(Stream utf8Json, JsonTypeInfo jsonTypeInfo, CancellationToken cancellationToken = default);
 
         public override async Task<string> SerializeWrapper(object value, Type inputType, JsonSerializerOptions options = null)
         {

src/libraries/System.Text.Json/tests/System.Text.Json.SourceGeneration.Tests/Serialization/CollectionTests.cs

@@ -452,6 +452,7 @@ public async Task DeserializeAsyncEnumerable()
         [JsonSerializable(typeof(KeyValuePair<string, KeyValuePair<string, string>>))]
         [JsonSerializable(typeof(KeyValuePair<string, KeyValuePair<string, object>>))]
         [JsonSerializable(typeof(KeyValuePair<string, KeyValuePair<string, SimpleClassWithKeyValuePairs>>))]
+        [JsonSerializable(typeof(SelfCancellingEnumerable))]
         [JsonSerializable(typeof(SimpleClassWithKeyValuePairs))]
         [JsonSerializable(typeof(KeyNameNullPolicy))]
         [JsonSerializable(typeof(ValueNameNullPolicy))]
@@ -873,6 +874,7 @@ public CollectionTests_Default()
         [JsonSerializable(typeof(KeyValuePair<string, KeyValuePair<string, string>>))]
         [JsonSerializable(typeof(KeyValuePair<string, KeyValuePair<string, object>>))]
         [JsonSerializable(typeof(KeyValuePair<string, KeyValuePair<string, SimpleClassWithKeyValuePairs>>))]
+        [JsonSerializable(typeof(SelfCancellingEnumerable))]
         [JsonSerializable(typeof(SimpleClassWithKeyValuePairs))]
         [JsonSerializable(typeof(KeyNameNullPolicy))]
         [JsonSerializable(typeof(ValueNameNullPolicy))]