[GR-73063] Attach raising thread to PException when escaping.

graalpython/com.oracle.graal.python.test/src/tests/test_traceback.py

@@ -37,9 +37,12 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 import ast
+import concurrent.futures
 import linecache
 import subprocess
 import sys
+import traceback
+import threading
 
 
 def assert_raises(err, fn, *args, **kwargs):
@@ -63,6 +66,34 @@ def test_traceback_dir():
         assert set(dir(tb)) == {'tb_frame', 'tb_next', 'tb_lasti', 'tb_lineno'}
 
 
+def test_traceback_frame_from_threadpool_exception_published_before_workitem_exit():
+    published = threading.Event()
+    release = threading.Event()
+    original_set_exception = concurrent.futures.Future.set_exception
+
+    def set_exception_and_wait(self, exc):
+        original_set_exception(self, exc)
+        published.set()
+        release.wait(timeout=5)
+
+    def target():
+        raise OSError("published before _WorkItem.run exits")
+
+    concurrent.futures.Future.set_exception = set_exception_and_wait
+    try:
+        executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+        try:
+            future = executor.submit(target)
+            assert published.wait(timeout=60)
+            exc = future.exception(timeout=0)
+            traceback.clear_frames(exc.__traceback__)
+        finally:
+            release.set()
+            executor.shutdown(wait=True)
+    finally:
+        concurrent.futures.Future.set_exception = original_set_exception
+
+
 def test_import():
     imported = True
     try:

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/traceback/TracebackBuiltins.java

@@ -270,13 +270,19 @@ static PFrame doOnStack(VirtualFrame frame, PTraceback tb,
             Reference frameInfo = tb.getFrameInfo();
             assert frameInfo.isEscaped() : createAssertionMessage("Cannot create traceback for non-escaped frame", frameInfo);
 
-            PFrame escapedFrame = readCallerFrame.getFrameForReference(frame, frameInfo, ReadFrameNode.AllPythonFramesSelector.INSTANCE, 0, 0);
+            PFrame escapedFrame = readCallerFrame.getFrameForReference(frame, frameInfo, ReadFrameNode.AllPythonFramesSelector.INSTANCE, 0, 0, getEscapedFrameThread(tb));
             assert escapedFrame != null : createAssertionMessage("Failed to find escaped frame on stack", frameInfo);
 
             tb.setFrame(escapedFrame);
             return escapedFrame;
         }
 
