Forward instrumentation materialization calls from TopLevelExceptionHandler

Author: msimacek

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/InstrumentationMaterializationForwarder.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright (c) 2022, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * The Universal Permissive License (UPL), Version 1.0
+ *
+ * Subject to the condition set forth below, permission is hereby granted to any
+ * person obtaining a copy of this software, associated documentation and/or
+ * data (collectively the "Software"), free of charge and under any and all
+ * copyright rights in the Software, and any and all patent rights owned or
+ * freely licensable by each licensor hereunder covering either (i) the
+ * unmodified Software as contributed to or provided by such licensor, or (ii)
+ * the Larger Works (as defined below), to deal in both
+ *
+ * (a) the Software, and
+ *
+ * (b) any piece of software and/or hardware listed in the lrgrwrks.txt file if
+ * one is included with the Software each a "Larger Work" to which the Software
+ * is contributed by such licensors),
+ *
+ * without restriction, including without limitation the rights to copy, create
+ * derivative works of, display, perform, and distribute the Software and make,
+ * use, sell, offer for sale, import, export, have made, and have sold the
+ * Software and the Larger Work(s), and to sublicense the foregoing rights on
+ * either these or other terms.
+ *
+ * This license is subject to the following condition:
+ *
+ * The above copyright notice and either this complete permission notice or at a
+ * minimum a reference to the UPL must be included in all copies or substantial
+ * portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+package com.oracle.graal.python.nodes.bytecode;
+
+import java.util.Set;
+
+import com.oracle.truffle.api.instrumentation.InstrumentableNode;
+import com.oracle.truffle.api.instrumentation.ProbeNode;
+import com.oracle.truffle.api.instrumentation.Tag;
+import com.oracle.truffle.api.nodes.Node;
+
+class InstrumentationMaterializationForwarder extends Node implements InstrumentableNode {
+    private final InstrumentableNode delegate;
+
+    InstrumentationMaterializationForwarder(InstrumentableNode delegate) {
+        this.delegate = delegate;
+    }
+
+    @Override
+    public boolean isInstrumentable() {
+        return true;
+    }
+
+    @Override
+    public WrapperNode createWrapper(ProbeNode probe) {
+        return new Wrapper(this, probe);
+    }
+
+    @Override
+    public InstrumentableNode materializeInstrumentableNodes(Set<Class<? extends Tag>> materializedTags) {
+        delegate.materializeInstrumentableNodes(materializedTags);
+        return this;
+    }
+
+    private static class Wrapper extends Node implements WrapperNode {
+        @Child private Node delegateNode;
+        @Child private ProbeNode probeNode;
+
+        Wrapper(Node delegateNode, ProbeNode probeNode) {
+            this.delegateNode = delegateNode;
+            this.probeNode = probeNode;
+        }
+
+        @Override
+        public Node getDelegateNode() {
+            return delegateNode;
+        }
+
+        @Override
+        public ProbeNode getProbeNode() {
+            return probeNode;
+        }
+    }
+}

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeRootNode.java

@@ -2450,6 +2450,10 @@ public void materializeContainedFunctionsForInstrumentation(Set<Class<? extends
         });
     }
 
+    public Node createInstrumentationMaterializationForwarder() {
+        return new InstrumentationMaterializationForwarder(instrumentationRoot);
+    }
+
     @InliningCutoff // Used only to print expressions in interactive mode
     private int bytecodePrintExpr(VirtualFrame virtualFrame, boolean useCachedNodes, int stackTop, int bci, Node[] localNodes, int bciSlot, int beginBci) {
         setCurrentBci(virtualFrame, bciSlot, bci);

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/control/TopLevelExceptionHandler.java

@@ -58,6 +58,7 @@
 import com.oracle.graal.python.lib.PyObjectStrAsObjectNode;
 import com.oracle.graal.python.nodes.BuiltinNames;
 import com.oracle.graal.python.nodes.ErrorMessages;
+import com.oracle.graal.python.nodes.bytecode.PBytecodeRootNode;
 import com.oracle.graal.python.nodes.object.GetClassNode;
 import com.oracle.graal.python.nodes.object.GetDictIfExistsNode;
 import com.oracle.graal.python.nodes.object.IsBuiltinClassProfile;
@@ -79,6 +80,7 @@
 import com.oracle.truffle.api.interop.ExceptionType;
 import com.oracle.truffle.api.interop.InteropLibrary;
 import com.oracle.truffle.api.interop.UnsupportedMessageException;
+import com.oracle.truffle.api.nodes.Node;
 import com.oracle.truffle.api.nodes.RootNode;
 import com.oracle.truffle.api.source.Source;
 import com.oracle.truffle.api.source.SourceSection;
@@ -91,12 +93,24 @@ public final class TopLevelExceptionHandler extends RootNode {
 
     @Child private GilNode gilNode = GilNode.create();
 
+    /*
+     * Necessary to forward materializeInstrumentableNodes calls to the root node. Truffle assumes
+     * that all source sections of a source are created when the first call target is initialized.
+     * That may sometimes happen before the bytecode root's call target is initialized, so we need
+     * to have a node that will trigger the materialization of all nested functions that haven't
+     * been created yet at this point.
+     */
+    @Child private Node instrumentationForwarder;
+
     public TopLevelExceptionHandler(PythonLanguage language, RootNode child, Source source) {
         super(language);
         this.sourceSection = child.getSourceSection();
         this.innerCallTarget = PythonUtils.getOrCreateCallTarget(child);
         this.exception = null;
         this.source = source;
+        if (child instanceof PBytecodeRootNode) {
+            instrumentationForwarder = ((PBytecodeRootNode) child).createInstrumentationMaterializationForwarder();
+        }
     }
 
     public TopLevelExceptionHandler(PythonLanguage language, PException exception) {