(PE-37376) Parallelize JRuby instance creation

.github/workflows/lein-test.yaml

@@ -19,14 +19,14 @@ jobs:
         java: [ '11', '17' ]
     steps:
       - name: Check out repository code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Setup java
-        uses: actions/setup-java@v3
+        uses: actions/setup-java@v4
         with:
           distribution: temurin
           java-version: ${{ matrix.java }}
       - name: Install Clojure tools
-        uses: DeLaGuardo/setup-clojure@10.2
+        uses: DeLaGuardo/setup-clojure@12.5
         with:
           cli: latest              # Clojure CLI based on tools.deps
           lein: latest                  # Leiningen

src/clj/puppetlabs/services/jruby_pool_manager/impl/jruby_agents.clj

@@ -7,11 +7,21 @@
             [puppetlabs.i18n.core :as i18n])
   (:import (clojure.lang IFn IDeref)
            (puppetlabs.services.jruby_pool_manager.jruby_schemas PoisonPill JRubyInstance)
-           (java.util.concurrent TimeUnit TimeoutException)))
+           (java.util.concurrent TimeUnit TimeoutException ExecutionException Future ExecutorService)))
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;;; Private
 
+(schema/defn execute-tasks!
+  [tasks :- [IFn]
+   task-executor :- ExecutorService]
+  (let [results (.invokeAll task-executor tasks)]
+    (try
+      (doseq [result results]
+        (.get ^Future result))
+      (catch ExecutionException ex
+        (throw (.getCause ex))))))
+
 (schema/defn ^:always-validate
   next-instance-id :- schema/Int
   [id :- schema/Int
@@ -68,15 +78,17 @@
                      (i18n/trs "Initializing JRubyInstances with the following settings:")
                      (ks/pprint-to-string config)))
   (let [pool (jruby-internal/get-pool pool-context)
-        count (.remainingCapacity pool)]
-    (dotimes [i count]
-      (let [id (inc i)]
-        (log/debug (i18n/trs "Priming JRubyInstance {0} of {1}"
-                             id count))
-        (add-instance pool-context id)
-        (log/info (i18n/trs "Finished creating JRubyInstance {0} of {1}"
-                            id count))))))
-
+        creation-service (jruby-internal/get-creation-service pool-context)
+        total (.remainingCapacity pool)
+        ids (->> total range (map inc))
+        add-instance* (fn [id]
+                          (log/debug (i18n/trs "Priming JRubyInstance {0} of {1}"
+                                               id count))
+                          (add-instance pool-context id)
+                          (log/info (i18n/trs "Finished creating JRubyInstance {0} of {1}"
+                                              id count)))
+        tasks (for [id ids] (fn [] (add-instance* id)))]
+    (execute-tasks! tasks creation-service)))
 
 (schema/defn ^:always-validate
   flush-instance!
@@ -148,23 +160,28 @@
    refill? :- schema/Bool]
   (let [pool (jruby-internal/get-pool pool-context)
         pool-size (jruby-internal/get-pool-size pool-context)
+        creation-service (jruby-internal/get-creation-service pool-context)
         new-instance-ids (map inc (range pool-size))
         config (:config pool-context)
-        cleanup-fn (get-in config [:lifecycle :cleanup])]
-    (doseq [[old-instance new-id] (zipmap old-instances new-instance-ids)]
-      (try
-        (jruby-internal/cleanup-pool-instance! old-instance cleanup-fn)
-        (when refill?
-          (jruby-internal/create-pool-instance! pool new-id config
-                                                (:splay-instance-flush config))
-          (log/info (i18n/trs "Finished creating JRubyInstance {0} of {1}"
-                               new-id pool-size)))
-        (catch Exception e
-          (.clear pool)
-          (jruby-internal/insert-poison-pill pool e)
-          (throw (IllegalStateException.
-                  (i18n/trs "There was a problem creating a JRubyInstance for the pool.")
-                  e))))))
+        cleanup-fn (get-in config [:lifecycle :cleanup])
+        cleanup-and-refill-instance
+          (fn [old-instance new-id]
+              (try
+                (jruby-internal/cleanup-pool-instance! old-instance cleanup-fn)
+                (when refill?
+                  (jruby-internal/create-pool-instance! pool new-id config
+                                                        (:splay-instance-flush config))
+                  (log/info (i18n/trs "Finished creating JRubyInstance {0} of {1}"
+                                       new-id pool-size)))
+                (catch Exception e
+                  (.clear pool)
+                  (jruby-internal/insert-poison-pill pool e)
+                  (throw (IllegalStateException.
+                          (i18n/trs "There was a problem creating a JRubyInstance for the pool.")
+                          e)))))
+        cleanup-and-refill-tasks (for [[old-instance new-id] (zipmap old-instances new-instance-ids)]
+                                   (fn [] (cleanup-and-refill-instance old-instance new-id)))]
+    (execute-tasks! cleanup-and-refill-tasks creation-service))
   (if refill?
     (log/info (i18n/trs "Finished draining and refilling pool."))
     (log/info (i18n/trs "Finished draining pool."))))

