Make Shutdown Idempotent

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

@@ -233,10 +233,19 @@
   (let [pool-state (jruby-internal/get-pool-state pool-context)
         pool (:pool pool-state)
         on-complete (promise)]
-    (drain-and-refill-pool! pool-context false on-complete)
-    (jruby-internal/insert-shutdown-poison-pill pool)
-    ; Wait for flush to complete
-    @on-complete
+    (try
+      (drain-and-refill-pool! pool-context false on-complete)
+      (jruby-internal/insert-shutdown-poison-pill pool)
+      ; Wait for flush to complete
+      @on-complete
+      (catch InterruptedException _e
+        ;; A poison pill is already present: either this pool was already shut
+        ;; down, or an instance failed to initialize and poisoned the pool
+        ;; while we were draining. Either way the pool is already unusable, so
+        ;; ensure a shutdown pill is in place and treat this as a completed
+        ;; (no-op) shutdown rather than failing the stop sequence.
+        (jruby-internal/insert-shutdown-poison-pill pool)
+        (log/debug (i18n/trs "JRuby pool already shut down or poisoned; skipping flush"))))
     (log/debug (i18n/trs "Finished flush of JRuby pools for shutdown"))))
 
 (schema/defn ^:always-validate

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

@@ -78,18 +78,22 @@
       ;; Lock the pool so no borrows or flushes can occur while we're shutting down
       (try
         (pool-protocol/lock-with-timeout pool-context flush-timeout TimeUnit/MILLISECONDS)
+        (try
+          (let [instance (.borrowItem pool)
+                _ (.releaseItem pool instance)]
+            ;; This will block until all borrows have been returned
+            (jruby-internal/cleanup-pool-instance! instance cleanup-fn))
+          ;; Insert a shutdown pill to ensure that all pending borrows and locks
+          ;; are rejected with the appropriate logging
+          (jruby-internal/insert-shutdown-poison-pill pool)
+          (finally
+            (pool-protocol/unlock pool-context)))
         (catch TimeoutException e
-          (jruby-internal/throw-jruby-lock-timeout e)))
-      (try
-        (let [instance (.borrowItem pool)
-              _ (.releaseItem pool instance)]
-          ;; This will block until all borrows have been returned
-          (jruby-internal/cleanup-pool-instance! instance cleanup-fn))
-        ;; Insert a shutdown pill to ensure that all pending borrows and locks
-        ;; are rejected with the appropriate logging
-        (jruby-internal/insert-shutdown-poison-pill pool)
-        (finally
-          (pool-protocol/unlock pool-context)))))
+          (jruby-internal/throw-jruby-lock-timeout e))
+        (catch InterruptedException _e
+          ;; A pill is already present; the pool is already shut down.
+          ;; Treat this as a completed no-op.
+          nil))))
 
   (lock
     [pool-context]

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

@@ -364,3 +364,69 @@
            pill (ShutdownPoisonPill. pool)]
        ; Returning a pill should be a noop
        (jruby-core/return-to-pool pool-context pill :test [])))))
+
+(deftest flush-and-repopulate-interrupted-by-shutdown-test
+  (testing "flush-and-repopulate-pool! raises InterruptedException when racing with flush-pool-for-shutdown!"
+    (let [config (jruby-testutils/jruby-config {:max-active-instances 2})
+          ;; Set up pool manually to avoid double-shutdown from with-pool-context
+          pool-context (jruby-pool-manager-core/create-pool-context config)
+          _ (jruby-agents/prime-pool! pool-context)
+          _ (jruby-testutils/wait-for-jrubies-from-pool-context pool-context)
+          pool (jruby-core/get-pool pool-context)
+          ;; Borrow all instances to block both shutdown and flush from acquiring the lock
+          instances (jruby-testutils/drain-pool pool-context 2)
+          ;; Start shutdown: acquires pool lock immediately, then waits for both
+          ;; instances to be returned before proceeding
+          shutdown-complete? (promise)
+          _ (future
+              (jruby-core/flush-pool-for-shutdown! pool-context)
+              (deliver shutdown-complete? true))
+          ;; Wait for shutdown to hold the pool lock
+          _ (jruby-testutils/wait-for-pool-to-be-locked pool)
+          ;; Start flush: will block in lockWithTimeout waiting for the lock
+          flush-thread (future
+                         (try
+                           (pool-protocol/flush-pool pool-context)
+                           :ok
+                           (catch InterruptedException e
+                             e)))]
+      ;; Both should be blocked at this point
+      (is (not (realized? shutdown-complete?)))
+      (is (not (realized? flush-thread)))
+      ;; Returning instances unblocks shutdown, which drains the pool, releases the
+      ;; lock, then inserts the pill. The flush racing for the lock receives
+      ;; InterruptedException because the pill is inserted before or during its attempt.
+      (jruby-core/return-to-pool pool-context (first instances) :test [])
+      (jruby-core/return-to-pool pool-context (second instances) :test [])
+      @shutdown-complete?
+      (let [result @flush-thread]
+        (is (instance? InterruptedException result))
+        (is (= "Lock can't be granted because a pill has been inserted"
+               (.getMessage result)))))))
+
+(deftest flush-pool-for-shutdown-tolerates-existing-pill-test
+  (testing "flush-pool-for-shutdown! completes without throwing when the pool has already been poisoned"
+    (let [config (jruby-testutils/jruby-config {:max-active-instances 2})
+          ;; Set up pool manually since we poison the pool and shut it down here
+          pool-context (jruby-pool-manager-core/create-pool-context config)
+          _ (jruby-agents/prime-pool! pool-context)
+          _ (jruby-testutils/wait-for-jrubies-from-pool-context pool-context)
+          pool (jruby-core/get-pool pool-context)]
+      ;; Simulate an instance failing to initialize, which clears the pool and
+      ;; inserts a poison pill (see add-instance's error handler). A subsequent
+      ;; shutdown must not fail the stop sequence with an uncaught
+      ;; InterruptedException from the pill.
+      (jruby-internal/insert-poison-pill pool (IllegalStateException. "simulated init failure"))
+      (is (nil? (jruby-core/flush-pool-for-shutdown! pool-context))))))
+
+(deftest flush-pool-for-shutdown-is-idempotent-test
+  (testing "calling flush-pool-for-shutdown! a second time is a safe no-op"
+    (let [config (jruby-testutils/jruby-config {:max-active-instances 2})
+          ;; Set up pool manually so we control the (repeated) shutdown
+          pool-context (jruby-pool-manager-core/create-pool-context config)
+          _ (jruby-agents/prime-pool! pool-context)
+          _ (jruby-testutils/wait-for-jrubies-from-pool-context pool-context)]
+      ;; First shutdown drains the pool and inserts a shutdown pill
+      (is (nil? (jruby-core/flush-pool-for-shutdown! pool-context)))
+      ;; Second shutdown finds the pill already present and must not throw
+      (is (nil? (jruby-core/flush-pool-for-shutdown! pool-context))))))

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

@@ -238,3 +238,58 @@
       (is (thrown+? [:kind :puppetlabs.services.jruby-pool-manager.impl.jruby-internal/jruby-lock-timeout
                      :msg "An attempt to lock the JRubyPool failed with a timeout"]
                     (pool-protocol/shutdown pool-context))))))
