Skip to content

Commit 62676f3

Browse files
committed
Implement async support in coroutine handling
1 parent 76e1224 commit 62676f3

10 files changed

Lines changed: 379 additions & 1 deletion

File tree

eng/net10.0.props

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
<Features>$(Features);FEATURE_METADATA_READER</Features>
2424
<Features>$(Features);FEATURE_MMAP</Features>
2525
<Features>$(Features);FEATURE_NATIVE</Features>
26+
<Features>$(Features);FEATURE_NET_ASYNC</Features>
2627
<Features>$(Features);FEATURE_OSPLATFORMATTRIBUTE</Features>
2728
<Features>$(Features);FEATURE_PIPES</Features>
2829
<Features>$(Features);FEATURE_PROCESS</Features>

eng/net8.0.props

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
<Features>$(Features);FEATURE_METADATA_READER</Features>
2424
<Features>$(Features);FEATURE_MMAP</Features>
2525
<Features>$(Features);FEATURE_NATIVE</Features>
26+
<Features>$(Features);FEATURE_NET_ASYNC</Features>
2627
<Features>$(Features);FEATURE_OSPLATFORMATTRIBUTE</Features>
2728
<Features>$(Features);FEATURE_PIPES</Features>
2829
<Features>$(Features);FEATURE_PROCESS</Features>

