From 753110147941a405219880c87611bfdb2588a592 Mon Sep 17 00:00:00 2001
From: Eirik Tsarpalis <eirik.tsarpalis@gmail.com>
Date: Mon, 12 Dec 2022 21:30:36 +0000
Subject: [PATCH] Fix #79556.

---
 .../Metadata/JsonTypeInfoOfT.WriteHelpers.cs  |  2 ++
 .../CollectionTests.AsyncEnumerable.cs        | 23 +++++++++++++++++++
 .../tests/Common/Utf8MemoryStream.cs          | 14 ++++++++++-
 3 files changed, 38 insertions(+), 1 deletion(-)

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
index 75b0b3f3e43809..db6dde95e15bce 100644
--- 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
@@ -121,6 +121,8 @@ rootValue is not null &&
                     supportContinuation: true,
                     supportAsync: true);
 
+                state.CancellationToken = cancellationToken;
+
                 using var bufferWriter = new PooledByteBufferWriter(Options.DefaultBufferSize);
                 using var writer = new Utf8JsonWriter(bufferWriter, Options.GetWriterOptions());
 
diff --git a/src/libraries/System.Text.Json/tests/Common/CollectionTests/CollectionTests.AsyncEnumerable.cs b/src/libraries/System.Text.Json/tests/Common/CollectionTests/CollectionTests.AsyncEnumerable.cs
index 52f3f2f32e84c1..0ae7c05046fd5d 100644
--- a/src/libraries/System.Text.Json/tests/Common/CollectionTests/CollectionTests.AsyncEnumerable.cs
+++ b/src/libraries/System.Text.Json/tests/Common/CollectionTests/CollectionTests.AsyncEnumerable.cs
@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
@@ -116,6 +117,28 @@ public async Task WriteNestedAsyncEnumerable_Nullable<TElement>(IEnumerable<TEle
             Assert.Equal(1, asyncEnumerable.TotalDisposedEnumerators);
         }
 
+        [Fact]
+        public async Task WriteAsyncEnumerable_CancellationToken_IsPassedToAsyncEnumerator()
+        {
+            // Regression test for https://github.com/dotnet/runtime/issues/79556
+            using var utf8Stream = new Utf8MemoryStream(ignoreCancellationTokenOnWriteAsync: true);
+            using var cts = new CancellationTokenSource();
+
+            IAsyncEnumerable<int> value = CreateEnumerable();
+            await JsonSerializer.SerializeAsync(utf8Stream, value, cancellationToken: cts.Token);
+            Assert.Equal("[1,2]", utf8Stream.AsString());
+
+            async IAsyncEnumerable<int> CreateEnumerable([EnumeratorCancellation] CancellationToken cancellationToken = default)
+            {
+                yield return 1;
+                await Task.Delay(20);
+                Assert.False(cancellationToken.IsCancellationRequested);
+                cts.Cancel();
+                Assert.True(cancellationToken.IsCancellationRequested);
+                yield return 2;
+            }
+        }
+
         [Theory, OuterLoop]
         [InlineData(5000, 1000, true)]
         [InlineData(5000, 1000, false)]
diff --git a/src/libraries/System.Text.Json/tests/Common/Utf8MemoryStream.cs b/src/libraries/System.Text.Json/tests/Common/Utf8MemoryStream.cs
index 1faa96058723eb..2904d744aa949a 100644
--- a/src/libraries/System.Text.Json/tests/Common/Utf8MemoryStream.cs
+++ b/src/libraries/System.Text.Json/tests/Common/Utf8MemoryStream.cs
@@ -2,19 +2,31 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace System.Text.Json.Serialization.Tests
 {
     public sealed class Utf8MemoryStream : MemoryStream
     {
-        public Utf8MemoryStream() : base()
+        private readonly bool _ignoreCancellationTokenOnWriteAsync;
+
+        public Utf8MemoryStream(bool ignoreCancellationTokenOnWriteAsync = false) : base()
         {
+            _ignoreCancellationTokenOnWriteAsync = ignoreCancellationTokenOnWriteAsync;
         }
 
         public Utf8MemoryStream(string text) : base(Encoding.UTF8.GetBytes(text))
         {
         }
 
+#if NETCOREAPP
+        public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
+            => base.WriteAsync(buffer, _ignoreCancellationTokenOnWriteAsync ? default : cancellationToken);
+#endif
+        public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+            => base.WriteAsync(buffer, offset, count, _ignoreCancellationTokenOnWriteAsync ? default : cancellationToken);
+
         public string AsString() => Encoding.UTF8.GetString(ToArray());
     }
 }