(SERVER-150) Add service function for flushing JRuby pool

This commit adds a new service function to the JRubyPuppet service.
The new function can be used to request a flush of the JRuby pool.

The function is implemented by creating a new (offline) JRuby pool
and starting to populate it; as soon as the first instance is ready,
we can swap it into the service and remove the old pool, and then
continue initializing the remaining instances and cleaning up the
old pool in parallel.  It uses clojure agents to coordinate the
creation of the new pool and the cleanup of the old pool.

dev/user_repl.clj

@@ -127,3 +127,9 @@
   []
   (jruby-protocol/mark-all-environments-expired!
     (tka/get-service system :JRubyPuppetService)))
+
+(defn flush-jruby-pool!
+  "Flush and repopulate the JRuby pool"
+  []
+  (jruby-protocol/flush-jruby-pool!
+    (tka/get-service system :JRubyPuppetService)))

src/clj/puppetlabs/services/jruby/jruby_puppet_agents.clj

@@ -1,8 +1,10 @@
 (ns puppetlabs.services.jruby.jruby-puppet-agents
   (:import (clojure.lang IFn Agent)
-           (com.puppetlabs.puppetserver PuppetProfiler))
+           (com.puppetlabs.puppetserver PuppetProfiler)
+           (puppetlabs.services.jruby.jruby_puppet_core PoisonPill))
   (:require [schema.core :as schema]
-            [puppetlabs.services.jruby.jruby-puppet-core :as jruby-core]))
+            [puppetlabs.services.jruby.jruby-puppet-core :as jruby-core]
+            [clojure.tools.logging :as log]))
 
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
@@ -33,11 +35,46 @@
                     agent-ctxt)]
     (send jruby-agent agent-fn)))
 
+(schema/defn ^:always-validate
+  flush-pool!
+  "Flush of the current JRuby pool.  NOTE: this function should never
+  be called except by the pool-agent."
+  [pool-context :- jruby-core/PoolContext]
+  ;; Since this function is only called by the pool-agent, and since this
+  ;; is the only entry point into the pool flushing code that is exposed by the
+  ;; service API, we know that if we receive multiple flush requests before the
+  ;; first one finishes, they will be queued up and the body of this function
+  ;; will be executed atomically; we don't need to worry about race conditions
+  ;; between the steps we perform here in the body.
+  (log/info "Flush request received; creating new JRuby pool.")
+  (let [{:keys [config profiler pool-state]} pool-context
+        new-pool-state (jruby-core/create-pool-from-config config)
+        new-pool (:pool new-pool-state)
+        old-pool @pool-state
+        count    (:size old-pool)]
+    (log/info "Replacing old JRuby pool with new instance.")
+    (reset! pool-state new-pool-state)
+    (log/info "Swapped JRuby pools, beginning cleanup of old pool.")
+    (doseq [i (range count)]
+      (let [id (inc i)
+            instance (jruby-core/borrow-from-pool (:pool old-pool))]
+        (.terminate (:scripting-container instance))
+        (log/infof "Cleaned up old JRuby instance %s of %s, creating replacement."
+                   id count)
+        (try
+          (jruby-core/create-pool-instance! #spy/d new-pool id config profiler)
+          (log/infof "Finished creating JRubyPuppet instance %d of %d"
+                     id count)
+          (catch Exception e
+            (.clear new-pool)
+            (.put new-pool (PoisonPill. e))
+            (throw (IllegalStateException. "There was a problem adding a JRubyPuppet instance to the pool." e))))))
+
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;;; Public
 
 (schema/defn ^:always-validate
-  jruby-pool-agent :- JRubyPoolAgent
+  pool-agent :- JRubyPoolAgent
   "Given a shutdown-on-error function, create an agent suitable for use in managing
   JRuby pools."
   [shutdown-on-error-fn :- IFn]
@@ -46,8 +83,14 @@
 (schema/defn ^:always-validate
   send-prime-pool! :- JRubyPoolAgent
   "Sends a request to the agent to prime the pool using the given pool context."
-  [prime-pool-agent :- JRubyPoolAgent
-   pool-state :- jruby-core/PoolStateContainer
-   config :- jruby-core/JRubyPuppetConfig
-   profiler :- (schema/maybe PuppetProfiler)]
-  (send-agent prime-pool-agent #(jruby-core/prime-pool! pool-state config profiler)))
+  [pool-context :- jruby-core/PoolContext
+   pool-agent :- JRubyPoolAgent]
+  (let [{:keys [pool-state config profiler]} pool-context]
+    (send-agent pool-agent #(jruby-core/prime-pool! pool-state config profiler))))
+
+(schema/defn ^:always-validate
+  send-flush-pool! :- JRubyPoolAgent
+  "Sends requests to the agent to flush the existing pool and create a new one."
+  [pool-context :- jruby-core/PoolContext
+   pool-agent :- JRubyPoolAgent]
+  (send-agent pool-agent #(flush-pool! pool-context)))

src/clj/puppetlabs/services/jruby/jruby_puppet_core.clj

@@ -146,8 +146,8 @@
     (.runScriptlet "require 'puppet/server/master'")))
 
 (schema/defn ^:always-validate
-  create-pool-instance :- JRubyPuppetInstance
-  "Creates a new pool instance."
+  create-pool-instance! :- JRubyPuppetInstance
+  "Creates a new JRubyPuppet instance and adds it to the pool."
   [pool     :- pool-queue-type
    id       :- schema/Int
    config   :- JRubyPuppetConfig
@@ -174,17 +174,19 @@
       (.put puppet-server-config "profiler" profiler)
       (.put puppet-server-config "environment_registry" env-registry)
 
-      (map->JRubyPuppetInstance
-        {:pool                 pool
-         :id                   id
-         :jruby-puppet         (.callMethod scripting-container
-                                            ruby-puppet-class
-                                            "new"
-                                            (into-array Object
-                                                        [puppet-config puppet-server-config])
-                                            JRubyPuppet)
-         :scripting-container  scripting-container
-         :environment-registry env-registry}))))
+      (let [instance (map->JRubyPuppetInstance
+                       {:pool                 pool
+                        :id                   id
+                        :jruby-puppet         (.callMethod scripting-container
+                                                           ruby-puppet-class
+                                                           "new"
+                                                           (into-array Object
+                                                                       [puppet-config puppet-server-config])
+                                                           JRubyPuppet)
+                        :scripting-container  scripting-container
+                        :environment-registry env-registry})]
+        (.put pool instance)
+        instance))))
 
 (schema/defn ^:always-validate
   get-pool-state :- PoolState
@@ -220,14 +222,13 @@
                                            "you did not specify the --config option?")))))
 
 (schema/defn ^:always-validate
-  create-pool-from-config :- PoolStateContainer
+  create-pool-from-config :- PoolState
   "Create a new PoolData based on the config input."
   [{size :max-active-instances} :- JRubyPuppetConfig]
   (let [size (or size default-pool-size)]
-    (atom
-      {:pool         (instantiate-free-pool size)
-       :size         size
-       :initialized? false})))
+    {:pool         (instantiate-free-pool size)
+     :size         size
+     :initialized? false}))
 
 (defn validate-instance-from-pool!
   "Validate an instance.  The main purpose of this function is to check for
@@ -264,7 +265,7 @@
         (dotimes [i count]
           (let [id (inc i)]
             (log/debugf "Priming JRubyPuppet instance %d of %d" id count)
-            (.put pool (create-pool-instance pool id config profiler))
+            (create-pool-instance! pool id config profiler)
             (log/infof "Finished creating JRubyPuppet instance %d of %d"
                        id count))
           (mark-as-initialized! pool-state)))
