Lazy async statemachine in WriteToAsync

Src/Newtonsoft.Json/JsonWriter.Async.cs

@@ -275,32 +275,120 @@ namespace Newtonsoft.Json
             }
         }
 
-        internal async Task InternalWriteEndAsync(JsonContainerType type, CancellationToken cancellationToken)
+        internal Task InternalWriteEndAsync(JsonContainerType type, CancellationToken cancellationToken)
         {
-            cancellationToken.ThrowIfCancellationRequested();
-            int levelsToComplete = CalculateLevelsToComplete(type);
+            if (cancellationToken.IsCancellationRequested)
+            {
+                return cancellationToken.FromCanceled();
+            }
 
+            int levelsToComplete = CalculateLevelsToComplete(type);
+            Task t;
             while (levelsToComplete-- > 0)
             {
                 JsonToken token = GetCloseTokenForType(Pop());
 
                 if (_currentState == State.Property)
                 {
-                    await WriteNullAsync(cancellationToken).ConfigureAwait(false);
+                    t = WriteNullAsync(cancellationToken);
+                    if (t.Status != TaskStatus.RanToCompletion)
+                    {
+                        return AwaitProperty(t, levelsToComplete, token, type, cancellationToken);
+                    }
                 }
 
                 if (_formatting == Formatting.Indented)
                 {
                     if (_currentState != State.ObjectStart && _currentState != State.ArrayStart)
                     {
-                        await WriteIndentAsync(cancellationToken).ConfigureAwait(false);
+                        t = WriteIndentAsync(cancellationToken);
+                        if (t.Status != TaskStatus.RanToCompletion)
+                        {
+                            return AwaitIndent(t, levelsToComplete, token, type, cancellationToken);
+                        }
                     }
                 }
 
-                await WriteEndAsync(token, cancellationToken).ConfigureAwait(false);
+                t = WriteEndAsync(token, cancellationToken);
+                if (t.Status != TaskStatus.RanToCompletion)
+                {
+                    return AwaitEnd(t, levelsToComplete, type, cancellationToken);
+                }
 
                 UpdateCurrentState();
             }
+
+            return AsyncUtils.CompletedTask;
+
+            // Local functions, params renamed (capitalized) so as not to capture and allocate when calling async
+            async Task AwaitProperty(Task task, int LevelsToComplete, JsonToken token, JsonContainerType Type, CancellationToken CancellationToken)
+            {
+                await task;
+
+                //  Finish current loop
+                if (_formatting == Formatting.Indented)
+                {
+                    if (_currentState != State.ObjectStart && _currentState != State.ArrayStart)
+                    {
+                        await WriteIndentAsync(CancellationToken).ConfigureAwait(false);
+                    }
+                }
+
+                await WriteEndAsync(token, CancellationToken).ConfigureAwait(false);
+
+                UpdateCurrentState();
+
+                await AwaitRemaining(LevelsToComplete, Type, CancellationToken).ConfigureAwait(false);
+            }
+
+            async Task AwaitIndent(Task task, int LevelsToComplete, JsonToken token, JsonContainerType Type, CancellationToken CancellationToken)
+            {
+                await task;
+
+                //  Finish current loop
+
+                await WriteEndAsync(token, CancellationToken).ConfigureAwait(false);
+
+                UpdateCurrentState();
+
+                await AwaitRemaining(LevelsToComplete, Type, CancellationToken).ConfigureAwait(false);
+            }
+
+            async Task AwaitEnd(Task task, int LevelsToComplete, JsonContainerType Type, CancellationToken CancellationToken)
+            {
+                await task;
+
+                //  Finish current loop
+
+                UpdateCurrentState();
+
+                await AwaitRemaining(LevelsToComplete, Type, CancellationToken).ConfigureAwait(false);
+            }
+
+            async Task AwaitRemaining(int LevelsToComplete, JsonContainerType Type, CancellationToken CancellationToken)
+            {
+                while (LevelsToComplete-- > 0)
+                {
+                    JsonToken token = GetCloseTokenForType(Pop());
+
+                    if (_currentState == State.Property)
+                    {
+                        await WriteNullAsync(CancellationToken).ConfigureAwait(false);
+                    }
+
+                    if (_formatting == Formatting.Indented)
+                    {
+                        if (_currentState != State.ObjectStart && _currentState != State.ArrayStart)
+                        {
+                            await WriteIndentAsync(CancellationToken).ConfigureAwait(false);
+                        }
+                    }
+
+                    await WriteEndAsync(token, CancellationToken).ConfigureAwait(false);
+
+                    UpdateCurrentState();
+                }
+            }
         }
 
         /// <summary>

