persist: proactively recycle connections on draining CRDB nodes [POC]

Cargo.toml

@@ -340,8 +340,8 @@ ctor = "1.0.10"
 custom-labels = "0.4.6"
 darling = "0.23.0"
 datadriven = "0.9.0"
-deadpool = { version = "0.9.5", default-features = false, features = ["managed"] }
-deadpool-postgres = "0.10.3"
+deadpool = { version = "0.12.3", default-features = false, features = ["managed"] }
+deadpool-postgres = "0.14.1"
 dec = "0.4.9"
 derivative = "2.2.0"
 differential-dataflow = "0.25.0"

deny.toml

@@ -141,8 +141,8 @@ skip = [
     { name = "quick-xml", version = "0.37.5" },
     # aws-lc-rs (via jsonwebtoken 10) and ring pull different `untrusted`.
     { name = "untrusted", version = "0.7.1" },
-    # Held back by lazy_static 1.4.0 (used by num-bigint-dig).
-    { name = "spin", version = "0.5.2" },
+    # Held back by lazy_static 1.5.0 (used by num-bigint-dig, deadpool).
+    { name = "spin", version = "0.9.9" },
     # held back by tower-lsp 0.20 (mz-deploy LSP server)
     { name = "dashmap", version = "5.5.3" },
     { name = "tower", version = "0.4.13" },
@@ -242,6 +242,7 @@ wrappers = [
 [[bans.deny]]
 name = "lazy_static"
 wrappers = [
+    "deadpool",
     "dynfmt",
     "findshlibs",
     "launchdarkly-server-sdk",

misc/python/materialize/mzcompose/__init__.py

@@ -160,6 +160,17 @@ def get_variable_system_parameters(
         ["true", "false"] if read_committed_safe else ["false"],
     )
 
+    # Drain-aware recycling stamps connections and polls gossip via
+    # crdb_internal, which only exists on CockroachDB. Default it on there for
+    # coverage; on other metadata stores the queries would fail (tolerated,
+    # but noisy), so keep it off.
+    drain_aware_safe = metadata_store == "cockroach"
+    persist_drain_aware_recycling = VariableSystemParameter(
+        "persist_consensus_connection_pool_drain_aware_recycling",
+        "true" if drain_aware_safe else "false",
+        ["true", "false"] if drain_aware_safe else ["false"],
+    )
+
     return [
         # -----
         # To reduce CRDB load as we are struggling with it in CI (values based on load test environment):
@@ -455,6 +466,7 @@ def get_variable_system_parameters(
             "persist_blob_cache_scale_with_threads", "true", ["true", "false"]
         ),
         persist_pg_consensus_read_committed,
+        persist_drain_aware_recycling,
         VariableSystemParameter(
             "persist_state_update_lease_timeout", "1s", ["0s", "1s", "10s"]
         ),
@@ -623,6 +635,7 @@ def get_default_system_parameters(
     "persist_consensus_connection_pool_max_wait",
     "persist_consensus_connection_pool_ttl",
     "persist_consensus_connection_pool_ttl_stagger",
+    "persist_consensus_connection_pool_drain_culls_per_tick",
     "persist_use_postgres_tuned_queries",
     "crdb_connect_timeout",
     "crdb_tcp_user_timeout",

misc/python/materialize/parallel_workload/action.py

@@ -1840,6 +1840,9 @@ def __init__(
         self.flags_with_values["persist_claim_unclaimed_compactions"] = (
             BOOLEAN_FLAG_VALUES
         )
+        self.flags_with_values[
+            "persist_consensus_connection_pool_drain_aware_recycling"
+        ] = BOOLEAN_FLAG_VALUES
         self.flags_with_values["persist_optimize_ignored_data_fetch"] = (
             BOOLEAN_FLAG_VALUES
         )
@@ -2109,6 +2112,7 @@ def __init__(
             "persist_consensus_connection_pool_max_wait",
             "persist_consensus_connection_pool_ttl",
             "persist_consensus_connection_pool_ttl_stagger",
+            "persist_consensus_connection_pool_drain_culls_per_tick",
             "persist_use_postgres_tuned_queries",
             "crdb_connect_timeout",
             "crdb_tcp_user_timeout",

src/persist-client/src/cfg.rs

@@ -299,6 +299,8 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet {
         .add(&BLOB_OPERATION_ATTEMPT_TIMEOUT)
         .add(&BLOB_CONNECT_TIMEOUT)
         .add(&BLOB_READ_TIMEOUT)
+        .add(&crate::cfg::CONSENSUS_CONNECTION_POOL_DRAIN_AWARE_RECYCLING)
+        .add(&crate::cfg::CONSENSUS_CONNECTION_POOL_DRAIN_CULLS_PER_TICK)
         .add(&crate::cfg::CONSENSUS_CONNECTION_POOL_MAX_SIZE)
         .add(&crate::cfg::CONSENSUS_CONNECTION_POOL_MAX_WAIT)
         .add(&crate::cfg::CONSENSUS_CONNECTION_POOL_TTL_STAGGER)
@@ -387,7 +389,8 @@ impl PersistConfig {
 
 /// Sets the maximum size of the connection pool that is used by consensus.
 ///
-/// Requires a restart of the process to take effect.
+/// Applies to existing pools without a restart. Each pool picks up a changed
+/// value the next time a connection is requested from it.
 pub const CONSENSUS_CONNECTION_POOL_MAX_SIZE: Config<usize> = Config::new(
     "persist_consensus_connection_pool_max_size",
     50,
@@ -428,6 +431,31 @@ const CONSENSUS_CONNECTION_POOL_TTL_STAGGER: Config<Duration> = Config::new(
     "The minimum time between TTLing Consensus connections to Postgres/CRDB.",
 );
 
+/// Whether consensus connections whose CockroachDB node is draining are
+/// proactively recycled, ahead of the server hard-closing them at the end of
+/// its drain. Only effective on CockroachDB backends, and only where the
+/// consensus role can read `crdb_internal.gossip_liveness` (see
+/// [`mz_postgres_client::PostgresClientKnobs::drain_aware_recycling`]). Leave
+/// off for vanilla Postgres.
+pub const CONSENSUS_CONNECTION_POOL_DRAIN_AWARE_RECYCLING: Config<bool> = Config::new(
+    "persist_consensus_connection_pool_drain_aware_recycling",
+    false,
+    "\
+    Proactively recycle consensus connections whose CockroachDB node is \
+    draining. Requires a role that can read crdb_internal.gossip_liveness.",
+);
+
+/// How many idle consensus connections on draining CockroachDB nodes are
+/// discarded per sweep, bounding how fast the pool sheds them. See
+/// [`CONSENSUS_CONNECTION_POOL_DRAIN_AWARE_RECYCLING`].
+const CONSENSUS_CONNECTION_POOL_DRAIN_CULLS_PER_TICK: Config<usize> = Config::new(
+    "persist_consensus_connection_pool_drain_culls_per_tick",
+    2,
+    "\
+    How many idle consensus connections on draining CockroachDB nodes are \
+    discarded per sweep.",
+);
+
 /// The duration to wait for a Consensus Postgres/CRDB connection to be made
 /// before retrying.
 pub const CRDB_CONNECT_TIMEOUT: Config<Duration> = Config::new(
@@ -600,6 +628,14 @@ impl PostgresClientKnobs for PersistConfig {
         CONSENSUS_CONNECTION_POOL_MAX_SIZE.get(self)
     }
 
+    fn drain_aware_recycling(&self) -> bool {
+        CONSENSUS_CONNECTION_POOL_DRAIN_AWARE_RECYCLING.get(self)
+    }
+
+    fn connection_pool_drain_culls_per_tick(&self) -> usize {
+        CONSENSUS_CONNECTION_POOL_DRAIN_CULLS_PER_TICK.get(self)
+    }
+
     fn connection_pool_max_wait(&self) -> Option<Duration> {
         Some(CONSENSUS_CONNECTION_POOL_MAX_WAIT.get(self))
     }

src/postgres-client/Cargo.toml

@@ -11,7 +11,6 @@ workspace = true
 
 [dependencies]
 anyhow.workspace = true
-async-trait.workspace = true
 deadpool.workspace = true
 deadpool-postgres.workspace = true
 mz-ore = { path = "../ore", default-features = false, features = ["metrics", "async", "bytes"] }
@@ -20,5 +19,8 @@ prometheus.workspace = true
 tokio.workspace = true
 tracing.workspace = true
 
+[dev-dependencies]
+mz-ore = { path = "../ore", default-features = false, features = ["test"] }
+
 [features]
 default = []