From 2e54cf174ef14edf8a423549001d3e3e6f9c73f9 Mon Sep 17 00:00:00 2001 From: Pavel Koneski Date: Wed, 4 Nov 2026 20:21:15 -0800 Subject: [PATCH] Keep generator frames on the thread running the body --- .../IronPython/Compiler/Ast/AstMethods.cs | 1 + .../Compiler/Ast/FunctionDefinition.cs | 26 ++-- src/core/IronPython/Compiler/Ast/Node.cs | 31 ++++- src/core/IronPython/Runtime/Generator.cs | 55 ++++++--- .../Runtime/Operations/PythonOps.cs | 8 ++ tests/suite/test_generator.py | 112 ++++++++++++++++++ 6 files changed, 207 insertions(+), 26 deletions(-) diff --git a/src/core/IronPython/Compiler/Ast/AstMethods.cs b/src/core/IronPython/Compiler/Ast/AstMethods.cs index ac64c8eb3..67ed3c147 100644 --- a/src/core/IronPython/Compiler/Ast/AstMethods.cs +++ b/src/core/IronPython/Compiler/Ast/AstMethods.cs @@ -77,6 +77,7 @@ internal static class AstMethods { public static readonly MethodInfo ForLoopDispose = GetMethod((Action>)PythonOps.ForLoopDispose); public static readonly MethodInfo GetClosureTupleFromContext = GetMethod((Func)PythonOps.GetClosureTupleFromContext); public static readonly MethodInfo PushFrame = GetMethod((Func>)PythonOps.PushFrame); + public static readonly MethodInfo PushGeneratorFrame = GetMethod((Action)PythonOps.PushGeneratorFrame); public static readonly MethodInfo FormatString = GetMethod((Func)PythonOps.FormatString); public static readonly MethodInfo GeneratorCheckThrowableAndReturnSendValue = GetMethod((Func)PythonOps.GeneratorCheckThrowableAndReturnSendValue); public static readonly MethodInfo MakeCoroutine = GetMethod((Func)PythonOps.MakeCoroutine); diff --git a/src/core/IronPython/Compiler/Ast/FunctionDefinition.cs b/src/core/IronPython/Compiler/Ast/FunctionDefinition.cs index 13b31ae4b..4aff6103b 100644 --- a/src/core/IronPython/Compiler/Ast/FunctionDefinition.cs +++ b/src/core/IronPython/Compiler/Ast/FunctionDefinition.cs @@ -686,16 +686,19 @@ private LightLambdaExpression CreateFunctionLambda() { statements.Add(MakeFunctionEntrySequencePoint(StartIndex)); - // For generators/coroutines, we need to do a check before the first statement for Generator.Throw() / Generator.Close(). - // The exception traceback needs to come from the generator's method body, and so we must do the check and throw - // from inside the generator. + // Whether the body runs in a PythonGenerator, which drives it through the generator rewriter's state machine. #if FEATURE_NET_ASYNC - // Async generators have no backing PythonGenerator (they lower to IAsyncEnumerable via AsyncEnumerableExpression), - // so skip the $generator.CheckThrowable() prologue for them. - if (IsGenerator && !IsAsync) { + // Async functions and async generators have no backing PythonGenerator: they lower to a Task via AsyncExpression, + // or to an IAsyncEnumerable via AsyncEnumerableExpression. + bool inPythonGenerator = IsGenerator && !IsAsync; #else - if (IsGenerator || IsAsync) { + bool inPythonGenerator = IsGenerator || IsAsync; #endif + + // For generators/coroutines, we need to do a check before the first statement for Generator.Throw() / Generator.Close(). + // The exception traceback needs to come from the generator's method body, and so we must do the check and throw + // from inside the generator. + if (inPythonGenerator) { MSAst.Expression s1 = YieldExpression.CreateCheckThrowExpression(SourceSpan.None); statements.Add(s1); } @@ -708,8 +711,13 @@ private LightLambdaExpression CreateFunctionLambda() { MSAst.Expression body = Ast.Block(statements); if (Body.CanThrow && GlobalParent.PyContext.PythonOptions.Frames) { - body = AddFrame(LocalContext, Ast.Property(_functionParam, typeof(PythonFunction).GetProperty(nameof(PythonFunction.__code__))), body); - locals.Add(FunctionStackVariable); + MSAst.Expression codeObject = Ast.Property(_functionParam, typeof(PythonFunction).GetProperty(nameof(PythonFunction.__code__))); + if (inPythonGenerator) { + body = AddGeneratorFrame(LocalContext, codeObject, body); + } else { + body = AddFrame(LocalContext, codeObject, body); + locals.Add(FunctionStackVariable); + } } body = AddProfiling(body); diff --git a/src/core/IronPython/Compiler/Ast/Node.cs b/src/core/IronPython/Compiler/Ast/Node.cs index c15820c82..d59808229 100644 --- a/src/core/IronPython/Compiler/Ast/Node.cs +++ b/src/core/IronPython/Compiler/Ast/Node.cs @@ -317,7 +317,18 @@ internal static MSAst.Expression RemoveDebugInfo(int prevStart, MSAst.Expression /// Creates a method frame for tracking purposes and enforces recursion /// internal static MSAst.Expression AddFrame(MSAst.Expression localContext, MSAst.Expression codeObject, MSAst.Expression body) - => new FramedCodeExpression(localContext, codeObject, body); + => new FramedCodeExpression(localContext, codeObject, body, inGenerator: false); + + /// + /// Creates a method frame for the body of a generator, which hands the frame over to its . + /// + /// + /// A generator body can suspend and resume on another thread, so it cannot pop its frame itself: it would pop it from the thread + /// it started on. The generator removes the frame from the current thread at each suspension, adds it back at each resumption, + /// and removes it when the body finishes. + /// + internal static MSAst.Expression AddGeneratorFrame(MSAst.Expression localContext, MSAst.Expression codeObject, MSAst.Expression body) + => new FramedCodeExpression(localContext, codeObject, body, inGenerator: true); /// /// Removes the frames from generated code for when we're compiling the tracing delegate @@ -337,10 +348,12 @@ private class FramedCodeVisitor : ExpressionVisitor { private sealed class FramedCodeExpression : MSAst.Expression { private readonly MSAst.Expression _localContext, _codeObject; + private readonly bool _inGenerator; - public FramedCodeExpression(MSAst.Expression localContext, MSAst.Expression codeObject, MSAst.Expression body) { + public FramedCodeExpression(MSAst.Expression localContext, MSAst.Expression codeObject, MSAst.Expression body, bool inGenerator) { _localContext = localContext; _codeObject = codeObject; + _inGenerator = inGenerator; Body = body; } @@ -349,6 +362,18 @@ public FramedCodeExpression(MSAst.Expression localContext, MSAst.Expression code public MSAst.Expression Body { get; } public override MSAst.Expression Reduce() { + if (_inGenerator) { + return Ast.Block( + Ast.Call( + AstMethods.PushGeneratorFrame, + _localContext, + GeneratorRewriter._generatorParam, + _codeObject + ), + Body + ); + } + return AstUtils.Try( Ast.Assign( FunctionStackVariable, @@ -384,7 +409,7 @@ protected override MSAst.Expression VisitChildren(ExpressionVisitor visitor) { var body = visitor.Visit(Body); if (localContext != _localContext || _codeObject != codeObject || body != Body) { - return new FramedCodeExpression(localContext, codeObject, body); + return new FramedCodeExpression(localContext, codeObject, body, _inGenerator); } return this; diff --git a/src/core/IronPython/Runtime/Generator.cs b/src/core/IronPython/Runtime/Generator.cs index 4cb008626..d2dc5f510 100644 --- a/src/core/IronPython/Runtime/Generator.cs +++ b/src/core/IronPython/Runtime/Generator.cs @@ -356,27 +356,54 @@ bool IEnumerator.MoveNext() { } } - private FunctionStack? fnStack = null; + /// + /// The body's entry on the frame stack, from the moment the body pushes it until the body finishes. + /// + /// + /// The body can suspend on one thread and resume on another, so the generator owns the entry once the body has pushed it: + /// the entry is on the frame stack of the thread running the body, and on no thread's stack while the body is suspended. + /// + /// It stays null for a body that doesn't push its own entry: with frames disabled, or when the body was compiled for tracing, + /// where the trace listener pushes and pops an entry at each step. This is decided when the body is compiled, so it doesn't change + /// when tracing is switched on or off while the generator is suspended. + /// + private FunctionStack? _stackEntry; + + /// + /// Called by the body when it starts, to push its entry on the current thread's frame stack and hand the entry over to the generator. + /// + internal void PushFrame(CodeContext context, FunctionCode code) { + Debug.Assert(_stackEntry is null, "the body of a generator starts only once"); + + // The entry never keeps a frame object, since the frame's f_back changes whenever the body resumes. + var entry = new FunctionStack(context, code); + PythonOps.GetFunctionStack().Add(entry); + _stackEntry = entry; + } private void SaveFunctionStack(bool done) { - if (!Context.LanguageContext.PythonOptions.Frames || Context.LanguageContext.EnableTracing) return; - if (!done) { - var stack = PythonOps.GetFunctionStack(); - var functionStack = stack[stack.Count - 1]; - stack.RemoveAt(stack.Count - 1); - Debug.Assert(functionStack.Context is not null); - fnStack = new FunctionStack(functionStack.Context!, functionStack.Code); // don't keep the frame since f_back may be invalid + if (_stackEntry is not { } entry) return; + + // The entry is on top, unless something the body called has left an entry of its own behind. Then remove only this + // body's entry, if it is there at all. sys._getframe replaces entries by copies with frame objects, so compare the parts. + var stack = PythonOps.GetFunctionStack(); + int index = stack.Count - 1; + while (index >= 0 && !(stack[index].Code == entry.Code && stack[index].Context == entry.Context)) { + index--; } - else { - fnStack = null; + Debug.Assert(index >= 0 && index == stack.Count - 1, "the entry of a suspending generator body should be on top of the frame stack"); + if (index >= 0) { + stack.RemoveAt(index); + } + + if (done) { + _stackEntry = null; } } private void RestoreFunctionStack() { - if (!Context.LanguageContext.PythonOptions.Frames) return; - if (fnStack != null) { - var stack = PythonOps.GetFunctionStack(); - stack.Add(fnStack.Value); + if (_stackEntry is { } entry) { + PythonOps.GetFunctionStack().Add(entry); } } diff --git a/src/core/IronPython/Runtime/Operations/PythonOps.cs b/src/core/IronPython/Runtime/Operations/PythonOps.cs index bdc02b079..abe7a066a 100644 --- a/src/core/IronPython/Runtime/Operations/PythonOps.cs +++ b/src/core/IronPython/Runtime/Operations/PythonOps.cs @@ -4421,6 +4421,14 @@ public static List PushFrame(CodeContext/*!*/ context, FunctionCo return stack; } + /// + /// Pushes the frame of a generator body on the current thread, and hands it over to , which removes + /// it at each suspension, adds it back at each resumption, and removes it when the body finishes. + /// + public static void PushGeneratorFrame(CodeContext/*!*/ context, PythonGenerator generator, FunctionCode function) { + generator.PushFrame(context, function); + } + internal static LightLambdaExpression ToGenerator(this LightLambdaExpression code, bool shouldInterpret, bool debuggable, int compilationThreshold, bool isCoroutine = false) { #pragma warning disable CA2263 // Prefer generic overload when type is known return Utils.LightLambda( diff --git a/tests/suite/test_generator.py b/tests/suite/test_generator.py index aec2c049c..21c991960 100644 --- a/tests/suite/test_generator.py +++ b/tests/suite/test_generator.py @@ -2,8 +2,21 @@ # The .NET Foundation licenses this file to you under the Apache 2.0 License. # See the LICENSE file in the project root for more information. +import sys +import threading + from iptest import IronPythonTestCase, is_cli, run_test +def _frame_names(): + """The names of the calling thread's frames, outermost first, without the frame of this function.""" + frame = sys._getframe(1) + names = [] + while frame is not None: + names.append(frame.f_code.co_name) + frame = frame.f_back + names.reverse() + return names + def ifilter(iterable): def predicate(x): return x % 3 @@ -583,4 +596,103 @@ def reraises(): next(g) self.assertEqual(cm.exception.args, ('k',)) + def test_generator_finished_on_other_thread(self): + # A generator started on a thread that has exited can be finished on another thread, whichever way it finishes + def gen(fail): + yield 1 + yield 2 + if fail: + raise KeyError('fail') + + def started_on_other_thread(g): + t = threading.Thread(target=lambda: next(g)) + t.start() + t.join() + return g + + expected = _frame_names() + + g = started_on_other_thread(gen(False)) + self.assertEqual(list(g), [2]) + self.assertEqual(_frame_names(), expected) + + g = started_on_other_thread(gen(True)) + self.assertEqual(next(g), 2) + self.assertRaises(KeyError, next, g) + self.assertEqual(_frame_names(), expected) + + g = started_on_other_thread(gen(False)) + g.close() + self.assertEqual(_frame_names(), expected) + + def test_generator_finished_while_starting_thread_runs(self): + # A generator finished on another thread than the one it started on leaves the frames of both threads intact + def gen(): + yield 1 + yield 2 + + def finish(g): + return list(g) + + g = gen() + started, done = threading.Event(), threading.Event() + seen = [] + + def wait_here(): + started.set() + done.wait() + seen.append(_frame_names()[-2:]) + + def worker(): + next(g) # the generator starts on this thread + wait_here() # and the thread keeps running Python code + + expected = _frame_names() + t = threading.Thread(target=worker) + t.start() + started.wait() + try: + self.assertEqual(finish(g), [2]) # the generator finishes on this thread + self.assertEqual(_frame_names(), expected) + finally: + done.set() + t.join() + self.assertEqual(seen, [['worker', 'wait_here']]) + + def test_generator_resumed_after_settrace(self): + # A generator suspended while tracing is switched on or off leaves the frames of the thread resuming it intact + def gen(): + yield 1 + yield 2 + + def tracer(frame, event, arg): + return None + + def resume(g): + next(g) + return _frame_names() + + expected = _frame_names() + ['resume'] + + # started untraced, resumed traced + g = gen() + next(g) + sys.settrace(tracer) + try: + names = resume(g) + finally: + sys.settrace(None) + self.assertEqual(names, expected) + self.assertEqual(_frame_names(), expected[:-1]) + + # started traced, resumed untraced + sys.settrace(tracer) + try: + g = gen() + next(g) + finally: + sys.settrace(None) + self.assertEqual(resume(g), expected) + self.assertEqual(_frame_names(), expected[:-1]) + run_test(__name__)