+        private static Thread getEscapedFrameThread(PTraceback tb) {
+            LazyTraceback lazyTraceback = tb.getLazyTraceback();
+            PException exception = lazyTraceback != null ? lazyTraceback.getException() : null;
+            return exception != null ? exception.getEscapedFrameThread() : null;
+        }
+
         @TruffleBoundary
         private static String createAssertionMessage(String prefix, Reference frameInfo) {
             String stackTrace;

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/frame/ReadFrameNode.java

@@ -171,13 +171,17 @@ public final PFrame getFrameForReference(Frame frame, PFrame.Reference startFram
     }
 
     public final PFrame getFrameForReference(Frame frame, PFrame.Reference startFrameInfo, FrameSelector selector, int level, int callerFlags) {
-        return execute(frame, startFrameInfo, FrameInstance.FrameAccess.READ_ONLY, selector, level, callerFlags | CallerFlags.NEEDS_PFRAME);
+        return getFrameForReference(frame, startFrameInfo, selector, level, callerFlags, null);
     }
 
-    protected abstract PFrame execute(Frame frame, PFrame.Reference startFrameInfo, FrameInstance.FrameAccess frameAccess, FrameSelector selector, int level, int callerFlags);
+    public final PFrame getFrameForReference(Frame frame, PFrame.Reference startFrameInfo, FrameSelector selector, int level, int callerFlags, Thread frameThread) {
+        return execute(frame, startFrameInfo, FrameInstance.FrameAccess.READ_ONLY, selector, level, callerFlags | CallerFlags.NEEDS_PFRAME, frameThread);
+    }
+
+    protected abstract PFrame execute(Frame frame, PFrame.Reference startFrameInfo, FrameInstance.FrameAccess frameAccess, FrameSelector selector, int level, int callerFlags, Thread frameThread);
 
     @Specialization
-    PFrame read(VirtualFrame frame, PFrame.Reference startFrameInfo, FrameInstance.FrameAccess frameAccess, FrameSelector selector, int level, int callerFlags,
+    PFrame read(VirtualFrame frame, PFrame.Reference startFrameInfo, FrameInstance.FrameAccess frameAccess, FrameSelector selector, int level, int callerFlags, Thread frameThread,
                     @Bind Node inliningTarget,
                     @Cached MaterializeFrameNode materializeFrameNode,
                     @Cached InlinedBranchProfile stackWalkProfile1,
@@ -233,17 +237,16 @@ PFrame read(VirtualFrame frame, PFrame.Reference startFrameInfo, FrameInstance.F
          * It is necessary to continue from where we stopped with the backref walk because the
          * original starting frame might not be on stack anymore
          */
-        return readSlowPath(curFrameInfo, frameAccess, selector, level - i, callerFlags, materializeFrameNode);
+        return readSlowPath(curFrameInfo, frameAccess, selector, level - i, callerFlags, frameThread, materializeFrameNode);
     }
 
     @TruffleBoundary
     @SuppressWarnings("try")
-    private PFrame readSlowPath(PFrame.Reference startFrameInfo, FrameInstance.FrameAccess frameAccess, FrameSelector selector, int level, int callerFlags,
+    private PFrame readSlowPath(PFrame.Reference startFrameInfo, FrameInstance.FrameAccess frameAccess, FrameSelector selector, int level, int callerFlags, Thread frameThread,
                     MaterializeFrameNode materializeFrameNode) {
-        if (level == 0 && startFrameInfo != null && startFrameInfo.getPyFrame() != null && !selector.skip(startFrameInfo.getRootNode()) && startFrameInfo.getPyFrame().getThread() != null &&
-                        startFrameInfo.getPyFrame().getThread() != Thread.currentThread()) {
+        Thread thread = getFrameThread(startFrameInfo, frameThread);
+        if (level == 0 && startFrameInfo != null && !selector.skip(startFrameInfo.getRootNode()) && thread != null && thread != Thread.currentThread()) {
             // We have the frame we're looking for, but it's on another thread
-            Thread thread = startFrameInfo.getPyFrame().getThread();
             if (thread.isAlive()) {
                 try (var gil = GilNode.uncachedRelease()) {
                     // Schedule a safepoint action on that thread
@@ -272,13 +275,23 @@ protected void perform(Access access) {
                     }, future);
                 }
             }
-            assert !startFrameInfo.getPyFrame().outdatedCallerFlags(callerFlags);
-            return startFrameInfo.getPyFrame();
+            PFrame pyFrame = startFrameInfo.getPyFrame();
+            if (pyFrame != null) {
+                assert !pyFrame.outdatedCallerFlags(callerFlags);
+                return pyFrame;
+            }
         }
         StackWalkResult callerFrameResult = getFrame(this, startFrameInfo, frameAccess, selector, level, callerFlags);
         return processStackWalkResult(materializeFrameNode, callerFlags, callerFrameResult);
     }
 
+    private static Thread getFrameThread(PFrame.Reference startFrameInfo, Thread frameThread) {
+        if (startFrameInfo != null && startFrameInfo.getPyFrame() != null) {
+            return startFrameInfo.getPyFrame().getThread();
+        }
+        return frameThread;
+    }
+
     private static PFrame processStackWalkResult(MaterializeFrameNode materializeFrameNode, int callerFlags, StackWalkResult callerFrameResult) {
         if (callerFrameResult != null) {
             Node location = callerFrameResult.callNode;

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/exception/PException.java

@@ -101,6 +101,7 @@ public final class PException extends AbstractTruffleException {
     private String message = null;
     final transient Object pythonException;
     private transient PFrame.Reference frameInfo;
+    private transient Thread escapedFrameThread;
 
     /**
      * Root node that caught this exception object. This node is a manual bytecode or Bytecode DSL
@@ -383,9 +384,16 @@ private void markFrameEscaped() {
         // shouldn't leak to the traceback
         if (this.frameInfo != null) {
             this.frameInfo.markAsEscaped();
+            if (escapedFrameThread == null) {
+                escapedFrameThread = Thread.currentThread();
+            }
         }
     }
 
+    public Thread getEscapedFrameThread() {
+        return escapedFrameThread;
+    }
+
     /**
      * Get the python exception while ensuring that the traceback frame is marked as escaped
      */