Src/Newtonsoft.Json/Linq/JObject.Async.cs

@@ -25,6 +25,7 @@
 
 #if HAVE_ASYNC
 
+using System;
 using System.Globalization;
 using System.Threading;
 using System.Threading.Tasks;
@@ -41,16 +42,54 @@ namespace Newtonsoft.Json.Linq
         /// <param name="cancellationToken">The token to monitor for cancellation requests.</param>
         /// <param name="converters">A collection of <see cref="JsonConverter"/> which will be used when writing the token.</param>
         /// <returns>A <see cref="Task"/> that represents the asynchronous write operation.</returns>
-        public override async Task WriteToAsync(JsonWriter writer, CancellationToken cancellationToken, params JsonConverter[] converters)
+        public override Task WriteToAsync(JsonWriter writer, CancellationToken cancellationToken, params JsonConverter[] converters)
         {
-            await writer.WriteStartObjectAsync(cancellationToken).ConfigureAwait(false);
+            var t = writer.WriteStartObjectAsync(cancellationToken);
+            if (t.Status != TaskStatus.RanToCompletion)
+            {
+                return AwaitStartObject(t, writer, cancellationToken, converters);
+            }
 
             for (int i = 0; i < _properties.Count; i++)
             {
-                await _properties[i].WriteToAsync(writer, cancellationToken, converters).ConfigureAwait(false);
+                t = _properties[i].WriteToAsync(writer, cancellationToken, converters);
+                if (t.Status != TaskStatus.RanToCompletion)
+                {
+                    return AwaitProperties(t, i + 1, writer, cancellationToken, converters);
+                }
             }
 
-            await writer.WriteEndObjectAsync(cancellationToken).ConfigureAwait(false);
+            t = writer.WriteEndObjectAsync(cancellationToken);
+
+            return (t.Status == TaskStatus.RanToCompletion) ? AsyncUtils.CompletedTask : AwaitEndObject(t);
+
+            // Local functions, params renamed (capitalized) so as not to capture and allocate when calling async
+            async Task AwaitStartObject(Task task, JsonWriter Writer, CancellationToken CancellationToken, JsonConverter[] Converters)
+            {
+                await task.ConfigureAwait(false);
+                for (int i = 0; i < _properties.Count; i++)
+                {
+                    await _properties[i].WriteToAsync(Writer, CancellationToken, Converters).ConfigureAwait(false);
+                }
+
+                await Writer.WriteEndObjectAsync(CancellationToken).ConfigureAwait(false);
+            }
+
+            async Task AwaitProperties(Task task, int i, JsonWriter Writer, CancellationToken CancellationToken, JsonConverter[] Converters)
+            {
+                await task.ConfigureAwait(false);
+                for (; i < _properties.Count; i++)
+                {
+                    await _properties[i].WriteToAsync(Writer, CancellationToken, Converters).ConfigureAwait(false);
+                }
+
+                await Writer.WriteEndObjectAsync(CancellationToken).ConfigureAwait(false);
+            }
+
+            async Task AwaitEndObject(Task task)
+            {
+                await task.ConfigureAwait(false);
+            }
         }
 
         /// <summary>