src/clj/puppetlabs/services/jruby_pool_manager/impl/jruby_internal.clj

@@ -15,7 +15,7 @@
            (com.puppetlabs.jruby_utils.jruby InternalScriptingContainer
                                              ScriptingContainer)
            (java.io File)
-           (java.util.concurrent TimeUnit)
+           (java.util.concurrent TimeUnit Executors ExecutorService)
            (org.jruby CompatVersion Main Ruby RubyInstanceConfig RubyInstanceConfig$CompileMode RubyInstanceConfig$ProfilingMode)
            (org.jruby.embed LocalContextScope)
            (org.jruby.runtime.profile.builtin ProfileOutput)
@@ -183,12 +183,16 @@
   "Create a new PoolState based on the config input."
   [config :- jruby-schemas/JRubyConfig]
   (let [multithreaded (:multithreaded config)
-        size (:max-active-instances config)]
+        size (:max-active-instances config)
+        creation-concurrency (get config :instance-creation-concurrency 4)
+        creation-service (Executors/newFixedThreadPool creation-concurrency)]
     (if multithreaded
       {:pool (instantiate-reference-pool size)
-       :size 1}
+       :size 1
+       :creation-service creation-service}
       {:pool (instantiate-instance-pool size)
-       :size size})))
+       :size size
+       :creation-service creation-service})))
 
 (schema/defn ^:always-validate
   cleanup-pool-instance!
@@ -281,6 +285,12 @@
   [context :- jruby-schemas/PoolContext]
   (:size (get-pool-state context)))
 
+(schema/defn
+  get-creation-service :- ExecutorService
+  "Gets the ExecutorService that will execute instance creation and termination."
+  [context :- jruby-schemas/PoolContext]
+  (:creation-service (get-pool-state context)))
+
 (schema/defn ^:always-validate
   get-flush-timeout :- schema/Int
   "Gets the size of the JRuby pool from the pool context."

src/clj/puppetlabs/services/jruby_pool_manager/jruby_schemas.clj

@@ -1,9 +1,10 @@
 (ns puppetlabs.services.jruby-pool-manager.jruby-schemas
   (:require [schema.core :as schema])
   (:import (clojure.lang Atom Agent IFn PersistentArrayMap PersistentHashMap)
+           (com.puppetlabs.jruby_utils.jruby ScriptingContainer)
            (com.puppetlabs.jruby_utils.pool LockablePool)
-           (org.jruby Main Main$Status RubyInstanceConfig)
-           (com.puppetlabs.jruby_utils.jruby ScriptingContainer)))
+           (java.util.concurrent ExecutorService)
+           (org.jruby Main Main$Status RubyInstanceConfig)))
 
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
@@ -72,7 +73,14 @@
 
     * :profiler-output-file - A target file to direct profiler output to. If
         not set, defaults to a random file relative to the working directory
