Async-streams: Disposal in async-iterator methods

docs/features/async-streams.md

@@ -36,8 +36,8 @@ An `await using` statement is just like a `using` statement, but it uses `IAsync
 The user can implement those interfaces manually, or can take advantage of the compiler generating a state-machine from a user-defined method (called an "async-iterator" method).
 An async-iterator method is a method that:
 1. is declared `async`,
-2. returns an `IAsyncEnumerable<T>` type,
-3. uses both `await` expressions and `yield` statements.
+2. returns an `IAsyncEnumerable<T>` or `IAsyncEnumerator<T>` type,
+3. uses both `await` syntax (`await` expression, `await foreach` or `await using` statements) and `yield` statements (`yield return`, `yield break`).
 
 For example:
 ```C#
@@ -56,7 +56,9 @@ async IAsyncEnumerable<int> GetValuesFromServer()
 }
 ```
 
-**open issue**: Design async LINQ
+Just like in iterator methods, there are several restrictions on where a yield statement can appear in async-iterator methods:
+- It is a compile-time error for a `yield` statement (of either form) to appear in the `finally` clause of a `try` statement.
+- It is a compile-time error for a `yield return` statement to appear anywhere in a `try` statement that contains any `catch` clauses.
 
 ### Detailed design for `await using` statement
 
@@ -92,31 +94,41 @@ finally
 An async-iterator method is replaced by a kick-off method, which initializes a state machine. It does not start running the state machine (unlike kick-off methods for regular async method).
 The kick-off method method is marked with both `AsyncStateMachineAttribute` and `IteratorStateMachineAttribute`.
 
-The state machine for an async-iterator method primarily implements `IAsyncEnumerable<T>` and `IAsyncEnumerator<T>`.
-It is similar to a state machine produced for an async method. It contains builder and awaiter fields, used to run the state machine in the background (when an `await` is reached in the async-iterator). It also captures parameter values (if any) or `this` (if needed).
+The state machine for an enumerable async-iterator method primarily implements `IAsyncEnumerable<T>` and `IAsyncEnumerator<T>`.
+For an enumerator async-iterator, it only implements `IAsyncEnumerator<T>`.
+It is similar to a state machine produced for an async method.
+It contains builder and awaiter fields, used to run the state machine in the background (when an `await` is reached in the async-iterator).
+It also captures parameter values (if any) or `this` (if needed).
+
 But it contains additional state:
 - a promise of a value-or-end,
 - a current yielded value of type `T`,
-- an `int` capturing the id of the thread that created it.
+- an `int` capturing the id of the thread that created it,
+- a `bool` flag indicating "dispose mode".
 
 The central method of the state machine is `MoveNext()`. It gets run by `MoveNextAsync()`, or as a background continuation initiated from these from an `await` in the method.
 
 The promise of a value-or-end is returned from `MoveNextAsync`. It can be fulfilled with either:
 - `true` (when a value becomes available following background execution of the state machine),
 - `false` (if the end is reached),
 - an exception.
-The promise is implemented as a `ManualResetValueTaskSourceLogic<bool>` (which is a re-usable and allocation-free way of producing and fulfilling `ValueTask<bool>` instances) and its surrounding interfaces on the state machine: `IValueTaskSource<bool>` and `IStrongBox<ManualResetValueTaskSourceLogic<bool>>`.
+The promise is implemented as a `ManualResetValueTaskSourceCore<bool>` (which is a re-usable and allocation-free way of producing and fulfilling `ValueTask<bool>` or `ValueTask` instances)
+and its surrounding interfaces on the state machine: `IValueTaskSource<bool>` and `IValueTaskSource`.
+See more details about those types at https://blogs.msdn.microsoft.com/dotnet/2018/11/07/understanding-the-whys-whats-and-whens-of-valuetask/
 
 Compared to the state machine for a regular async method, the `MoveNext()` for an async-iterator method adds logic:
 - to support handling a `yield return` statement, which saves the current value and fulfills the promise with result `true`,
-- to support handling a `yield break` statement, which fulfills the promise with result `false`,
+- to support handling a `yield break` statement, which sets the dispose mode on and jumps to the closest `finally` or exit,
+- to dispatch execution to `finally` blocks (when disposing),
 - to exit the method, which fulfills the promise with result `false`,
-- to the handling of exceptions, to set the exception into the promise.
+- to catch exceptions, which set the exception into the promise.
 (The handling of an `await` is unchanged)
 
 This is reflected in the implementation, which extends the lowering machinery for async methods to:
-1. handle `yield return` and `yield break` statements (add methods `VisitYieldReturnStatement` and `VisitYieldBreakStatement` to `AsyncMethodToStateMachineRewriter`),
-2. produce additional state and logic for the promise itself (we use `AsyncIteratorRewriter` instead of `AsyncRewriter` to drive the lowering, and produces the other members: `MoveNextAsync`, `Current`, `DisposeAsync`, and some members supporting the resettable `ValueTask`, namely `GetResult`, `SetStatus`, `OnCompleted` and `Value.get`).
+1. handle `yield return` and `yield break` statements (see methods `VisitYieldReturnStatement` and `VisitYieldBreakStatement` to `AsyncIteratorMethodToStateMachineRewriter`),
+2. handle `try` statements (see methods `VisitTryStatement` and `VisitExtractedFinallyBlock` in `AsyncIteratorMethodToStateMachineRewriter`)
+3. produce additional state and logic for the promise itself (see `AsyncIteratorRewriter`, which produces various other members: `MoveNextAsync`, `Current`, `DisposeAsync`,
+and some members supporting the resettable `ValueTask` behavior, namely `GetResult`, `SetStatus`, `OnCompleted`).
 
 ```C#
 ValueTask<bool> MoveNextAsync()