@@ -283,7 +284,7 @@
   [config profiler]
   {:config     config
    :profiler   profiler
-   :pool-state (create-pool-from-config config)})
+   :pool-state (atom (create-pool-from-config config))})
 
 (schema/defn ^:always-validate
   free-instance-count
@@ -327,3 +328,19 @@
   "Return a borrowed pool instance to its free pool."
   [instance :- JRubyPuppetInstance]
   (.put (:pool instance) instance))
+
+(defmacro with-jruby-puppet
+  "Encapsulates the behavior of borrowing and returning an instance of
+  JRubyPuppet.  Example usage:
+
+  (let [pool (get-pool pool-context)]
+    (with-jruby-puppet
+      jruby-puppet
+      pool
+      (do-something-with-a-jruby-puppet-instance jruby-puppet)))"
+  [jruby-puppet pool & body]
+  `(let [~jruby-puppet (borrow-from-pool ~pool)]
+     (try
+       ~@body
+       (finally
+         (return-to-pool ~jruby-puppet)))))

src/clj/puppetlabs/services/jruby/jruby_puppet_service.clj

@@ -27,15 +27,15 @@
                                      (get-in-config [:http-client :cipher-suites])))
           service-id        (tk-services/service-id this)
           agent-shutdown-fn (partial shutdown-on-error service-id)
-          prime-pool-agent  (jruby-agents/jruby-pool-agent agent-shutdown-fn)
+          pool-agent  (jruby-agents/pool-agent agent-shutdown-fn)
           profiler          (get-profiler)]
       (core/verify-config-found! config)
       (log/info "Initializing the JRuby service")
       (let [pool-context (core/create-pool-context config profiler)]
-        (jruby-agents/send-prime-pool! prime-pool-agent (:pool-state pool-context) config profiler)
+        (jruby-agents/send-prime-pool! pool-context pool-agent)
         (-> context
             (assoc :pool-context pool-context)
-            (assoc :prime-pool-agent prime-pool-agent)))))
+            (assoc :pool-agent pool-agent)))))
 
   (borrow-instance
     [this]
@@ -62,7 +62,13 @@
   (mark-all-environments-expired!
     [this]
     (let [pool-context (:pool-context (tk-services/service-context this))]
-      (core/mark-all-environments-expired! pool-context))))
+      (core/mark-all-environments-expired! pool-context)))
+
+  (flush-jruby-pool!
+    [this]
+    (let [service-context (tk-services/service-context this)
+          {:keys [pool-context pool-agent]} service-context]
+      (jruby-agents/send-flush-pool! pool-context pool-agent))))
 
 (defmacro with-jruby-puppet
   "Encapsulates the behavior of borrowing and returning an instance of

src/clj/puppetlabs/services/protocols/jruby_puppet.clj

@@ -23,4 +23,10 @@
 
   (mark-all-environments-expired!
     [this]
-    "Mark all cached environments expired, in all JRuby instances."))
+    "Mark all cached environments expired, in all JRuby instances.")
+
+  (flush-jruby-pool!
+    [this]
+    "Flush all the current JRuby instances and repopulate the pool."))
+
+

test/unit/puppetlabs/services/config/puppet_server_config_service_test.clj

@@ -85,7 +85,7 @@
              "Providing config values that should be read from Puppet results "
              "in an error that mentions all offending config keys.")
     (with-redefs
-      [jruby-puppet-core/create-pool-instance
+      [jruby-puppet-core/create-pool-instance!
          jruby-testutils/create-mock-pool-instance]
       (with-test-logging
         (is (thrown-with-msg?

test/unit/puppetlabs/services/jruby/jruby_interpreter_test.clj

@@ -18,7 +18,7 @@
                          :gem-home        testutils/gem-home
                          :master-conf-dir testutils/conf-dir}
           pool          (instantiate-free-pool 1)
-          pool-instance (create-pool-instance pool 1 config testutils/default-profiler)
+          pool-instance (create-pool-instance! pool 1 config testutils/default-profiler)
           jruby-puppet  (:jruby-puppet pool-instance)
           var-dir       (.getSetting jruby-puppet "vardir")]
       (is (not (nil? var-dir)))))

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

@@ -77,7 +77,7 @@
         pool-context  (create-pool-context config profiler)
         pool          (get-pool pool-context)
         err-msg       (re-pattern "Unable to borrow JRuby instance from pool")]
-    (with-redefs [core/create-pool-instance (fn [_] (throw (IllegalStateException. "BORK!")))]
+    (with-redefs [core/create-pool-instance! (fn [_] (throw (IllegalStateException. "BORK!")))]
                  (is (thrown? IllegalStateException (prime-pool! (:pool-state pool-context) config profiler))))
     (testing "borrow and borrow-with-timeout both throw an exception if the pool failed to initialize"
       (is (thrown-with-msg? IllegalStateException

test/unit/puppetlabs/services/jruby/jruby_puppet_service_test.clj

@@ -23,7 +23,7 @@
       (str "If there as an exception while putting a JRubyPuppet instance in "
            "the pool the application should shut down.")
     (logging/with-test-logging
-      (with-redefs [jruby-puppet-core/create-pool-instance
+      (with-redefs [jruby-puppet-core/create-pool-instance!
                     (fn [& _] (throw (Exception. "42")))]
                    (let [got-expected-exception (atom false)]
                      (try

test/unit/puppetlabs/services/jruby/jruby_testutils.clj

@@ -2,6 +2,7 @@
   (:import (com.puppetlabs.puppetserver JRubyPuppet JRubyPuppetResponse)
            (org.jruby.embed ScriptingContainer))
   (:require [puppetlabs.services.jruby.jruby-puppet-core :as jruby-core]
+            [puppetlabs.services.puppet-profiler.puppet-profiler-core :as profiler-core]
             [me.raynes.fs :as fs]
             [puppetlabs.services.jruby.puppet-environments :as puppet-env]))
 
@@ -68,7 +69,7 @@
    (create-pool-instance (jruby-puppet-config 1)))
   ([config]
    (let [pool (jruby-core/instantiate-free-pool 1)]
-     (jruby-core/create-pool-instance pool 1 config default-profiler))))
+     (jruby-core/create-pool-instance! pool 1 config default-profiler))))
 
 (defn create-mock-jruby-instance
   "Creates a mock implementation of the JRubyPuppet interface."
@@ -81,18 +82,20 @@
 
 (defn create-mock-pool-instance
   [pool _ _ _]
-  (jruby-core/map->JRubyPuppetInstance
-    {:pool                 pool
-     :id                   1
-     :jruby-puppet         (create-mock-jruby-instance)
-     :scripting-container  (ScriptingContainer.)
-     :environment-registry (puppet-env/environment-registry)}))
+  (let [instance (jruby-core/map->JRubyPuppetInstance
+                   {:pool                 pool
+                    :id                   1
+                    :jruby-puppet         (create-mock-jruby-instance)
+                    :scripting-container  (ScriptingContainer.)
+                    :environment-registry (puppet-env/environment-registry)})]
+    (.put pool instance)
+    instance))
 
 (defn mock-pool-instance-fixture
   "Test fixture which changes the behavior of the JRubyPool to create
   mock JRubyPuppet instances."
   [f]
   (with-redefs
-    [jruby-core/create-pool-instance create-mock-pool-instance]
+    [jruby-core/create-pool-instance! create-mock-pool-instance]
     (f)))