+
+(deftest shutdown-is-idempotent-test
+  (testing "calling shutdown a second time is a safe no-op"
+    (let [config (jruby-test-config 0 2)
+          ;; Set up the pool manually, since we trigger shutdown directly here
+          ;; and the macro would also try to shutdown.
+          pool-context (jruby-pool-manager-core/create-pool-context config)
+          _ (jruby-agents/prime-pool! pool-context)
+          _ (jruby-testutils/wait-for-jrubies-from-pool-context pool-context)]
+      ;; First shutdown drains the instance and inserts a shutdown pill
+      (is (nil? (pool-protocol/shutdown pool-context)))
+      ;; Second shutdown finds the pill already present; lock-with-timeout
+      ;; throws InterruptedException, which must be absorbed rather than
+      ;; propagated out of the stop sequence.
+      (is (nil? (pool-protocol/shutdown pool-context))))))
+
+(deftest flush-if-at-max-borrows-interrupted-by-shutdown-test
+  (testing "InterruptedException from flush-if-at-max-borrows during shutdown propagates uncaught through the agent"
+    ;; A custom shutdown-on-error captures the exception that escapes from
+    ;; flush-if-at-max-borrows when it races with a concurrent shutdown.
+    (let [captured-exception (promise)
+          config (jruby-test-config
+                   1 1
+                   {:lifecycle {:shutdown-on-error (fn [f]
+                                                     (try (f)
+                                                          (catch InterruptedException e
+                                                            (deliver captured-exception e))))}})
+          ;; Set up pool manually to avoid double-shutdown from with-pool-context
+          pool-context (jruby-pool-manager-core/create-pool-context config)
+          _ (jruby-agents/prime-pool! pool-context)
+          _ (jruby-testutils/wait-for-jrubies-from-pool-context pool-context)
+          pool (jruby-core/get-pool pool-context)
+          ;; Borrow the instance to prevent shutdown from proceeding immediately
+          instance (first (pool-protocol/borrow pool-context))
+          ;; Start shutdown: acquires pool lock, then waits for borrow count to drop to zero
+          shutdown-complete? (promise)
+          _ (future
+              (pool-protocol/shutdown pool-context)
+              (deliver shutdown-complete? true))]
+      ;; Wait for shutdown to hold the pool lock
+      (is (jruby-testutils/wait-for-pool-to-be-locked pool))
+      ;; Returning the instance simultaneously unblocks shutdown (currentBorrowCount
+      ;; drops to 0) and triggers flush-if-at-max-borrows to be sent to the agent
+      ;; (borrow-count atom crosses max-borrows-per-instance). Shutdown inserts its
+      ;; pill while flush-if-at-max-borrows is waiting to acquire the lock, causing
+      ;; an InterruptedException to propagate through the agent's shutdown-on-error.
+      (pool-protocol/return pool-context instance)
+      @shutdown-complete?
+      (let [err (deref captured-exception 5000 nil)]
+        (is (some? err) "expected InterruptedException from flush-if-at-max-borrows")
+        (when err
+          (is (instance? InterruptedException err))
+          (is (= "Lock can't be granted because a pill has been inserted"
+                 (.getMessage err))))))))
+