Force a Thread's block to optimize eagerly

The block of code passed to Thread.new is frequently only executed
once during a program's runtime, and in such cases there's often
heavy lifting performed within the block such as an IO or Queue
loop. Because we do not support any form of on-stack replacement,
these blocks may never fully optimize and will remain in whatever
execution mode JRuby starts them in (usually the slowest "startup"
interpreter).

This patch adds logic from Thread.new into the block body to force
completion of builds steps that would optimize the body. This is
akin to the behavior of the command-line target script, which we
also eagerly compile due to the great number of single-script
utilities and benchmarks in the Ruby ecosystem.

This patch also plumbs that force-ability through the other
"buildable" execution units: methods and block-based methods like
define_method. This may become useful in the future to have more
direct control over when a target block or method optimizes.

core/src/main/java/org/jruby/RubyThread.java

@@ -67,6 +67,7 @@
 import org.jruby.anno.JRubyClass;
 import org.jruby.anno.JRubyMethod;
 import org.jruby.api.Create;
+import org.jruby.compiler.Compilable;
 import org.jruby.exceptions.MainExitException;
 import org.jruby.exceptions.RaiseException;
 import org.jruby.exceptions.ThreadKill;
@@ -658,6 +659,12 @@ public IRubyObject initialize(ThreadContext context, IRubyObject[] args, Block b
         if (threadImpl != ThreadLike.DUMMY) throw context.runtime.newThreadError("already initialized thread");
 
         BlockBody body = block.getBody();
+
+        // Force top-level body for thread to compile, since it may be executed only once and never get to JIT
+        if (body instanceof Compilable compilable) {
+            compilable.forceBuild(context);
+        }
+
         startThread(context, new RubyRunnable(this, context, args, block, callInfo), body.getFile(), body.getLine());
 
         return context.nil;

core/src/main/java/org/jruby/compiler/Compilable.java

@@ -22,6 +22,30 @@ default void completeBuild(T buildResult) {
     default void completeBuild(ThreadContext context, T buildResult) {
         completeBuild(buildResult);
     }
+
+    /**
+     * Force this Compilable to finish its build process.
+     *
+     * Override this and {@link #isBuildComplete()} for implementations that have a build process.
+     *
+     * @param context the current thread context
+     * @return true if the forced build completed; false otherwise
+     */
+    default boolean forceBuild(ThreadContext context) {
+        return false;
+    }
+
+    /**
+     * Indicates whether this Compilable has completed its build process.
+     *
+     * Override this and {@link #forceBuild(ThreadContext)} for implementations that have a build process.
+     *
+     * @return true if the build has completed; false otherwise.
+     */
+    default boolean isBuildComplete() {
+        return true;
+    }
+
     IRScope getIRScope();
     InterpreterContext ensureInstrsReady();
 

core/src/main/java/org/jruby/compiler/JITCompiler.java

@@ -204,7 +204,7 @@ public Runnable getTaskFor(ThreadContext context, Compilable method) {
         return new FullBuildTask(this, method);
     }
 
-    public void buildThresholdReached(ThreadContext context, final Compilable method) {
+    public void buildThresholdReached(ThreadContext context, final Compilable method, boolean force) {
         final RubyInstanceConfig config = instanceConfig(context);
 
         final Runnable task = getTaskFor(context, method);
@@ -217,7 +217,7 @@ public void buildThresholdReached(ThreadContext context, final Compilable method
         }
 
         try {
-            if (config.getJitBackground() && config.getJitThreshold() > 0) {
+            if (!force && config.getJitBackground() && config.getJitThreshold() > 0) {
                 try {
                     executor.submit(task);
                 } catch (RejectedExecutionException ree) {

core/src/main/java/org/jruby/internal/runtime/AbstractIRMethod.java

@@ -86,28 +86,38 @@ public AbstractIRMethod(StaticScope scope, String id, int line, Visibility visib
         }
     }
 
-    public static <T extends AbstractIRMethod & Compilable> void tryJit(ThreadContext context, T self) {
+    protected static <T extends AbstractIRMethod & Compilable> void tryJit(ThreadContext context, T self, boolean force) {
         if (context.runtime.isBooting() && !Options.JIT_KERNEL.load()) return; // don't JIT during runtime boot
 
         if (self.callCount < 0) return;
         // we don't synchronize callCount++ it does not matter if count isn't accurate
         if (self.callCount++ >= instanceConfig(context).getJitThreshold()) {
             synchronized (self) { // disable same jit tasks from entering queue twice
                 if (self.callCount >= 0) {
-                    self.callCount = Integer.MIN_VALUE; // so that callCount++ stays < 0
-
-                    context.runtime.getJITCompiler().buildThresholdReached(context, self);
+                    build(context, self, false);
                 }
             }
         }
     }
 
+    protected static <T extends AbstractIRMethod & Compilable> void build(ThreadContext context, T self, boolean force) {
+        self.callCount = Integer.MIN_VALUE; // so that callCount++ stays < 0
+
+        context.runtime.getJITCompiler().buildThresholdReached(context, self, force);
+    }
+
     public final void setCallCount(int callCount) {
         synchronized (this) {
             this.callCount = callCount;
         }
     }
 
+    // Overrides method in Compilable but this class does not implement that.
+    public boolean isBuildComplete() {
+        // Successful build and disabled build both set callCount to -1, indicating no further build is possible.
+        return callCount < 0;
+    }
+
     public IRScope getIRScope() {
         try {
             if (method == null) method = staticScope.getIRScope();

core/src/main/java/org/jruby/internal/runtime/methods/InterpretedIRMethod.java

@@ -260,7 +260,7 @@ public void completeBuild(ThreadContext context, InterpreterContext interpreterC
     // executor pool it may take a while
     // and replace interpreterContext asynchronously.
     private void promoteToFullBuild(ThreadContext context) {
-        tryJit(context, this);
+        tryJit(context, this, false);
     }
 
     @Deprecated

core/src/main/java/org/jruby/internal/runtime/methods/MixedModeIRMethod.java

@@ -105,7 +105,7 @@ public IRubyObject call(ThreadContext context, IRubyObject self, RubyModule claz
 
         // try jit before checking actualMethod, so we use jitted version immediately if
         // it's ready
-        if (callCount >= 0) tryJit(context, this);
+        if (callCount >= 0) tryJit(context, this, false);
 
         DynamicMethod jittedMethod = actualMethod;
         if (jittedMethod != null) {
@@ -141,7 +141,7 @@ public IRubyObject call(ThreadContext context, IRubyObject self, RubyModule claz
 
         // try jit before checking actualMethod, so we use jitted version immediately if
         // it's ready
-        if (callCount >= 0) tryJit(context, this);
+        if (callCount >= 0) tryJit(context, this, false);
 
         DynamicMethod jittedMethod = actualMethod;
         if (jittedMethod != null) {
@@ -178,7 +178,7 @@ public IRubyObject call(ThreadContext context, IRubyObject self, RubyModule claz
 
         // try jit before checking actualMethod, so we use jitted version immediately if
         // it's ready
-        if (callCount >= 0) tryJit(context, this);
+        if (callCount >= 0) tryJit(context, this, false);
 
         DynamicMethod jittedMethod = actualMethod;
         if (jittedMethod != null) {
@@ -215,7 +215,7 @@ public IRubyObject call(ThreadContext context, IRubyObject self, RubyModule claz
 
         // try jit before checking actualMethod, so we use jitted version immediately if
         // it's ready
-        if (callCount >= 0) tryJit(context, this);
+        if (callCount >= 0) tryJit(context, this, false);
 
         DynamicMethod jittedMethod = actualMethod;
         if (jittedMethod != null) {
@@ -252,7 +252,7 @@ public IRubyObject call(ThreadContext context, IRubyObject self, RubyModule claz
 
         // try jit before checking actualMethod, so we use jitted version immediately if
         // it's ready
-        if (callCount >= 0) tryJit(context, this);
+        if (callCount >= 0) tryJit(context, this, false);
 
         DynamicMethod jittedMethod = actualMethod;
         if (jittedMethod != null) {
@@ -360,4 +360,12 @@ public DynamicMethod dup() {
         return x;
     }
 
+    @Override
+    public boolean forceBuild(ThreadContext context) {
+        build(context, this, true);
+
+        // Force = true should trigger jit to run synchronously, so we'll be optimistic
+        return true;
+    }
+
 }

core/src/main/java/org/jruby/runtime/InterpretedIRBlockBody.java

@@ -110,9 +110,20 @@ protected IRubyObject yieldDirect(ThreadContext context, Block block, IRubyObjec
         return Interpreter.INTERPRET_BLOCK(context, block, self, fullInterpreterContext, args, block.getBinding().getMethod(), Block.NULL_BLOCK);
     }
 
+    // TODO: Duplicated in InterpretedIRBlockBody
+    private static void tryJIT(InterpretedIRBlockBody body, ThreadContext context) {
+        // don't JIT during runtime boot
+        if (body.callCount >= 0 && (!context.runtime.isBooting() || Options.JIT_KERNEL.load())) {
+            // we don't synchronize callCount++ it does not matter if count isn't accurate
+            if (body.callCount++ >= instanceConfig(context).getJitThreshold()) {
+                body.promoteToFullBuild(context, false);
+            }
+        }
+    }
+
     @Override
     protected IRubyObject commonYieldPath(ThreadContext context, Block block, IRubyObject[] args, IRubyObject self, Block blockArg) {
-        if (callCount >= 0) promoteToFullBuild(context);
+        tryJIT(this, context);
 
         InterpreterContext ic = ensureInstrsReady();
 
@@ -141,20 +152,28 @@ protected IRubyObject commonYieldPath(ThreadContext context, Block block, IRubyO
         }
     }
 
+    @Override
+    public boolean forceBuild(ThreadContext context) {
+        promoteToFullBuild(context, true);
+
+        // Force = true should trigger jit to run synchronously, so we'll be optimistic
+        return true;
+    }
+
+    @Override
+    public boolean isBuildComplete() {
+        // Successful build and disabled build both set callCount to -1, indicating no further build is possible.
+        return callCount < 0;
+    }
+
     // Unlike JIT in MixedMode this will always successfully build but if using executor pool it may take a while
     // and replace interpreterContext asynchronously.
-    private void promoteToFullBuild(ThreadContext context) {
-        if (context.runtime.isBooting() && !Options.JIT_KERNEL.load()) return; // don't JIT during runtime boot
-
-        if (this.callCount < 0) return;
-        // we don't synchronize callCount++ it does not matter if count isn't accurate
-        if (this.callCount++ >= instanceConfig(context).getJitThreshold()) {
-            synchronized (this) { // disable same jit tasks from entering queue twice
-                if (this.callCount >= 0) {
-                    this.callCount = Integer.MIN_VALUE; // so that callCount++ stays < 0
-
-                    context.runtime.getJITCompiler().buildThresholdReached(context, this);
-                }
+    private void promoteToFullBuild(ThreadContext context, boolean force) {
+        synchronized (this) { // disable same jit tasks from entering queue twice
+            if (this.callCount >= 0) {
+                this.callCount = Integer.MIN_VALUE; // so that callCount++ stays < 0
+
+                context.runtime.getJITCompiler().buildThresholdReached(context, this, force);
             }
         }
     }

core/src/main/java/org/jruby/runtime/MixedModeIRBlockBody.java

@@ -141,35 +141,54 @@ protected IRubyObject commonYieldPath(ThreadContext context, Block block, IRubyO
             postYield(context, ic, binding, oldVis, prevFrame);
 
             // trigger JIT on the trailing edge, so we make a best effort to not interpret again after jitting
-            if (callCount >= 0) promoteToFullBuild(context);
+            tryJIT(this, context);
         }
     }
 
-    private void promoteToFullBuild(ThreadContext context) {
-        if (context.runtime.isBooting() && !Options.JIT_KERNEL.load()) return; // don't JIT during runtime boot
+    // TODO: Duplicated in InterpretedIRBlockBody
+    private static void tryJIT(MixedModeIRBlockBody body, ThreadContext context) {
+        // don't JIT during runtime boot
+        if (body.callCount >= 0 && (!context.runtime.isBooting() || Options.JIT_KERNEL.load())) {
+            // we don't synchronize callCount++ it does not matter if count isn't accurate
+            if (body.callCount++ >= instanceConfig(context).getJitThreshold()) {
+                body.promoteToFullBuild(context, false);
+            }
+        }
+    }
+
+    @Override
+    public boolean forceBuild(ThreadContext context) {
+        promoteToFullBuild(context, true);
+
+        // Force = true should trigger jit to run synchronously, so we'll be optimistic
+        return true;
+    }
+
+    @Override
+    public boolean isBuildComplete() {
+        // Successful build and disabled build both set callCount to -1, indicating no further build is possible.
+        return callCount < 0;
+    }
 
-        if (this.callCount < 0) return;
-        // we don't synchronize callCount++ it does not matter if count isn't accurate
-        if (this.callCount++ >= instanceConfig(context).getJitThreshold()) {
-            synchronized (this) { // disable same jit tasks from entering queue twice
-                if (this.callCount >= 0) {
-                    this.callCount = Integer.MIN_VALUE; // so that callCount++ stays < 0
+    private void promoteToFullBuild(ThreadContext context, boolean force) {
+        synchronized (this) { // disable same jit tasks from entering queue twice
+            if (this.callCount >= 0) {
+                this.callCount = Integer.MIN_VALUE; // so that callCount++ stays < 0
 
-                    // ensure we've got code ready for JIT
-                    ensureInstrsReady();
-                    closure.getNearestTopLocalVariableScope().prepareForCompilation();
+                // ensure we've got code ready for JIT
+                ensureInstrsReady();
+                closure.getNearestTopLocalVariableScope().prepareForCompilation();
 
-                    FullInterpreterContext fic = closure.getFullInterpreterContext();
+                FullInterpreterContext fic = closure.getFullInterpreterContext();
 
-                    if (fic == null || !fic.hasExplicitCallProtocol()) {
-                        if (Options.JIT_LOGGING.load()) {
-                            LOG.info("JIT failed; no full IR or no call protocol found in block: " + closure);
-                        }
-                        return; // do not JIT if we don't have an explicit protocol
+                if (fic == null || !fic.hasExplicitCallProtocol()) {
+                    if (Options.JIT_LOGGING.load()) {
+                        LOG.info("JIT failed; no full IR or no call protocol found in block: " + closure);
                     }
-
-                    context.runtime.getJITCompiler().buildThresholdReached(context, this);
+                    return; // do not JIT if we don't have an explicit protocol
                 }
+
+                context.runtime.getJITCompiler().buildThresholdReached(context, this, force);
             }
         }
     }