-        of the service."
+        of the service.
+
+    * :multithreaded - Instead of managing the number of JRuby Instances create
+        a single JRuby instance and manage the number of threads that may access it.
+
+    * :instance-creation-concurrency - How many instances to create at once. This
+        will improve start up and potentially reload times, but if too high may
+        create unaceptable load on the system during startup or reload."
   {:ruby-load-path [schema/Str]
    :gem-home schema/Str
    :gem-path (schema/maybe schema/Str)
@@ -86,7 +94,8 @@
    :environment-vars {schema/Keyword schema/Str}
    :profiling-mode SupportedJRubyProfilingModes
    :profiler-output-file schema/Str
-   :multithreaded schema/Bool})
+   :multithreaded schema/Bool
+   (schema/optional-key :instance-creation-concurrency) schema/Int})
 
 (def JRubyPoolAgent
   "An agent configured for use in managing JRuby pools"
@@ -100,8 +109,9 @@
 
 (def PoolState
   "A map that describes all attributes of a particular JRuby pool."
-  {:pool         pool-queue-type
-   :size schema/Int})
+  {:pool             pool-queue-type
+   :size             schema/Int
+   :creation-service ExecutorService})
 
 (def PoolStateContainer
   "An atom containing the current state of all of the JRuby pool."

test/unit/puppetlabs/services/jruby_pool_manager/jruby_agents_test.clj

@@ -4,11 +4,28 @@
             [puppetlabs.services.jruby-pool-manager.jruby-testutils :as jruby-testutils]
             [puppetlabs.services.jruby-pool-manager.jruby-core :as jruby-core]
             [puppetlabs.services.jruby-pool-manager.impl.jruby-agents :as jruby-agents]
+            [puppetlabs.services.jruby-pool-manager.impl.jruby-internal :as jruby-internal]
             [puppetlabs.services.jruby-pool-manager.impl.jruby-pool-manager-core :as jruby-pool-manager-core])
   (:import (puppetlabs.services.jruby_pool_manager.jruby_schemas JRubyInstance)))
 
 (use-fixtures :once schema-test/validate-schemas)
 
+(deftest execute-tasks!-test
+  (let [pool-context (jruby-pool-manager-core/create-pool-context
+                      (jruby-testutils/jruby-config {:instance-creation-concurrency 3}))
+        creation-service (jruby-internal/get-creation-service pool-context)]
+    (testing "creation-service is a FixedThreadPool of configured number of threads"
+      (is (= 3 (.getMaximumPoolSize creation-service))))
+    ;; this isn't a requirement and should be able to change in the future without issue,
+    ;; but none of the current callers require the result, so explictly test the assumption.
+    (testing "does not return results of task execution"
+      (let [tasks [(fn [] :foo) (fn [] :bar)]
+            results (jruby-agents/execute-tasks! tasks creation-service)]
+        (is (nil? results))))
+    (testing "throws original execptions"
+      (let [tasks [(fn [] (throw (IllegalStateException. "BOOM")))]]
+        (is (thrown? IllegalStateException (jruby-agents/execute-tasks! tasks creation-service)))))))
+
 (deftest next-instance-id-test
   (let [pool-context (jruby-pool-manager-core/create-pool-context
                       (jruby-testutils/jruby-config {:max-active-instances 8}))]

test/unit/puppetlabs/services/jruby_pool_manager/jruby_pool_test.clj

@@ -189,7 +189,7 @@
         config        (jruby-testutils/jruby-config {:max-active-instances pool-size})
         pool-context (jruby-pool-manager-core/create-pool-context config)
         err-msg       (re-pattern "Unable to borrow JRubyInstance from pool")]
-   (is (thrown? IllegalStateException (jruby-agents/prime-pool!
+    (is (thrown? IllegalStateException (jruby-agents/prime-pool!
                                        (assoc-in pool-context [:config :lifecycle :initialize-pool-instance]
                                                  (fn [_] (throw (IllegalStateException. "BORK!")))))))
     (testing "borrow and borrow-with-timeout both throw an exception if the pool failed to initialize"