(PDB-5645) add query monitor

.clj-kondo/config.edn

@@ -16,14 +16,17 @@
    puppetlabs.trapperkeeper.testutils.logging/with-logged-event-maps hooks/one-binding-then-body
    puppetlabs.trapperkeeper.testutils.webserver/with-test-webserver hooks/with-test-webserver}}
  :lint-as
- {puppetlabs.puppetdb.schema/defn-validated schema.core/defn
+ {murphy/try! clojure.core/try
+  puppetlabs.puppetdb.schema/defn-validated schema.core/defn
   puppetlabs.puppetdb.testutils/dotestseq clojure.core/doseq
   puppetlabs.puppetdb.testutils/with-wrapped-fn-args clojure.core/let
   puppetlabs.puppetdb.scf.storage-test/deftest-db clojure.test/deftest}
  :linters
  {:deprecated-var
   {:exclude
-   {puppetlabs.puppetdb.jdbc/call-with-array-converted-query-rows {:namespaces [".*"]}}}
+   {puppetlabs.puppetdb.jdbc/call-with-array-converted-query-rows {:namespaces [".*"]}
+    puppetlabs.puppetdb.testutils.services/call-with-puppetdb-instance {:namespaces [".*"]}
+    puppetlabs.puppetdb.testutils.services/with-puppetdb-instance {:namespaces [".*"]}}}
   :refer-all {:exclude [clojure.test]}
   :unresolved-symbol
   {:exclude

eastwood.clj

@@ -21,6 +21,8 @@
  {:linter :deprecations
   :symbol-matches
   #{#"^#'puppetlabs\.puppetdb\.jdbc/call-with-array-converted-query-rows$"
+    #"^#'puppetlabs\.puppetdb\.testutils\.services/call-with-puppetdb-instance$"
+    #"^#'puppetlabs\.puppetdb\.testutils\.services/with-puppetdb-instance$"
     #"^#'puppetlabs\.trapperkeeper\.testutils\.logging/atom-appender$"
     #"^#'puppetlabs\.trapperkeeper\.testutils\.logging/atom-logger$"
     #"^#'puppetlabs\.trapperkeeper\.testutils\.logging/logs-matching$"}})

project.clj

@@ -210,7 +210,7 @@
 
                  ;; WebAPI support libraries.
                  [bidi]
-                 [clj-http "2.0.1"]
+                 [clj-http "3.12.3"]
                  [commons-io]
                  [compojure]
                  [ring/ring-core]

src/puppetlabs/puppetdb/cli/services.clj

@@ -59,6 +59,7 @@
             [puppetlabs.puppetdb.utils.metrics :as mutils]
             [puppetlabs.puppetdb.nio :refer [get-path]]
             [puppetlabs.puppetdb.query-eng :as qeng]
+            [puppetlabs.puppetdb.query.monitor :as qmon]
             [puppetlabs.puppetdb.query.population :as pop]
             [puppetlabs.puppetdb.scf.migrate
              :refer [desired-schema-version
@@ -100,6 +101,12 @@
    [java.util.concurrent.locks ReentrantLock]
    [org.joda.time Period]))
 
+(def env-monitor-queries?
+  ;; Note: theres'a also a :conf/test switch
+  (->> (or (System/getenv "PDB_PROMPTLY_END_QUERIES") "true")
+       (re-matches #"yes|true|1")
+       seq))
+
 (defn stop-reporters [registries]
   (letfn [(stop [[registry & registries]]
             (when registry
@@ -505,6 +512,7 @@
     (log/debug (trs "Skipping update check on Puppet Enterprise"))))
 
 (def stop-gc-wait-ms (constantly 5000))
+(def stop-query-monitor-wait-ms (constantly 5000))
 
 (defn ready-to-stop? [{:keys [stopping collecting-garbage]}]
   (and stopping (not (seq collecting-garbage))))
@@ -527,11 +535,12 @@
 (defn shut-down-after-scheduler-unresponsive [f]
   (f))
 
-(defn shut-down-service-scheduler-or-die [s request-shutdown]
+(defn shut-down-scheduler-or-die [s request-shutdown name]
   (request-scheduler-shutdown s :interrupt)
   (if (await-scheduler-shutdown s (stop-gc-wait-ms))
-    (log/info (trs "Periodic activities halted"))
-    (let [msg (trs "Unable to shut down scheduled service tasks, requesting server shutdown")]
+    (log/info (trs "Shut down {0} scheduler" name))
+    (let [msg (trs "Unable to shut down {0} scheduler; requesting server shutdown"
+                   name)]
       (log/error msg)
       (shut-down-after-scheduler-unresponsive
        #(request-shutdown {:puppetlabs.trapperkeeper.core/exit
@@ -545,9 +554,12 @@
    (log/info (trs "Shutdown request received; puppetdb exiting."))
    context
    (finally (some-> (:job-pool context)
-                    (shut-down-service-scheduler-or-die request-shutdown)))
+                    (shut-down-scheduler-or-die request-shutdown "job")))
    (finally (some-> (:gc-pool context)
-                    (shut-down-service-scheduler-or-die request-shutdown)))
+                    (shut-down-scheduler-or-die request-shutdown "gc")))
+   (finally (when-let [m (get-in context [:shared-globals :query-monitor])]
+              (when-not (qmon/stop m (stop-query-monitor-wait-ms))
+                (log/error (trs "Unable to stop query monitor")))))
    (finally (close-write-dbs (get-in context [:shared-globals :scf-write-dbs])))
    (finally (some-> (get-in context [:shared-globals :scf-read-db :datasource])
                     .close))
@@ -923,16 +935,16 @@
 (defn start-garbage-collection
   "Starts garbage collection of the databases represented in db-configs"
   [{:keys [clean-lock] :as _context}
-   job-pool db-configs db-pools db-lock-statuses shutdown-for-ex]
+   sched db-configs db-pools db-lock-statuses shutdown-for-ex]
   (let [dbs-with-gc-enabled (filter (fn [[cfg _ _]] (-> cfg :gc-interval to-millis pos?))
                                     (map vector db-configs db-pools db-lock-statuses))]
     (doseq [[cfg db lock-status] dbs-with-gc-enabled]
       (let [request-cfg (db-config->clean-request cfg)]
         ;; Start GC job pool with initial and subsequent delay of :gc-interval
         (doseq [[request interval] request-cfg]
-          (schedule-with-fixed-delay job-pool #(invoke-periodic-gc db cfg request
-                                                                   shutdown-for-ex
-                                                                   clean-lock lock-status)
+          (schedule-with-fixed-delay sched #(invoke-periodic-gc db cfg request
+                                                                shutdown-for-ex
+                                                                clean-lock lock-status)
                                      (to-millis interval) (to-millis interval)))))))
 
 
@@ -954,26 +966,30 @@
   (when-let [v (version/version)]
     (log/info (trs "PuppetDB version {0}" v)))
 
-  (let [{:keys [database developer read-database emit-cmd-events?]} config
+  (let [{:keys [puppetdb database developer read-database emit-cmd-events?]} config
         {:keys [cmd-event-mult cmd-event-ch]} context
         ;; Assume that the exception has already been reported.
         shutdown-for-ex (exceptional-shutdown-requestor request-shutdown [] 2)
         write-dbs-config (conf/write-databases config)
         emit-cmd-events? (or (conf/pe? config) emit-cmd-events?)
         maybe-send-cmd-event! (partial maybe-send-cmd-event! emit-cmd-events? cmd-event-ch)
-        context (assoc context  ;; context may be augmented further below
-                       :shared-globals {:pretty-print (:pretty-print developer)
-                                        :node-purge-ttl (:node-purge-ttl database)
-                                        :add-agent-report-filter (get-in config [:puppetdb :add-agent-report-filter])
-                                        :query-timeout-default (get-in config [:puppetdb :query-timeout-default])
-                                        :query-timeout-max (get-in config [:puppetdb :query-timeout-max])
-                                        :cmd-event-mult cmd-event-mult
-                                        :maybe-send-cmd-event! maybe-send-cmd-event!
-                                        ;; FIXME: remove this if/when
-                                        ;; we add immediate ::tk/exit
-                                        ;; support to trapperkeeper.
-                                        :shutdown-request (:shutdown-request context)
-                                        :log-queries (get-in config [:puppetdb :log-queries])}
+        test-config (puppetdb ::conf/test)
+        monitor-queries? (::qmon/monitor-queries? test-config true)
+        context (assoc context ;; context may be augmented further below
+                       :shared-globals
+                       (cond-> {:pretty-print (:pretty-print developer)
+                                :node-purge-ttl (:node-purge-ttl database)
+                                :add-agent-report-filter (puppetdb :add-agent-report-filter)
+                                :query-timeout-default (puppetdb :query-timeout-default)
+                                :query-timeout-max (puppetdb :query-timeout-max)
+                                :cmd-event-mult cmd-event-mult
+                                :maybe-send-cmd-event! maybe-send-cmd-event!
+                                ;; FIXME: remove this if/when we add
+                                ;; immediate ::tk/exit support to
+                                ;; trapperkeeper.
+                                :shutdown-request (:shutdown-request context)
+                                :log-queries (puppetdb :log-queries)}
+                         test-config (assoc ::conf/test test-config))
                        :clean-lock (ReentrantLock.))]
 
     (when (> (count write-dbs-config) 1)
@@ -1007,13 +1023,22 @@
 
                    _ command-loader :error #(some-> % future-cancel)
                    ;; schema checks, update checks
-                   job-pool (scheduler 4) :error #(shut-down-service-scheduler-or-die
-                                                   % request-shutdown)
+                   job-pool (scheduler 4) :error #(shut-down-scheduler-or-die
+                                                   % request-shutdown "job")
                    ;; schedule gc routines on a fixed delay, but ensure that we limit
                    ;; database contention by "serializing" the jobs on a thread pool
                    ;; with 1 thread
-                   gc-pool (scheduler 1) :error #(shut-down-service-scheduler-or-die
-                                                   % request-shutdown)
+                   gc-pool (scheduler 1) :error #(shut-down-scheduler-or-die
+                                                  % request-shutdown "gc")
+
+                   query-monitor (when (and env-monitor-queries? monitor-queries?)
+                                   (-> (qmon/monitor :on-fatal-error request-shutdown)
+                                       qmon/start))
+                   :error #(when %
+                             (or (qmon/stop % (stop-query-monitor-wait-ms))
+                                 (log/error
+                                  (trs "Unable to stop query monitor after failed puppetdb startup"))))
+
                    ; Create connection pools for each DB we want to write to
                    {:keys [write-db-cfgs write-db-names write-db-pools]}
                    (init-write-dbs write-dbs-config)
@@ -1046,6 +1071,7 @@
                      :dlo dlo
                      :maybe-send-cmd-event! maybe-send-cmd-event!
                      :q q
+                     :query-monitor query-monitor
                      :scf-read-db read-db
                      :scf-write-dbs write-db-pools
                      :scf-write-db-cfgs write-db-cfgs
@@ -1172,7 +1198,7 @@
   that trapperkeeper will call on exit."
   PuppetDBServer
   [[:DefaultedConfig get-config]
-   [:WebroutingService add-ring-handler get-registered-endpoints]
+   [:WebroutingService get-registered-endpoints]
    [:ShutdownService get-shutdown-reason request-shutdown]]
 
   (init

src/puppetlabs/puppetdb/config.clj

@@ -691,11 +691,13 @@
   "Accepts a map containing all of the user-provided configuration values
   and configures the various PuppetDB subsystems."
   [config]
-  (-> config
-      configure-globals
-      configure-developer
-      validate-vardir
-      convert-config))
+  (let [test-config (get-in config [:puppetdb ::test])]
+    (-> config
+        configure-globals
+        configure-developer
+        validate-vardir
+        convert-config
+        (assoc-in [:puppetdb ::test] test-config))))
 
 (defn foss? [config]
   (= "puppetdb" (get-in config [:global :product-name])))

src/puppetlabs/puppetdb/http/query.clj

@@ -6,10 +6,13 @@
   (:require [puppetlabs.puppetdb.cheshire :as json]
             [clojure.java.io]
             [clojure.core.match :as cm]
+            [clojure.set :as set]
             [clojure.tools.logging :as log]
             [clojure.walk :refer [keywordize-keys stringify-keys]]
+            [murphy :refer [try!]]
             [puppetlabs.puppetdb.query-eng :as qeng]
-            [clojure.set :as set]
+            [puppetlabs.puppetdb.query.monitor :as qmon]
+            [puppetlabs.trapperkeeper.services.webserver.jetty9 :as jetty9]
             [puppetlabs.i18n.core :refer [trs tru]]
             [puppetlabs.kitchensink.core :as kitchensink]
             [schema.core :as s]
@@ -24,7 +27,7 @@
                                                       parse-order-by-json]]
             [puppetlabs.puppetdb.pql :as pql]
             [puppetlabs.puppetdb.time :refer [ephemeral-now-ns to-timestamp]]
-            [puppetlabs.puppetdb.utils :refer [update-when]]
+            [puppetlabs.puppetdb.utils :refer [response->channel update-when]]
             [puppetlabs.puppetdb.utils.string-formatter :refer [pprint-json-parse-exception]])
   (:import
    (clojure.lang ExceptionInfo)
@@ -398,28 +401,49 @@
      (handler
       (if puppetdb-query
         req
-        (try
+        (try!
           (let [start-ns (ephemeral-now-ns) ;; capture at the top
                 query-uuid (str (java.util.UUID/randomUUID))
+
+                {:keys [pretty-print query-monitor query-timeout-default
+                        query-timeout-max scf-read-db]}
+                (:globals req)
+
                 req-with-query-uuid (assoc req :puppetdb-query-uuid query-uuid)
                 query-map (create-query-map req-with-query-uuid param-spec parse-fn)
                 ;; Right now, query parsing (above) has no timeouts.
-                max-timeout (get-in req [:globals :query-timeout-max])
-                default-timeout (get-in req [:globals :query-timeout-default])
                 ;; sync is expected to override (based on its own deadlines)
                 sync-timeout (when (#{"puppet:puppetdb-sync-batch"
                                       "puppet:puppetdb-sync-summary"}
                                     (:origin query-map))
-                               (:timeout query-map default-timeout))
+                               (:timeout query-map query-timeout-default))
                 timeout (or sync-timeout
-                            (min (:timeout query-map default-timeout)
-                                 max-timeout))
+                            (min (:timeout query-map query-timeout-default)
+                                 query-timeout-max))
                 deadline (+ start-ns (* timeout 1000000000))
-                pretty-print (:pretty query-map (get-in req [:globals :pretty-print]))]
+                ;; Wait one extra second for time-limited-seq and
+                ;; statement timeouts since the monitor kills the
+                ;; entire pg worker.
+                monitor-deadline (+ deadline 1000000000)
+
+                ;; May have no response because some tests (e.g. some
+                ;; with-http-app based tests) don't add one right now.
+                monitor-id (when-let [chan (and query-monitor
+                                                (some-> (::jetty9/response req)
+                                                        response->channel))]
+                             (qmon/stop-query-at-deadline-or-disconnect query-monitor
+                                                                        query-uuid
+                                                                        chan
+                                                                        monitor-deadline
+                                                                        scf-read-db))]
             (-> req-with-query-uuid
                 (assoc :puppetdb-query query-map)
-                (assoc-in [:globals :pretty-print] pretty-print)
-                (assoc-in [:globals :query-deadline-ns] deadline)))
+                (update :globals merge
+                        {:pretty-print (:pretty query-map pretty-print)
+                         :query-deadline-ns deadline
+                         :query-monitor query-monitor}
+                        (when monitor-id
+                          {:query-monitor-id monitor-id}))))
           (catch ExceptionInfo ex
             (when-not (= :puppetlabs.puppetdb.query/timeout (:kind (ex-data ex)))
               (throw ex))
@@ -476,7 +500,8 @@
   [globals]
   (select-keys globals [:scf-read-db :warn-experimental :url-prefix
                         :pretty-print :node-purge-ttl :add-agent-report-filter
-                        :log-queries :query-deadline-ns]))
+                        :log-queries :query-deadline-ns :query-monitor
+                        :query-monitor-id :puppetlabs.puppetdb.config/test]))
 
 (defn valid-query?
   [scf-read-db version query-map query-options]

src/puppetlabs/puppetdb/jdbc.clj

@@ -716,6 +716,9 @@
 (defn current-database []
   (-> "select current_database();" query-to-vec first :current_database))
 
+(defn current-pid []
+  (-> "select pg_backend_pid();" query-to-vec first :pg_backend_pid))
+
 (defn-validated has-database-privilege? :- s/Bool
   [user db privilege]
   (-> ["select has_database_privilege(?, ?, ?)" user db privilege]

src/puppetlabs/puppetdb/middleware.clj

@@ -24,6 +24,7 @@
             [puppetlabs.puppetdb.constants :as constants]
             [puppetlabs.puppetdb.command.constants :as const])
   (:import
+   (clojure.lang ExceptionInfo)
    (java.net HttpURLConnection)))
 
 (def handler-schema (s/=> s/Any {s/Any s/Any}))
@@ -112,6 +113,14 @@
   (fn [req]
     (try
       (app req)
+      (catch ExceptionInfo e
+        (if (= :puppetlabs.puppetdb.query/terminated (:kind (ex-data e)))
+          (http/error-response (tru "Query backend terminated")
+                               HttpURLConnection/HTTP_INTERNAL_ERROR)
+          (do
+            (log/error e)
+            (http/error-response (cause-finder e)
+              HttpURLConnection/HTTP_INTERNAL_ERROR))))
       (catch Exception e
         (log/error e)
         (http/error-response (cause-finder e)

src/puppetlabs/puppetdb/pdb_routing.clj

@@ -130,7 +130,8 @@
                                    clean
                                    delete-node))
          (mid/wrap-cert-authn cert-allowlist)
-         mid/wrap-with-puppetdb-middleware))
+         mid/wrap-with-puppetdb-middleware)
+     {:include-response true})
 
     (enable-maint-mode)
     (pdb-status/register-pdb-status