Save and restore the state of stack overflow checks for continuations.

Author: peter-hofer

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/graal/snippets/StackOverflowCheckImpl.java

@@ -140,23 +140,44 @@ public void initialize(IsolateThread thread) {
         yellowZoneStateTL.set(thread, STATE_YELLOW_ENABLED);
     }
 
+    @Override
+    public int getState() {
+        return yellowZoneStateTL.get();
+    }
+
+    @Override
+    @Uninterruptible(reason = "Atomically manipulating state of multiple thread local variables.")
+    public void setState(int newState) {
+        int oldState = yellowZoneStateTL.get();
+        yellowZoneStateTL.set(newState);
+
+        if (newState > oldState) {
+            onYellowZoneMadeAvailable(oldState, newState);
+        } else if (newState < oldState) {
+            onYellowZoneProtected(oldState, newState);
+        }
+    }
+
     @Uninterruptible(reason = "Atomically manipulating state of multiple thread local variables.")
     @Override
     public void makeYellowZoneAvailable() {
-        /*
-         * Even though "yellow zones" and "recurring callbacks" are orthogonal features, running a
-         * recurring callback in the yellow zone is dangerous because a stack overflow in the
-         * recurring callback would then lead to a fatal error.
-         */
-        ThreadingSupportImpl.pauseRecurringCallback("Recurring callbacks are considered user code and must not run in yellow zone");
+        setState(yellowZoneStateTL.get() + 1);
+    }
 
-        int state = yellowZoneStateTL.get();
-        VMError.guarantee(state >= STATE_YELLOW_ENABLED, "StackOverflowSupport.disableYellowZone: Illegal state");
+    @Uninterruptible(reason = "Atomically manipulating state of multiple thread local variables.")
+    private static void onYellowZoneMadeAvailable(int oldState, int newState) {
+        VMError.guarantee(newState > oldState && newState > STATE_YELLOW_ENABLED, "StackOverflowCheckImpl.onYellowZoneMadeAvailable: Illegal state");
+
+        if (oldState == STATE_YELLOW_ENABLED) {
+            /*
+             * Even though "yellow zones" and "recurring callbacks" are orthogonal features, running
+             * a recurring callback in the yellow zone is dangerous because a stack overflow in the
+             * recurring callback would then lead to a fatal error.
+             */
+            ThreadingSupportImpl.pauseRecurringCallback("Recurring callbacks are considered user code and must not run in yellow zone");
 
-        if (state == STATE_YELLOW_ENABLED) {
             stackBoundaryTL.set(stackBoundaryTL.get().subtract(Options.StackYellowZoneSize.getValue()));
         }
-        yellowZoneStateTL.set(state + 1);
 
         /*
          * Check that after enabling the yellow zone there is actually stack space available again.
@@ -180,14 +201,16 @@ public boolean isYellowZoneAvailable() {
     @Uninterruptible(reason = "Atomically manipulating state of multiple thread local variables.")
     @Override
     public void protectYellowZone() {
-        ThreadingSupportImpl.resumeRecurringCallbackAtNextSafepoint();
+        setState(yellowZoneStateTL.get() - 1);
+    }
 
-        int state = yellowZoneStateTL.get();
-        VMError.guarantee(state > STATE_YELLOW_ENABLED, "StackOverflowSupport.enableYellowZone: Illegal state");
+    @Uninterruptible(reason = "Atomically manipulating state of multiple thread local variables.")
+    private static void onYellowZoneProtected(int oldState, int newState) {
+        VMError.guarantee(newState < oldState && newState >= STATE_YELLOW_ENABLED, "StackOverflowCheckImpl.onYellowZoneProtected: Illegal state");
 
-        int newState = state - 1;
-        yellowZoneStateTL.set(newState);
         if (newState == STATE_YELLOW_ENABLED) {
+            ThreadingSupportImpl.resumeRecurringCallbackAtNextSafepoint();
+
             stackBoundaryTL.set(stackBoundaryTL.get().add(Options.StackYellowZoneSize.getValue()));
         }
     }

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/stack/StackOverflowCheck.java

@@ -145,6 +145,15 @@ static StackOverflowCheck singleton() {
      */
     int yellowAndRedZoneSize();
 
+    /** @see #setState */
+    int getState();
+
+    /**
+     * Restore the specified state of the stack overflow checks obtained from {@link #getState}.
+     * This is intended for yielding and resuming continuations on a thread.
+     */
+    void setState(int state);
+
     /**
      * Disables all stack overflow checks for this thread. This operation is not reversible, i.e.,
      * it must only be called in the case of a fatal error where the VM is going to exit soon and

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/thread/Continuation.java

@@ -37,6 +37,7 @@
 import com.oracle.svm.core.heap.StoredContinuation;
 import com.oracle.svm.core.heap.StoredContinuationImpl;
 import com.oracle.svm.core.snippets.KnownIntrinsics;
+import com.oracle.svm.core.stack.StackOverflowCheck;
 import com.oracle.svm.core.util.VMError;
 
 /**
@@ -65,6 +66,7 @@ public static boolean isSupported() {
     Pointer sp;
     CodePointer ip;
     private boolean done;
+    private int overflowCheckState;
 
     Continuation(Runnable target) {
         this.target = target;
@@ -75,7 +77,19 @@ public void setIP(CodePointer ip) {
     }
 
     void enter() {
-        enter0(ip.isNonNull());
+        int stateBefore = StackOverflowCheck.singleton().getState();
+        VMError.guarantee(!StackOverflowCheck.singleton().isYellowZoneAvailable(), "Stack overflow checks must be active when entering a continuation");
+
+        boolean isContinue = ip.isNonNull();
+        if (isContinue) {
+            StackOverflowCheck.singleton().setState(overflowCheckState);
+        }
+        try {
+            enter0(isContinue);
+        } finally {
+            overflowCheckState = StackOverflowCheck.singleton().getState();
+            StackOverflowCheck.singleton().setState(stateBefore);
+        }
     }
 
     @NeverInline("access stack pointer")