Keep generator frames on the thread running the body

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

@@ -77,6 +77,7 @@ internal static class AstMethods {
         public static readonly MethodInfo ForLoopDispose = GetMethod((Action<KeyValuePair<IEnumerator, IDisposable>>)PythonOps.ForLoopDispose);
         public static readonly MethodInfo GetClosureTupleFromContext = GetMethod((Func<CodeContext, MutableTuple>)PythonOps.GetClosureTupleFromContext);
         public static readonly MethodInfo PushFrame = GetMethod((Func<CodeContext, FunctionCode, List<FunctionStack>>)PythonOps.PushFrame);
+        public static readonly MethodInfo PushGeneratorFrame = GetMethod((Action<CodeContext, PythonGenerator, FunctionCode>)PythonOps.PushGeneratorFrame);
         public static readonly MethodInfo FormatString = GetMethod((Func<CodeContext, string, object, string>)PythonOps.FormatString);
         public static readonly MethodInfo GeneratorCheckThrowableAndReturnSendValue = GetMethod((Func<object, object>)PythonOps.GeneratorCheckThrowableAndReturnSendValue);
         public static readonly MethodInfo MakeCoroutine = GetMethod((Func<PythonFunction, MutableTuple, object, PythonCoroutine>)PythonOps.MakeCoroutine);

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);

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
         /// </summary>
         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);
+
+        /// <summary>
+        ///   Creates a method frame for the body of a generator, which hands the frame over to its <see cref="PythonGenerator"/>.
+        /// </summary>
+        /// <remarks>
+        ///   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.
+        /// </remarks>
+        internal static MSAst.Expression AddGeneratorFrame(MSAst.Expression localContext, MSAst.Expression codeObject, MSAst.Expression body)
+            => new FramedCodeExpression(localContext, codeObject, body, inGenerator: true);
 
         /// <summary>
         /// 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;

src/core/IronPython/Runtime/Generator.cs

@@ -356,27 +356,54 @@ bool IEnumerator.MoveNext() {
             }
         }
 
-        private FunctionStack? fnStack = null;
+        /// <summary>
+        ///   The body's entry on the frame stack, from the moment the body pushes it until the body finishes.
+        /// </summary>
+        /// <remarks>
+        ///   <para>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.</para>
+        ///
+        ///   <para>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.</para>
+        /// </remarks>
+        private FunctionStack? _stackEntry;
+
+        /// <summary>
+        ///   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.
+        /// </summary>
+        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);
             }
         }
 

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

@@ -4421,6 +4421,14 @@ public static List<FunctionStack> PushFrame(CodeContext/*!*/ context, FunctionCo
             return stack;
         }
 
+        /// <summary>
+        ///   Pushes the frame of a generator body on the current thread, and hands it over to <paramref name="generator"/>, which removes
+        ///   it at each suspension, adds it back at each resumption, and removes it when the body finishes.
+        /// </summary>
+        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(

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__)