@@ -128,7 +140,7 @@ ValueTask<bool> MoveNextAsync()
     valueOrEndPromise.Reset();
     var inst = this;
     builder.Start(ref inst);
-    return new ValueTask<bool>(this, valueOrEndPromise.Version);
+    return new ValueTask<bool>(this, valueOrEndPromise.Version); // note this leverages the state machine's implementation of IValueTaskSource<bool>
 }
 ```
 
@@ -137,7 +149,7 @@ T Current => current;
 ```
 
 The kick-off method and the initialization of the state machine for an async-iterator method follows those for regular iterator methods.
-In particular, the synthesized `GetAsyncEnumerator()` method is like `GetEnuemrator()` except that it sets the initial state to to StateMachineStates.NotStartedStateMachine (-1):
+In particular, the synthesized `GetAsyncEnumerator()` method is like `GetEnumerator()` except that it sets the initial state to to StateMachineStates.NotStartedStateMachine (-1):
 ```C#
 IAsyncEnumerator<T> GetAsyncEnumerator()
 {
@@ -164,3 +176,137 @@ Similarly, the kick-off method is much like those of regular iterator methods:
     return result;
 }
 ```
+
+#### Disposal
+
+Iterator and async-iterator methods need disposal because their execution steps are controlled by the caller, which could choose to dispose the enumerator before getting all of its elements.
+For example, `foreach (...) { if (...) break; }`.
+In contrast, async methods continue running autonomously until they are done. They are never left suspended in the middle of execution from the caller's perspective, so they don't need to be disposed.
+
+In summary, disposal of an async-iterator works based on four design elements:
+- `yield return` (jumps to finally when resuming in dispose mode)
+- `yield break` (enters dispose mode and jumps to finally)
+- `finally` (after a `finally` we jump to the next one)
+- `DisposeAsync` (enters dispose mode and resumes execution)
+
+The caller of an async-iterator method should only call `DisposeAsync()` when the method completed or was suspended by a `yield return`.
+`DisposeAsync` sets a flag on the state machine ("dispose mode") and (if the method wasn't completed) resumes the execution from the current state.
+The state machine can resume execution from a given state (even those located within a `try`).
+When the execution is resumed in dispose mode, it jumps straight to the relevant `finally`.
+`finally` blocks may involve pauses and resumes, but only for `await` expressions. As a result of the restrictions imposed on `yield return` (described above), dispose mode never runs into a `yield return`.
+Once a `finally` block completes, the execution in dispose mode jumps to the next relevant `finally`, or the end of the method once we reach the top-level.
+
+Reaching a `yield break` also sets the dispose mode flag and jumps to the next relevant `finally` (or end of the method).
+By the time we return control to the caller (completing the promise as `false` by reaching the end of the method) all disposal was completed,
+and the state machine is left in finished state. So `DisposeAsync()` has no work left to do.
+
+Looking at disposal from the perspective of a given `finally` block, the code in that block can get executed:
+- by normal execution (ie. after the code in the `try` block),
+- by raising an exception inside the `try` block (which will execute the necessary `finally` blocks and terminate the method in Finished state),
+- by calling `DisposeAsync()` (which resumes execution in dispose mode and jumps to the relevant finally),
+- following a `yield break` (which enters dispose mode and jumps to the relevant finally),
+- in dispose mode, following a nested `finally`.
+
+A `yield return` is lowered as:
+```C#
+_current = expression;
+_state = <next_state>;
+goto <exprReturnTruelabel>; // which does _valueOrEndPromise.SetResult(true); return;
+
+// resuming from state=<next_state> will dispatch execution to this label
+<next_state_label>: ;
+this.state = cachedState = NotStartedStateMachine;
+if (disposeMode) /* jump to enclosing finally or exit */
+```
+
+A `yield break` is lowered as:
+```C#
+disposeMode = true;
+/* jump to enclosing finally or exit */
+```
+
+```C#
+ValueTask IAsyncDisposable.DisposeAsync()
+{
+    disposeMode = true;
+    if (state == StateMachineStates.FinishedStateMachine ||
+        state == StateMachineStates.NotStartedStateMachine)
+    {
+        return default;
+    }
+    _valueOrEndPromise.Reset();
+    var inst = this;
+    _builder.Start(ref inst);
+    return new ValueTask(this, _valueOrEndPromise.Version); // note this leverages the state machine's implementation of IValueTaskSource
+}
+```
+
+##### Regular versus extracted finally
+
+When the `finally` clause contains no `await` expressions, a `try/finally` is lowered as:
+```C#
+try
+{
+    ...
+    finallyEntryLabel:
+}
+finally
+{
+    ...
+}
+if (disposeMode) /* jump to enclosing finally or exit */
+```
+
+When a `finally` contains `await` expressions, it is extracted before async rewriting (by AsyncExceptionHandlerRewriter). In those cases, we get:
+```C#
+try
+{
+    ...
+    goto finallyEntryLabel;
+}
+catch (Exception e)
+{
+    ... save exception ...
+}
+finallyEntryLabel:
+{
+    ... original code from finally and additional handling for exception ...
+}
+```
+
+In both cases, we will add a `if (disposeMode) /* jump to next finally or exit */` after the block for `finally` logic.
+
+#### State values and transitions
+
+The enumerable starts with state -2.
+Calling GetAsyncEnumerator sets the state to -1, or returns a fresh enumerator (also with state -1).
+
+From there, MoveNext will either:
+- reach the end of the method (-2)
+- reach a `yield break` (-1, dispose mode = true)
+- reach a `yield return` or `await` (N)
+
+From suspended state N, MoveNext will resume execution (-1).
+But if the suspension was a `yield return`, you could also call DisposeAsync, which resumes execution (-1) in dispose mode.
+
+When in dispose mode, MoveNext continues to suspend (N) and resume (-1) until the end of the method is reached (-2).
+
+```
+   GetAsyncEnumerator    suspension (yield return, await)
+-2 -----------------> -1 -------------------------------> N
+ ^                   |  ^                                 |   Dispose mode = false
+ | done and disposed |  |          resuming               |
+ +-------------------+  +---------------------------------+
+ |                   |                                    |
+ |                   |                                    |
+ |             yield |                                    |
+ |             break |           DisposeAsync             |
+ |                   |  +---------------------------------+
+ |                   |  |
+ |                   |  |
+ | done and disposed v  v     suspension (await)
+ +------------------- -1 -------------------------------> N
+                        ^                                 |   Dispose mode = true
+                        |          resuming               |
+                        +---------------------------------+
+```

src/Compilers/CSharp/Portable/BoundTree/BoundNodes.xml

@@ -148,6 +148,12 @@
     <Field Name="ChildBoundNodes" Type="ImmutableArray&lt;BoundNode&gt;"/>
   </Node>
 
+  <!-- This node is used to wrap the BoundBlock for a finally extracted by AsyncExceptionHandlerRewriter.
+    It is processed and removed by AsyncIteratorRewriter. -->
+  <Node Name="BoundExtractedFinallyBlock" Base="BoundStatement">
+    <Field Name="FinallyBlock" Type="BoundBlock"/>
+  </Node>
+
   <Node Name="BoundTypeExpression" Base="BoundExpression">
     <!-- Non-null type is required for this node kind -->
     <Field Name="Type" Type="TypeSymbol" Override="true" Null="disallow"/>
@@ -710,8 +716,6 @@
     <Field Name="StatementOpt" Type="BoundStatement" Null="allow"/>
     <Field Name="Span" Type="TextSpan"/>
   </Node>
-
-
 
   <!--  
   BoundBlock contains 
@@ -1000,10 +1004,17 @@
     <Field Name="TryBlock" Type="BoundBlock"/>
     <Field Name="CatchBlocks" Type="ImmutableArray&lt;BoundCatchBlock&gt;"/>
     <Field Name="FinallyBlockOpt" Type="BoundBlock" Null="allow"/>
-
+
+    <!--
+    When lowering trys, we sometimes extract the finally clause out of the try.
+    For example, `try { } finally { await expr; }` becomes something like `try { } catch { ... } finallyLabel: { await expr; ... }`.
+    We need to save the label for the finally so that async-iterator rewriting can implement proper disposal.
+    -->
+    <Field Name="FinallyLabelOpt" Type="LabelSymbol" Null="allow"/>
+
     <!--
     PreferFaultHandler is a hint to the codegen to emit Finally in the following shape -
-    
+
     try
     {    
     }

src/Compilers/CSharp/Portable/BoundTree/Constructors.cs

@@ -556,8 +556,8 @@ public BoundDefaultExpression(SyntaxNode syntax, TypeSymbol type, bool hasErrors
 
     internal partial class BoundTryStatement
     {
-        public BoundTryStatement(SyntaxNode syntax, BoundBlock tryBlock, ImmutableArray<BoundCatchBlock> catchBlocks, BoundBlock finallyBlockOpt)
-            : this(syntax, tryBlock, catchBlocks, finallyBlockOpt, preferFaultHandler: false, hasErrors: false)
+        public BoundTryStatement(SyntaxNode syntax, BoundBlock tryBlock, ImmutableArray<BoundCatchBlock> catchBlocks, BoundBlock finallyBlockOpt, LabelSymbol finallyLabelOpt = null)
+            : this(syntax, tryBlock, catchBlocks, finallyBlockOpt, finallyLabelOpt, preferFaultHandler: false, hasErrors: false)
         {
         }
     }

src/Compilers/CSharp/Portable/CodeGen/Optimizer.cs

@@ -1586,7 +1586,7 @@ public override BoundNode VisitTryStatement(BoundTryStatement node)
 
             EnsureOnlyEvalStack();
 
-            return node.Update(tryBlock, catchBlocks, finallyBlock, node.PreferFaultHandler);
+            return node.Update(tryBlock, catchBlocks, finallyBlock, finallyLabelOpt: node.FinallyLabelOpt, node.PreferFaultHandler);
         }
 
         public override BoundNode VisitCatchBlock(BoundCatchBlock node)

src/Compilers/CSharp/Portable/Compiler/MethodCompiler.cs

@@ -1285,7 +1285,6 @@ internal static BoundStatement LowerBodyOrInitializer(
                     // The reason why this rewrite happens before the lambda rewrite 
                     // is that we may need access to exception locals and it would be fairly hard to do
                     // if these locals are captured into closures (possibly nested ones).
-                    Debug.Assert(!method.IsIterator);
                     loweredBody = AsyncExceptionHandlerRewriter.Rewrite(
                         method,
                         method.ContainingType,
@@ -1335,6 +1334,7 @@ internal static BoundStatement LowerBodyOrInitializer(
                     return bodyWithoutIterators;
                 }
 
+                // Rewrite async and async-iterator methods
                 AsyncStateMachine asyncStateMachine;
                 BoundStatement bodyWithoutAsync = AsyncRewriter.Rewrite(bodyWithoutIterators, method, methodOrdinal, lazyVariableSlotAllocator, compilationState, diagnostics, out asyncStateMachine);
 

src/Compilers/CSharp/Portable/FlowAnalysis/AbstractFlowPass.cs

@@ -1651,6 +1651,11 @@ protected virtual void VisitFinallyBlock(BoundStatement finallyBlock, ref TLocal
             VisitStatement(finallyBlock); // this should generate no pending branches
         }
 
+        public override BoundNode VisitExtractedFinallyBlock(BoundExtractedFinallyBlock node)
+        {
+            return VisitBlock(node.FinallyBlock);
+        }
+
         public sealed override BoundNode VisitReturnStatement(BoundReturnStatement node)
         {
             var result = VisitReturnStatementNoAdjust(node);
@@ -2286,7 +2291,7 @@ public override BoundNode VisitForEachStatement(BoundForEachStatement node)
             LoopTail(node);
             ResolveBreaks(breakState, node.BreakLabel);
 
-            if (((CommonForEachStatementSyntax)node.Syntax).AwaitKeyword != default)
+            if (AwaitUsingAndForeachAddsPendingBranch && ((CommonForEachStatementSyntax)node.Syntax).AwaitKeyword != default)
             {
                 PendingBranches.Add(new PendingBranch(node, this.State));
             }
@@ -2623,14 +2628,14 @@ public override BoundNode VisitUsingStatement(BoundUsingStatement node)
 
             VisitStatement(node.Body);
 
-            if (AwaitUsingAddsPendingBranch && node.AwaitOpt != null)
+            if (AwaitUsingAndForeachAddsPendingBranch && node.AwaitOpt != null)
             {
                 PendingBranches.Add(new PendingBranch(node, this.State));
             }
             return null;
         }
 
-        public abstract bool AwaitUsingAddsPendingBranch { get; }
+        public abstract bool AwaitUsingAndForeachAddsPendingBranch { get; }
 
         public override BoundNode VisitFixedStatement(BoundFixedStatement node)
         {

src/Compilers/CSharp/Portable/FlowAnalysis/ControlFlowPass.cs

@@ -290,8 +290,8 @@ protected override void VisitFinallyBlock(BoundStatement finallyBlock, ref Local
             RestorePending(oldPending1);
         }
 
-        // For purpose of control flow analysis, awaits do not create pending branches, so asynchronous usings don't either
-        public sealed override bool AwaitUsingAddsPendingBranch => false;
+        // For purpose of control flow analysis, awaits do not create pending branches, so asynchronous usings and foreachs don't either
+        public sealed override bool AwaitUsingAndForeachAddsPendingBranch => false;
 
         protected override void VisitLabel(BoundLabeledStatement node)
         {

src/Compilers/CSharp/Portable/FlowAnalysis/DefiniteAssignment.cs

@@ -274,8 +274,8 @@ private static bool HasAwait(PendingBranch pending)
             }
         }
 
-        // For purpose of data flow analysis, awaits create pending branches, so async usings do too
-        public sealed override bool AwaitUsingAddsPendingBranch => true;
+        // For purpose of definite assignment analysis, awaits create pending branches, so async usings and foreachs do too
+        public sealed override bool AwaitUsingAndForeachAddsPendingBranch => true;
 
         protected virtual void ReportUnassignedOutParameter(ParameterSymbol parameter, SyntaxNode node, Location location)
         {

src/Compilers/CSharp/Portable/FlowAnalysis/NullableWalker.cs

@@ -164,8 +164,8 @@ private NullableWalker(
             }
         }
 
-        // For purpose of nullability analysis, awaits create pending branches, so async usings do too
-        public sealed override bool AwaitUsingAddsPendingBranch => true;
+        // For purpose of nullability analysis, awaits create pending branches, so async usings and foreachs do too
+        public sealed override bool AwaitUsingAndForeachAddsPendingBranch => true;
 
         protected override bool ConvertInsufficientExecutionStackExceptionToCancelledByStackGuardException()
         {