Add readiness coordination for systemd notifications

Signed-off-by: Tim Meusel <tim@bastelfreak.de>

Author: bastelfreak

project.clj

@@ -95,7 +95,7 @@
             [org.openvoxproject/i18n ~i18n-version]]
 
   :eastwood {:ignored-faults {:reflection {puppetlabs.trapperkeeper.logging [{:line 92}]
-                                           puppetlabs.trapperkeeper.internal [{:line 177}]
+                                           puppetlabs.trapperkeeper.internal [{:line 194}]
                                            puppetlabs.trapperkeeper.testutils.logging true
                                            puppetlabs.trapperkeeper.testutils.logging-test true
                                            puppetlabs.trapperkeeper.services.nrepl.nrepl-service-test true

src/puppetlabs/trapperkeeper/core.clj

@@ -13,17 +13,19 @@
   (:import
    (clojure.lang ExceptionInfo)))
 
-(def #^{:macro true
-        :doc "An alias for the `puppetlabs.trapperkeeper.services/service` macro
-             so that it is accessible from the core namespace along with the
-             rest of the API."}
-       service #'services/service)
-
-(def #^{:macro true
-        :doc "An alias for the `puppetlabs.trapperkeeper.services/defservice` macro
-             so that it is accessible from the core namespace along with the
-             rest of the API."}
-       defservice #'services/defservice)
+(defmacro service
+  "An alias for the `puppetlabs.trapperkeeper.services/service` macro
+  so that it is accessible from the core namespace along with the
+  rest of the API."
+  [& forms]
+  `(services/service ~@forms))
+
+(defmacro defservice
+  "An alias for the `puppetlabs.trapperkeeper.services/defservice` macro
+  so that it is accessible from the core namespace along with the
+  rest of the API."
+  [svc-name & forms]
+  `(services/defservice ~svc-name ~@forms))
 
 (defn build-app
   "Given a list of services and a map of configuration data, build an instance

src/puppetlabs/trapperkeeper/internal.clj

@@ -10,6 +10,7 @@
 
   (:require [clojure.tools.logging :as log]
             [beckon]
+            [clojure.set :as set]
             [plumbing.graph :as graph]
             [slingshot.slingshot :refer [throw+]]
             [puppetlabs.trapperkeeper.config :refer [config-service get-in-config]]
@@ -77,6 +78,22 @@
 (defn notice-service-reloading [] (maybe-notify-systemd "RELOADING=1\n"))
 (defn notice-service-stopping [] (maybe-notify-systemd "STOPPING=1\n"))
 
+(defn maybe-send-ready-notice!
+  [readiness-state]
+  (let [send-notice? (atom false)]
+    (swap! readiness-state
+           (fn [{:keys [notifications-enabled? registered ready notice-sent?] :as state}]
+             (if (and notifications-enabled?
+                      (not notice-sent?)
+                      (seq registered)
+                      (set/subset? registered ready))
+               (do
+                 (reset! send-notice? true)
+                 (assoc state :notice-sent? true))
+               state)))
+    (when @send-notice?
+      (notice-service-ready))))
+
 ;; This is (eww) a global variable that holds a reference to all of the running
 ;; Trapperkeeper apps in the process. It can be used when connecting via nrepl
 ;; to allow you to do useful things, and also may be used for other things
@@ -469,6 +486,20 @@
     "Higher-order function to execute application logic and trigger shutdown in
     the event of an exception"))
 
+(defprotocol ReadinessService
+  (register-ready! [this service-id]
+    "Register a service that will signal readiness explicitly.")
+  (signal-ready! [this service-id]
+    "Mark a registered service as ready.")
+  (readiness-coordinated? [this]
+    "Returns true if any services have registered explicit readiness coordination.")
+  (enable-ready-notifications! [this]
+    "Allow readiness notifications to be emitted once all registered services are ready.")
+  (reset-readiness! [this]
+    "Reset readiness tracking for a boot or restart cycle.")
+  (readiness-state [this]
+    "Return the current readiness tracking state."))
+
 (schema/defn shutdown-service
   "Provides various functions for triggering application shutdown programatically.
   Primarily intended to serve application services, though TrapperKeeper also uses
@@ -497,6 +528,60 @@
     (shutdown-on-error [this svc-id f] (shutdown-on-error* shutdown-reason-promise app-context svc-id f))
     (shutdown-on-error [this svc-id f on-error] (shutdown-on-error* shutdown-reason-promise app-context svc-id f on-error))))
 
+(schema/defn readiness-service :- (schema/protocol s/ServiceDefinition)
+  "Provides explicit readiness coordination for services that need to delay
+  systemd READY=1 until after their own startup work is complete. Services that
+  participate should call `register-ready!` before app startup completes and
+  `signal-ready!` when they are actually ready to serve traffic."
+  []
+  (let [state (atom {:notifications-enabled? false
+                     :registered #{}
+                     :ready #{}
+                     :notice-sent? false})]
+    (s/service ReadinessService
+      []
+      (register-ready! [_ service-id]
+        (schema/validate schema/Keyword service-id)
+        (swap! state update :registered conj service-id)
+        (maybe-send-ready-notice! state)
+        nil)
+      (signal-ready! [_ service-id]
+        (schema/validate schema/Keyword service-id)
+        (swap! state update :ready conj service-id)
+        (maybe-send-ready-notice! state)
+        nil)
+      (readiness-coordinated? [_]
+        (boolean (seq (:registered @state))))
+      (enable-ready-notifications! [_]
+        (swap! state assoc :notifications-enabled? true)
+        (maybe-send-ready-notice! state)
+        nil)
+      (reset-readiness! [_]
+        (reset! state {:notifications-enabled? false
+                       :registered #{}
+                       :ready #{}
+                       :notice-sent? false})
+        nil)
+      (readiness-state [_]
+        @state))))
+
+(defn reset-readiness-tracking!
+  [services-by-id]
+  (when-let [readiness-service (services-by-id :ReadinessService)]
+    (reset-readiness! readiness-service)))
+
+(defn enable-ready-notifications-for-app!
+  [services-by-id]
+  (when-let [readiness-service (services-by-id :ReadinessService)]
+    (enable-ready-notifications! readiness-service)))
+
+(defn notice-service-ready-if-uncoordinated!
+  [services-by-id]
+  (if-let [readiness-service (services-by-id :ReadinessService)]
+    (when-not (readiness-coordinated? readiness-service)
+      (notice-service-ready))
+    (notice-service-ready)))
+
 (schema/defn ^:always-validate shutdown! :- [Throwable]
   "Perform shutdown calling the `stop` lifecycle function on each service,
    in reverse order (to account for dependency relationships).
@@ -633,6 +718,7 @@
         service-refs (atom {})
         services (conj services
                        (config-service config-data-fn)
+                       (readiness-service)
                        (initialize-shutdown-service! app-context
                                                      shutdown-reason-promise))
         service-map (apply merge (map s/service-map services))
@@ -660,11 +746,13 @@
       (a/check-for-errors! [this] (throw-app-error-if-exists!
                                    this))
       (a/init [this]
+        (reset-readiness-tracking! services-by-id)
         (run-lifecycle-fns app-context s/init "init" ordered-services)
         this)
       (a/start [this]
         (run-lifecycle-fns app-context s/start "start" ordered-services)
-        (notice-service-ready)
+        (enable-ready-notifications-for-app! services-by-id)
+        (notice-service-ready-if-uncoordinated! services-by-id)
         (inc-restart-counter! this)
         this)
       (a/stop [this]
@@ -682,9 +770,11 @@
           (notice-service-reloading)
           (run-lifecycle-fns app-context s/stop "stop" (reverse ordered-services))
           (doseq [svc-id (keys services-by-id)] (swap! app-context assoc-in [:service-contexts svc-id] {}))
+          (reset-readiness-tracking! services-by-id)
           (run-lifecycle-fns app-context s/init "init" ordered-services)
           (run-lifecycle-fns app-context s/start "start" ordered-services)
-          (notice-service-ready)
+          (enable-ready-notifications-for-app! services-by-id)
+          (notice-service-ready-if-uncoordinated! services-by-id)
           (inc-restart-counter! this)
           this
           (catch Throwable t

src/puppetlabs/trapperkeeper/services.clj

@@ -103,7 +103,7 @@
                          (get-services [this#]
                            (-> ~'@tk-app-context
                                :services-by-id
-                               (dissoc :ConfigService :ShutdownService)
+                               (dissoc :ConfigService :ShutdownService :ReadinessService)
                                vals))
                          (service-symbol [this#] '~service-sym)
                          (service-included? [this# service-id#]

test/puppetlabs/trapperkeeper/services_test.clj

@@ -55,6 +55,10 @@
 (defprotocol Service3
   (service3-fn [this]))
 
+(defprotocol CoordinatedReadyService
+  (mark-ready [this])
+  (current-readiness-state [this]))
+
 (deftest test-services-not-required
   (testing "services are not required to define lifecycle functions"
     (let [service1 (service Service1
@@ -520,6 +524,66 @@
               (is (= #{:EmptyService :HelloService}
                      (set (map svcs/service-id all-services)))))))))))
 
+(deftest readiness-service-falls-back-to-central-ready-notice
+  (let [ready-notices (atom 0)
+        service1 (service Service1
+                          []
+                          (service1-fn [_] "hi"))]
+    (with-redefs [internal/notice-service-ready #(swap! ready-notices inc)]
+      (with-app-with-empty-config app [service1]
+        (is (= 1 @ready-notices))))))
+
+(deftest readiness-service-defers-ready-notice-until-signaled
+  (let [ready-notices (atom 0)
+        coordinated-service
+        (service CoordinatedReadyService
+                 [[:ReadinessService register-ready! signal-ready! readiness-state]]
+                 (init [_ context]
+                   (register-ready! :CoordinatedReadyService)
+                   context)
+                 (start [_ context]
+                   context)
+                 (mark-ready [_]
+                   (signal-ready! :CoordinatedReadyService))
+                 (current-readiness-state [_]
+                   (readiness-state)))]
+    (with-redefs [internal/notice-service-ready #(swap! ready-notices inc)]
+      (with-app-with-empty-config app [coordinated-service]
+        (let [service (app/get-service app :CoordinatedReadyService)]
+          (is (zero? @ready-notices))
+          (is (= #{:CoordinatedReadyService}
+                 (get-in (current-readiness-state service) [:registered])))
+          (mark-ready service)
+          (is (= 1 @ready-notices))
+          (mark-ready service)
+          (is (= 1 @ready-notices)))))))
+
+(deftest readiness-service-resets-across-restart
+  (let [ready-notices (atom 0)
+        coordinated-service
+        (service CoordinatedReadyService
+                 [[:ReadinessService register-ready! signal-ready! readiness-state]]
+                 (init [_ context]
+                   (register-ready! :CoordinatedReadyService)
+                   context)
+                 (start [_ context]
+                   context)
+                 (mark-ready [_]
+                   (signal-ready! :CoordinatedReadyService))
+                 (current-readiness-state [_]
+                   (readiness-state)))]
+    (with-redefs [internal/notice-service-ready #(swap! ready-notices inc)]
+      (with-app-with-empty-config app [coordinated-service]
+        (let [service (app/get-service app :CoordinatedReadyService)]
+          (mark-ready service)
+          (is (= 1 @ready-notices))
+          (app/restart app)
+          (is (= 1 @ready-notices))
+          (is (= #{:CoordinatedReadyService}
+                 (get-in (current-readiness-state service) [:registered])))
+          (mark-ready service)
+          (is (= 2 @ready-notices)))))))
+
 (deftest minimal-services-test
   (testing "minimal services can be defined without a protocol"
     (let [call-seq (atom [])