Update DLR; adapt to changed async API; add tests

Author: BCSharp

src/core/IronPython/Compiler/Ast/AstMethods.cs

@@ -83,7 +83,7 @@ internal static class AstMethods {
         public static readonly MethodInfo MakeCoroutine = GetMethod((Func<PythonFunction, MutableTuple, object, PythonCoroutine>)PythonOps.MakeCoroutine);
 #if FEATURE_NET_ASYNC
         public static readonly MethodInfo AsTaskForAwait = GetMethod((Func<object, System.Threading.Tasks.Task<object>>)PythonOps.AsTaskForAwait);
-        public static readonly MethodInfo MakeAsyncCoroutine = GetMethod((Func<PythonFunction, Func<System.Threading.Tasks.Task<object>>, System.Threading.CancellationTokenSource, System.Runtime.CompilerServices.StrongBox<System.Exception>, PythonCoroutine>)PythonOps.MakeAsyncCoroutine);
+        public static readonly MethodInfo MakeAsyncCoroutine = GetMethod((Func<PythonFunction, Func<System.Threading.Tasks.Task<object>>, Microsoft.Scripting.Runtime.AsyncInterruptSource, PythonCoroutine>)PythonOps.MakeAsyncCoroutine);
         public static readonly MethodInfo MakeAsyncGenerator = GetMethod((Func<PythonFunction, System.Collections.Generic.IAsyncEnumerable<object>, System.Runtime.CompilerServices.StrongBox<object>, System.Runtime.CompilerServices.StrongBox<System.Exception>, System.Threading.CancellationTokenSource, PythonAsyncGenerator>)PythonOps.MakeAsyncGenerator);
 #endif
 

src/core/IronPython/Compiler/Ast/FunctionDefinition.cs

@@ -727,33 +727,27 @@ private LightLambdaExpression CreateFunctionLambda() {
 
 #if FEATURE_NET_ASYNC
             // Under .NET-async, an `async def` body returns a PythonCoroutine wrapping a Task<object?>.
-            // We pre-allocate a CancellationTokenSource and a StrongBox<Exception?> here
-            // so the same instances are shared with both AsyncExpression, which threads them into AsyncHelpers.DriveAsync
-            // and PythonCoroutine, which uses them to implement coro.throw(exc) on a running coroutine:
-            // write the exception to the box, cancel the CTS, and DriveAsync surfaces that exception in place of OperationCanceledException.
             if (IsAsync) {
-                var cts = MSAst.Expression.Variable(typeof(CancellationTokenSource), "$cts");
-                var excBox = MSAst.Expression.Variable(typeof(StrongBox<Exception>), "$cancelExc");
-                var ctToken = MSAst.Expression.Property(cts, nameof(CancellationTokenSource.Token));
                 if (IsGenerator) {
                     // Async generator: the body has both `await` and `yield`. Lower it to an
                     // IAsyncEnumerable<object?> via AsyncEnumerableExpression, sharing the generator label so
                     // the body's yields and the rewritten awaits land in one generator, then wrap it in a
                     // PythonAsyncGenerator. The send/throw slots are per-generator StrongBoxes captured by the
                     // generator (the body's yields read them) AND handed to the wrapper, which writes them
                     // before each resume — see AsyncSendSlot / AsyncThrowSlot.
+                    var cts = MSAst.Expression.Variable(typeof(CancellationTokenSource), "$cts");
+                    var ctToken = MSAst.Expression.Property(cts, nameof(CancellationTokenSource.Token));
                     var sendSlot = AsyncSendSlot;
                     var throwSlot = AsyncThrowSlot;
                     body = MSAst.Expression.Block(
-                        [cts, excBox, sendSlot, throwSlot],
+                        [cts, sendSlot, throwSlot],
                         MSAst.Expression.Assign(cts, MSAst.Expression.New(typeof(CancellationTokenSource))),
-                        MSAst.Expression.Assign(excBox, MSAst.Expression.New(typeof(StrongBox<Exception>))),
                         MSAst.Expression.Assign(sendSlot, MSAst.Expression.New(typeof(StrongBox<object>))),
                         MSAst.Expression.Assign(throwSlot, MSAst.Expression.New(typeof(StrongBox<Exception>))),
                         Ast.Call(
                             AstMethods.MakeAsyncGenerator,
                             _functionParam,
-                            AstUtils.AsyncEnumerable(Name, body, GeneratorLabel, ctToken, excBox),
+                            AstUtils.AsyncEnumerable(Name, body, GeneratorLabel, interrupts: null, ctToken),
                             sendSlot,
                             throwSlot,
                             cts));
@@ -763,17 +757,18 @@ private LightLambdaExpression CreateFunctionLambda() {
                     // so the body doesn't execute until the coroutine is first driven (send/AsTask).
                     // This makes calling an async def side-effect-free (PEP 492) and lets the body's first await capture the driver's SynchronizationContext
                     // rather than whatever context happened to be current at construction.
+                    // The AsyncInterruptSource is shared by AsyncExpression, which threads it into AsyncHelpers.DriveAsync,
+                    // and PythonCoroutine, which uses it to implement coro.throw(exc).
+                    var interrupts = MSAst.Expression.Variable(typeof(Microsoft.Scripting.Runtime.AsyncInterruptSource), "$interrupts");
                     body = MSAst.Expression.Block(
-                        [cts, excBox],
-                        MSAst.Expression.Assign(cts, MSAst.Expression.New(typeof(CancellationTokenSource))),
-                        MSAst.Expression.Assign(excBox, MSAst.Expression.New(typeof(StrongBox<Exception>))),
+                        [interrupts],
+                        MSAst.Expression.Assign(interrupts, MSAst.Expression.New(typeof(Microsoft.Scripting.Runtime.AsyncInterruptSource))),
                         Ast.Call(
                             AstMethods.MakeAsyncCoroutine,
                             _functionParam,
                             MSAst.Expression.Lambda<Func<Task<object>>>(
-                                AstUtils.Async(Name, body, ctToken, excBox)),
-                            cts,
-                            excBox));
+                                AstUtils.Async(Name, body, interrupts)),
+                            interrupts));
                 }
             }
 #endif