eng/net9.0.props

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
<Features>$(Features);FEATURE_METADATA_READER</Features>
2424
<Features>$(Features);FEATURE_MMAP</Features>
2525
<Features>$(Features);FEATURE_NATIVE</Features>
26+
<Features>$(Features);FEATURE_NET_ASYNC</Features>
2627
<Features>$(Features);FEATURE_OSPLATFORMATTRIBUTE</Features>
2728
<Features>$(Features);FEATURE_PIPES</Features>
2829
<Features>$(Features);FEATURE_PROCESS</Features>

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,10 @@ internal static class AstMethods {
8080
public static readonly MethodInfo FormatString = GetMethod((Func<CodeContext, string, object, string>)PythonOps.FormatString);
8181
public static readonly MethodInfo GeneratorCheckThrowableAndReturnSendValue = GetMethod((Func<object, object>)PythonOps.GeneratorCheckThrowableAndReturnSendValue);
8282
public static readonly MethodInfo MakeCoroutine = GetMethod((Func<PythonFunction, MutableTuple, object, PythonCoroutine>)PythonOps.MakeCoroutine);
83+
#if FEATURE_NET_ASYNC
84+
public static readonly MethodInfo AsTaskForAwait = GetMethod((Func<object, System.Threading.Tasks.Task<object>>)PythonOps.AsTaskForAwait);
85+
public static readonly MethodInfo MakeAsyncCoroutine = GetMethod((Func<PythonFunction, System.Threading.Tasks.Task<object>, System.Threading.CancellationTokenSource, System.Runtime.CompilerServices.StrongBox<System.Exception>, PythonCoroutine>)PythonOps.MakeAsyncCoroutine);
86+
#endif
8387

8488
// builtins
8589
public static readonly MethodInfo Format = GetMethod((Func<CodeContext, object, string, string>)PythonOps.Format);

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

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,15 +6,43 @@
66

77
using MSAst = System.Linq.Expressions;
88

9+
using IronPython.Runtime.Operations;
10+
911
using AstUtils = Microsoft.Scripting.Ast.Utils;
1012

1113
namespace IronPython.Compiler.Ast {
1214
using Ast = MSAst.Expression;
1315

1416
/// <summary>
15-
/// Represents an await expression. Implemented as yield from expr.__await__().
17+
/// Represents <c>await expr</c>. Under <c>FEATURE_NET_ASYNC</c> this compiles
18+
/// directly to a DLR runtime-async suspension point. Otherwise it is
19+
/// desugared into <c>yield from expr.__await__()</c> against the enclosing
20+
/// generator-shaped coroutine state machine.
1621
/// </summary>
1722
public class AwaitExpression : Expression {
23+
#if FEATURE_NET_ASYNC
24+
public AwaitExpression(Expression expression) {
25+
Expression = expression;
26+
}
27+
28+
public Expression Expression { get; }
29+
30+
public override MSAst.Expression Reduce() {
31+
// await x -> AsyncHelpers-driven suspension on the Task produced by
32+
// PythonOps.AsTaskForAwait(x).
33+
return AstUtils.Await(
34+
Ast.Call(
35+
AstMethods.AsTaskForAwait,
36+
AstUtils.Convert(Expression, typeof(object))));
37+
}
38+
39+
public override void Walk(PythonWalker walker) {
40+
if (walker.Walk(this)) {
41+
Expression?.Walk(walker);
42+
}
43+
walker.PostWalk(this);
44+
}
45+
#else
1846
private readonly Statement _statement;
1947
private readonly NameExpression _result;
2048

@@ -60,6 +88,7 @@ public override void Walk(PythonWalker walker) {
6088
}
6189
walker.PostWalk(this);
6290
}
91+
#endif
6392

6493
public override string NodeName => "await expression";
6594
}

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,15 @@ internal override int KwOnlyArgCount {
117117

118118
public Expression ReturnAnnotation { get; internal set; }
119119

120+
#if FEATURE_NET_ASYNC
121+
// Under runtime-async, async functions are compiled directly to a
122+
// Task<object?> via the DLR's AsyncExpression rather than reused
123+
// through the generator state machine, so IsAsync no longer implies
124+
// generator-shaped emission.
125+
internal override bool IsGeneratorMethod => IsGenerator;
126+
#else
120127
internal override bool IsGeneratorMethod => IsGenerator || IsAsync;
128+
#endif
121129

122130
/// <summary>
123131
/// The function is a generator
@@ -182,9 +190,15 @@ internal override FunctionAttributes Flags {
182190
fa |= FunctionAttributes.ContainsTryFinally;
183191
}
184192

193+
#if FEATURE_NET_ASYNC
194+
if (IsGenerator) {
195+
fa |= FunctionAttributes.Generator;
196+
}
197+
#else
185198
if (IsGenerator || IsAsync) {
186199
fa |= FunctionAttributes.Generator;
187200
}
201+
#endif
188202

189203
if (IsAsync) {
190204
fa |= FunctionAttributes.Coroutine;
@@ -357,9 +371,15 @@ internal MSAst.Expression MakeFunctionExpression() {
357371
annotations
358372
)
359373
),
374+
#if FEATURE_NET_ASYNC
375+
IsGenerator ?
376+
(MSAst.Expression)new PythonGeneratorExpression(code, GlobalParent.PyContext.Options.CompilationThreshold, IsAsync) :
377+
(MSAst.Expression)code
378+
#else
360379
(IsGenerator || IsAsync) ?
361380
(MSAst.Expression)new PythonGeneratorExpression(code, GlobalParent.PyContext.Options.CompilationThreshold, IsAsync) :
362381
(MSAst.Expression)code
382+
#endif
363383
);
364384
} else {
365385
ret = Ast.Call(
@@ -659,10 +679,17 @@ private LightLambdaExpression CreateFunctionLambda() {
659679
// For generators/coroutines, we need to do a check before the first statement for Generator.Throw() / Generator.Close().
660680
// The exception traceback needs to come from the generator's method body, and so we must do the check and throw
661681
// from inside the generator.
682+
#if FEATURE_NET_ASYNC
683+
if (IsGenerator) {
684+
MSAst.Expression s1 = YieldExpression.CreateCheckThrowExpression(SourceSpan.None);
685+
statements.Add(s1);
686+
}
687+
#else
662688
if (IsGenerator || IsAsync) {
663689
MSAst.Expression s1 = YieldExpression.CreateCheckThrowExpression(SourceSpan.None);
664690
statements.Add(s1);
665691
}
692+
#endif
666693

667694
if (Body.CanThrow && !(Body is SuiteStatement) && Body.StartIndex != -1) {
668695
statements.Add(UpdateLineNumber(GlobalParent.IndexToLocation(Body.StartIndex).Line));
@@ -681,6 +708,31 @@ private LightLambdaExpression CreateFunctionLambda() {
681708
body = Ast.Block(body, AstUtils.Empty());
682709
body = AddReturnTarget(body);
683710

711+
#if FEATURE_NET_ASYNC
712+
// Under runtime-async, an `async def` body returns a PythonCoroutine wrapping a
713+
// Task<object?>. We pre-allocate a CancellationTokenSource and a StrongBox<Exception?>
714+
// here so the same instances are shared with both AsyncExpression (which threads them into
715+
// AsyncHelpers.DriveAsync) and PythonCoroutine (which uses them to implement
716+
// coro.throw(exc) on a running coroutine: write the exception to the box, cancel the CTS,
717+
// and DriveAsync surfaces that exception in place of OperationCanceledException).
718+
if (IsAsync) {
719+
var cts = MSAst.Expression.Variable(typeof(System.Threading.CancellationTokenSource), "$cts");
720+
var excBox = MSAst.Expression.Variable(typeof(System.Runtime.CompilerServices.StrongBox<Exception>), "$cancelExc");
721+
body = MSAst.Expression.Block(
722+
new[] { cts, excBox },
723+
MSAst.Expression.Assign(cts, MSAst.Expression.New(typeof(System.Threading.CancellationTokenSource))),
724+
MSAst.Expression.Assign(excBox, MSAst.Expression.New(typeof(System.Runtime.CompilerServices.StrongBox<Exception>))),
725+
Ast.Call(
726+
AstMethods.MakeAsyncCoroutine,
727+
_functionParam,
728+
AstUtils.Async(Name, body,
729+
MSAst.Expression.Property(cts, nameof(System.Threading.CancellationTokenSource.Token)),
730+
excBox),
731+
cts,
732+
excBox));
733+
}
734+
#endif
735+
684736
MSAst.Expression bodyStmt = body;
685737
if (localContext != null) {
686738
var createLocal = CreateLocalContext(_parentContext);

src/core/IronPython/Compiler/Parser.cs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1999,10 +1999,17 @@ private Expression ParseAtomExpr() {
19991999
if (current is null || !current.IsAsync) {
20002000
ReportSyntaxError("'await' outside async function");
20012001
}
2002+
#if !FEATURE_NET_ASYNC
2003+
// Under the async-generator path, `await` desugars to `yield from`, so the
2004+
// enclosing function must be marked a generator. Under FEATURE_NET_ASYNC
2005+
// the body is lowered through AsyncExpression and is *not*
2006+
// a Python generator — leaving IsGenerator false here is what keeps
2007+
// ReturnStatement out of the yield branch.
20022008
if (current is not null) {
20032009
current.IsGenerator = true;
20042010
current.GeneratorStop = GeneratorStop;
20052011
}
2012+
#endif
20062013
}
20072014

20082015
Expression ret = ParseAtom();

src/core/IronPython/Runtime/Coroutine.cs

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,9 @@
44

55
#nullable enable
66

7+
using System;
78
using System.Runtime.CompilerServices;
9+
using System.Threading;
810
using System.Threading.Tasks;
911

1012
using Microsoft.Scripting.Runtime;
@@ -14,6 +16,142 @@
1416
using IronPython.Runtime.Types;
1517

1618
namespace IronPython.Runtime {
19+
#if FEATURE_NET_ASYNC
20+
[PythonType("coroutine")]
21+
[DontMapIDisposableToContextManager, DontMapIEnumerableToContains]
22+
public sealed class PythonCoroutine : ICodeFormattable, IWeakReferenceable {
23+
// Under .NET async, `async def` is compiled directly to a Task<object?>
24+
// via the DLR's AsyncExpression + AsyncHelpers. PythonCoroutine is a
25+
// facade over the Task with two cancellation channels pre-bound at the
26+
// codegen level (see FunctionDefinition):
27+
// _cts — CancellationTokenSource whose Token was passed
28+
// into AsyncExpression. throw(exc) cancels this.
29+
// _cancellationException — StrongBox<Exception?> shared with DriveAsync's
30+
// `cancellationException` parameter; when non-null
31+
// at cancellation time, that exception is surfaced
32+
// to the body instead of OperationCanceledException.
33+
// throw(non-OCE) writes here before cancelling.
34+
private readonly Task<object?> _task;
35+
private readonly string _name;
36+
private readonly FunctionCode? _code;
37+
private readonly CancellationTokenSource _cts;
38+
private readonly StrongBox<Exception?> _cancellationException;
39+
private WeakRefTracker? _tracker;
40+
41+
internal PythonCoroutine(Task<object?> task, string name, FunctionCode? code,
42+
CancellationTokenSource cts,
43+
StrongBox<Exception?> cancellationException) {
44+
_task = task;
45+
_name = name;
46+
_code = code;
47+
_cts = cts;
48+
_cancellationException = cancellationException;
49+
}
50+
51+
[LightThrowing]
52+
public object send(object? value) {
53+
// Fast path: task already completed. Dispatch on terminal state without
54+
// any try/catch — Result on a Faulted/Canceled task throws, so we must
55+
// check IsCanceled / IsFaulted before reading it.
56+
//
57+
// Slow path: not completed. Block via AsyncWaitHandle.WaitOne() so we do
58+
// not pay for exception unwinding (GetAwaiter().GetResult() would observe
59+
// and rethrow cancel/fault). After waiting, the same terminal-state
60+
// dispatch below applies.
61+
if (!_task.IsCompleted) {
62+
((IAsyncResult)_task).AsyncWaitHandle.WaitOne();
63+
}
64+
if (_task.IsCanceled) {
65+
// Surfaces as Python CancelledError via the OCE -> CancelledError mapping in PythonExceptions.
66+
// We need an OperationCanceledException instance because the Canceled task carries no Exception object.
67+
return LightExceptions.Throw(new TaskCanceledException(_task));
68+
}
69+
if (_task.IsFaulted) {
70+
return LightExceptions.Throw(_task.Exception?.InnerException ?? _task.Exception);
71+
}
72+
return LightExceptions.Throw(new PythonExceptions._StopIteration().InitAndGetClrException(_task.Result!));
73+
}
74+
75+
[LightThrowing]
76+
public object @throw(object? type) => @throw(type, null, null);
77+
78+
[LightThrowing]
79+
public object @throw(object? type, object? value) => @throw(type, value, null);
80+
81+
[LightThrowing]
82+
public object @throw(object? type, object? value, object? traceback) {
83+
// Validate the shape of (type, value, traceback) up front. Mirrors PythonGenerator.@throw
84+
// so we don't mutate any state on bad input.
85+
if (type is Exception || type is PythonExceptions.BaseException) {
86+
if (value is not null)
87+
return LightExceptions.Throw(PythonOps.TypeError("instance exception may not have a separate value"));
88+
} else if (type is PythonType pt && typeof(PythonExceptions.BaseException).IsAssignableFrom(pt.UnderlyingSystemType)) {
89+
// ok — class form, MakeExceptionForGenerator will construct.
90+
} else {
91+
return LightExceptions.Throw(PythonOps.TypeError(
92+
"exceptions must be classes or instances deriving from BaseException, not {0}",
93+
PythonOps.GetPythonTypeName(type)));
94+
}
95+
96+
// Construct the actual exception object. DefaultContext.Default suffices here — throw()'s
97+
// CodeContext role is limited to resolving class-form exception construction, and the
98+
// coroutine's own context will re-bind the chain when the exception unwinds inside the body.
99+
Exception ex = PythonOps.MakeExceptionForGenerator(
100+
DefaultContext.Default, type, value, traceback, cause: null);
101+
102+
if (_task.IsCompleted) {
103+
// Regime 1: coroutine has finished. throw() raises the exception immediately —
104+
// the body has nothing left to observe it. Matches CPython.
105+
return LightExceptions.Throw(ex);
106+
}
107+
108+
// Regime 2: body is still running (or suspended at an await). Stash the exception in the
109+
// shared StrongBox so DriveAsync surfaces it (instead of OCE) the moment it observes the
110+
// cancellation, then cancel. The body sees `ex` rethrown at its next resume point via the
111+
// rewriter's per-await ExceptionDispatchInfo block.
112+
_cancellationException.Value = ex;
113+
_cts.Cancel();
114+
115+
// Wait for the body to settle. After this the task is in a terminal state — dispatch via
116+
// the same logic as send() so the body's reaction (catch-and-return, propagate, etc.) is
117+
// reflected back to the caller of throw().
118+
((IAsyncResult)_task).AsyncWaitHandle.WaitOne();
119+
return SettleCompletedTask();
120+
}
121+
122+
private object SettleCompletedTask() {
123+
if (_task.IsCanceled) {
124+
return LightExceptions.Throw(new TaskCanceledException(_task));
125+
}
126+
if (_task.IsFaulted) {
127+
return LightExceptions.Throw(_task.Exception?.InnerException ?? _task.Exception);
128+
}
129+
return LightExceptions.Throw(new PythonExceptions._StopIteration().InitAndGetClrException(_task.Result!));
130+
}
131+
132+
[LightThrowing]
133+
public object? close() => null;
134+
135+
public object __await__() => new CoroutineWrapper(this);
136+
137+
public FunctionCode? cr_code => _code;
138+
139+
public int cr_running => _task.IsCompleted ? 0 : 1;
140+
141+
public TraceBackFrame? cr_frame => null;
142+
143+
public string __name__ => _name;
144+
145+
public string __qualname__ => _name;
146+
147+
/// <summary>Returns the underlying <see cref="Task{Object}"/>.</summary>
148+
public Task<object?> AsTask() => _task;
149+
150+
/// <summary>Enables <c>await coroutine</c> from C# code.</summary>
151+
public TaskAwaiter<object?> GetAwaiter() => _task.GetAwaiter();
152+
153+
internal Task<object?> Task => _task;
154+
#else
17155
[PythonType("coroutine")]
18156
[DontMapIDisposableToContextManager, DontMapIEnumerableToContains]
19157
public sealed class PythonCoroutine : ICodeFormattable, IWeakReferenceable {
@@ -98,6 +236,7 @@ public string __qualname__ {
98236
public TaskAwaiter<object?> GetAwaiter() => AsTask().GetAwaiter();
99237

100238
internal PythonGenerator Generator => _generator;
239+
#endif
101240

102241
#region ICodeFormattable Members
103242

0 commit comments

Comments
 (0)