src/core/IronPython/Runtime/Coroutine.cs

@@ -6,7 +6,6 @@
 
 using System;
 using System.Runtime.CompilerServices;
-using System.Threading;
 using System.Threading.Tasks;
 
 using Microsoft.Scripting.Runtime;
@@ -24,12 +23,8 @@ public sealed class PythonCoroutine : ICodeFormattable, IWeakReferenceable {
 #if FEATURE_NET_ASYNC
 
         // Under .NET async, `async def` is compiled directly to a Task<object?> via the DLR's AsyncExpression + AsyncHelpers.
-        // PythonCoroutine is a wrapper over the Task with two cancellation channels pre-bound at the codegen level:
-        //
-        //   _cts                   — CancellationTokenSource whose Token was passed into AsyncExpression. throw(exc) cancels this.
-        //   _cancellationException — StrongBox<Exception?> shared with DriveAsync's `cancellationException` parameter;
-        //                            when non-null at cancellation time, that exception is surfaced to the body
-        //                            instead of OperationCanceledException. 
+        // PythonCoroutine is a wrapper over the Task. _interrupts is the AsyncInterruptSource passed into AsyncExpression;
+        // throw(exc) interrupts the body with it.
         //
         // Lazy start: the body is not executed at construction. _factory is the thunk that creates (and thereby starts driving) the Task on
         // first access; _task caches the materialized Task. Calling an async def is therefore side-effect free until the coroutine is first
@@ -38,18 +33,14 @@ public sealed class PythonCoroutine : ICodeFormattable, IWeakReferenceable {
         private Task<object?>? _task;
         private readonly string _name;
         private readonly FunctionCode? _code;
-        private readonly CancellationTokenSource _cts;
-        private readonly StrongBox<Exception?> _cancellationException;
+        private readonly AsyncInterruptSource _interrupts;
         private WeakRefTracker? _tracker;
 
-        internal PythonCoroutine(Func<Task<object?>> factory, string name, FunctionCode? code,
-                                  CancellationTokenSource cts,
-                                  StrongBox<Exception?> cancellationException) {
+        internal PythonCoroutine(Func<Task<object?>> factory, string name, FunctionCode? code, AsyncInterruptSource interrupts) {
             _factory = factory;
             _name = name;
             _code = code;
-            _cts = cts;
-            _cancellationException = cancellationException;
+            _interrupts = interrupts;
         }
 
         private Task<object?> EnsureTask() => _task ??= _factory();
@@ -106,22 +97,18 @@ public object @throw(object? type, object? value, object? traceback) {
             Exception ex = PythonOps.MakeExceptionForGenerator(
                 DefaultContext.Default, type, value, traceback, cause: null);
 
-            var task = EnsureTask();
-            if (task.IsCompleted) {
-                // Case 1: coroutine has finished. throw() raises the exception immediately —
+            // The body sees `ex` raised at the await it is suspended at, or at its next await if it is running.
+            // If the body hasn't started, it doesn't run at all: its task fails with `ex`, as in CPython.
+            if (!_interrupts.Interrupt(ex)) {
+                // The coroutine has finished. throw() raises the exception immediately —
                 // the body has nothing left to observe it. Matches CPython.
                 return LightExceptions.Throw(ex);
             }
 
-            // Case 2: body is still running (or suspended at an await). Stash the exception in the
-            // shared StrongBox so DriveAsync surfaces it (instead of OCE) the moment it observes the
-            // cancellation, then cancel. The body sees `ex` rethrown at its next resume point.
-            _cancellationException.Value = ex;
-            _cts.Cancel();
-
             // Wait for the body to settle. After this the task is in a terminal state — dispatch via
             // the same logic as send() so the body's reaction (catch-and-return, propagate, etc.) is
             // reflected back to the caller of throw().
+            var task = EnsureTask();
             ((IAsyncResult)task).AsyncWaitHandle.WaitOne();
             return SettleCompletedTask(task);
         }

src/core/IronPython/Runtime/Operations/PythonOps.cs

@@ -3345,12 +3345,11 @@ public static object MakeCoroutineWrapper(PythonGenerator generator) {
         /// <remarks>
         ///   The thunk is invoked lazily on first drive (send/AsTask), so calling an <c>async def</c> is side-effect-free until the coroutine
         ///   is awaited (PEP 492).
-        ///   The supplied <see cref="CancellationTokenSource"/> and exception-override box are pre-bound
-        ///   to <see cref="Microsoft.Scripting.Runtime.AsyncHelpers.DriveAsync"/>'s cancellation channel — the coroutine uses them to implement
-        ///   <c>coro.throw(exc)</c> on a running body (set the box, cancel the CTS).
+        ///   The supplied <see cref="AsyncInterruptSource"/> is the one the body's driver,
+        ///   <see cref="Microsoft.Scripting.Runtime.AsyncHelpers.DriveAsync"/>, uses — the coroutine uses it to implement <c>coro.throw(exc)</c>.
         /// </remarks>
-        public static PythonCoroutine MakeAsyncCoroutine(PythonFunction function, Func<Task<object?>> taskFactory, CancellationTokenSource cts, StrongBox<System.Exception?> cancellationException) {
-            return new PythonCoroutine(taskFactory, function.__name__, function.__code__, cts, cancellationException);
+        public static PythonCoroutine MakeAsyncCoroutine(PythonFunction function, Func<Task<object?>> taskFactory, AsyncInterruptSource interrupts) {
+            return new PythonCoroutine(taskFactory, function.__name__, function.__code__, interrupts);
         }
 
 

src/dlr

@@ -519,6 +519,98 @@ def test_cancellation_token_cancel(self):
         self.assertTrue(token.IsCancellationRequested)
 
 
+@skipUnlessIronPython()
+@unittest.skipUnless(is_netcoreapp, "tests the .NET async implementation (FEATURE_NET_ASYNC)")
+class CoroutineThrowTest(unittest.TestCase):
+    """coro.throw(exc) raises exc once, at the await the coroutine is suspended at, and can be repeated."""
+
+    @staticmethod
+    def pending():
+        from System.Threading.Tasks import TaskCompletionSource
+        return TaskCompletionSource[object]().Task  # never completes
+
+    @staticmethod
+    def throw_at_pending_await(coro, exc):
+        import clr
+        coro.AsTask()   # runs the body up to its first await of a task that hasn't completed
+        return coro.throw(exc)
+
+    def test_caught(self):
+        async def coro():
+            try:
+                await self.pending()
+            except KeyError as e:
+                return 'caught %s' % e.args[0]
+        with self.assertRaises(StopIteration) as cm:
+            self.throw_at_pending_await(coro(), KeyError('k'))
+        self.assertEqual(cm.exception.value, 'caught k')
+
+    def test_not_caught(self):
+        async def coro():
+            await self.pending()
+        with self.assertRaises(KeyError):
+            self.throw_at_pending_await(coro(), KeyError('k'))
+
+    def test_awaits_after_caught(self):
+        # The exception interrupts the body once: its later awaits complete normally
+        from System.Threading.Tasks import Task
+        async def coro():
+            try:
+                await self.pending()
+            except KeyError:
+                pass
+            await Task.Delay(10)
+            return await Task.FromResult(42)
+        with self.assertRaises(StopIteration) as cm:
+            self.throw_at_pending_await(coro(), KeyError('k'))
+        self.assertEqual(cm.exception.value, 42)
+
+    def test_repeated(self):
+        # The body handles a throw() and continues; a later throw() interrupts it again
+        import threading
+        at_second = threading.Event()
+        async def coro():
+            caught = []
+            try:
+                await self.pending()
+            except KeyError:
+                caught.append('KeyError')
+            at_second.set()
+            try:
+                await self.pending()
+            except ValueError:
+                caught.append('ValueError')
+            return caught
+
+        c = coro()
+        c.AsTask()
+        first = []
+        def throw_first():
+            # throw() returns only when the coroutine finishes
+            try:
+                c.throw(KeyError('k'))
+            except StopIteration as e:
+                first.append(e.value)
+        t = threading.Thread(target=throw_first)
+        t.start()
+        self.assertTrue(at_second.wait(10))
+        with self.assertRaises(StopIteration) as cm:
+            c.throw(ValueError('v'))
+        t.join(10)
+        self.assertEqual(cm.exception.value, ['KeyError', 'ValueError'])
+        self.assertEqual(first, [['KeyError', 'ValueError']])
+
+    def test_not_started(self):
+        # As in CPython, the body doesn't run
+        ran = []
+        async def coro():
+            ran.append(True)
+            await self.pending()
+        with self.assertRaises(KeyError):
+            coro().throw(KeyError('k'))
+        self.assertEqual(ran, [])
+
+
 @unittest.skipUnless(is_netcoreapp, "ValueTask is not available on .NET Framework")
 class ValueTaskInteropTest(unittest.TestCase):
     def test_await_valuetask(self):