[persist] Postgres specific consensus queries (#32389)

This PR adds some new queries to `PostgresConsensus` that are tuned
against vanilla postgres to obtain row level locks as opposed to the
table level locks that the current queries get. These new queries are
only used if `persist_use_postgres_tuned_queries` is set to true,
__and__ we detect we're running against vanilla Postgres.

Locks acquired on `main`
```
  pid  |       mode       | granted
-------+------------------+---------
 21004 | AccessShareLock  | t
 21004 | RowExclusiveLock | t
```

Locks acquired with the new feature enabled
```
  pid  |       mode       | granted
-------+------------------+---------
 21004 | RowShareLock     | t
 21004 | RowExclusiveLock | t
```

It's not entirely clear that this improves performance, but some
simulation testing seems to indicate that it does.

### Testing

Also included in this PR are new CI workflows that run platform-checks
and maelstrom against Postgres consensus, using the new Postgres
queries.


### Motivation

Improve performance of self-hosted which runs against vanilla postgres

### Tips for reviewer

<!--
Leave some tips for your reviewer, like:

    * The diff is much smaller if viewed with whitespace hidden.
    * [Some function/module/file] deserves extra attention.
* [Some function/module/file] is pure code movement and only needs a
skim.

Delete this section if no tips.
-->

### Checklist

- [ ] This PR has adequate test coverage / QA involvement has been duly
considered. ([trigger-ci for additional test/nightly
runs](https://trigger-ci.dev.materialize.com/))
- [ ] This PR has an associated up-to-date [design
doc](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/README.md),
is a design doc
([template](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/00000000_template.md)),
or is sufficiently small to not require a design.
  <!-- Reference the design in the description. -->
- [ ] If this PR evolves [an existing `$T ⇔ Proto$T`
mapping](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/command-and-response-binary-encoding.md)
(possibly in a backwards-incompatible way), then it is tagged with a
`T-proto` label.
- [ ] If this PR will require changes to cloud orchestration or tests,
there is a companion cloud PR to account for those changes that is
tagged with the release-blocker label
([example](MaterializeInc/cloud#5021)).
<!-- Ask in #team-cloud on Slack if you need help preparing the cloud
PR. -->
- [ ] If this PR includes major [user-facing behavior
changes](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/guide-changes.md#what-changes-require-a-release-note),
I have pinged the relevant PM to schedule a changelog post.

Author: ParkMyCar

Cargo.lock

@@ -787,6 +787,24 @@ steps:
                   "--seed=$BUILDKITE_JOB_ID",
                 ]
 
+      - id: checks-parallel-drop-create-default-replica-postgres
+        label: "Checks parallel + DROP/CREATE replica (:postgres: Consensus)"
+        depends_on: build-aarch64
+        timeout_in_minutes: 180
+        parallelism: 2
+        agents:
+          queue: hetzner-aarch64-16cpu-32gb
+        plugins:
+          - ./ci/plugins/mzcompose:
+              composition: platform-checks
+              args:
+                [
+                  --scenario=DropCreateDefaultReplica,
+                  --execution-mode=parallel,
+                  --features=postgres_consensus,
+                  --system-param=persist_use_postgres_tuned_queries=true,
+                  "--seed=$BUILDKITE_JOB_ID",
+                ]
 
       - id: checks-parallel-restart-clusterd-compute
         label: "Checks parallel + restart compute clusterd"
@@ -824,6 +842,24 @@ steps:
               composition: platform-checks
               args: [--scenario=RestartEnvironmentdClusterdStorage, --execution-mode=parallel, "--seed=$BUILDKITE_JOB_ID"]
 
+      - id: checks-parallel-restart-environmentd-clusterd-storage-postgres-consensus
+        label: "Checks parallel + restart of environmentd & storage clusterd (:postgres: Consensus)"
+        depends_on: build-aarch64
+        timeout_in_minutes: 180
+        parallelism: 2
+        agents:
+          queue: hetzner-aarch64-16cpu-32gb
+        plugins:
+          - ./ci/plugins/mzcompose:
+              composition: platform-checks
+              args: [
+                --scenario=RestartEnvironmentdClusterdStorage,
+                --execution-mode=parallel,
+                --features=postgres_consensus,
+                --system-param=persist_use_postgres_tuned_queries=true,
+                "--seed=$BUILDKITE_JOB_ID",
+              ]
+
       - id: checks-parallel-kill-clusterd-storage
         label: "Checks parallel + kill storage clusterd"
         depends_on: build-aarch64
@@ -1107,7 +1143,7 @@ steps:
               args: [--node-count=1, --consensus=mem, --blob=mem, --time-limit=600, --concurrency=4, --rate=500, --max-txn-length=16, --unreliability=0.1]
 
       - id: persist-maelstrom-multi-node
-        label: Long multi-node Maelstrom coverage of persist with postgres consensus
+        label: Long multi-node Maelstrom coverage of persist with CockroachDB consensus
         depends_on: build-aarch64
         timeout_in_minutes: 40
         agents:
@@ -1118,8 +1154,20 @@ steps:
               composition: persist
               args: [--node-count=4, --consensus=cockroach, --blob=maelstrom, --time-limit=300, --concurrency=4, --rate=500, --max-txn-length=16, --unreliability=0.1]
 
+      - id: persist-maelstrom-multi-node-postgres
+        label: "Long multi-node Maelstrom coverage of persist with :postgres: Consensus"
+        depends_on: build-aarch64
+        timeout_in_minutes: 40
+        agents:
+          queue: hetzner-aarch64-4cpu-8gb
+        artifact_paths: [test/persist/maelstrom/**/*.log]
+        plugins:
+          - ./ci/plugins/mzcompose:
+              composition: persist
+              args: [--node-count=4, --consensus=postgres, --blob=maelstrom, --time-limit=300, --concurrency=4, --rate=500, --max-txn-length=16, --unreliability=0.1]
+
       - id: txn-wal-maelstrom
-        label: Maelstrom coverage of txn-wal
+        label: Maelstrom coverage of txn-wal with CockroachDB Consensus
         depends_on: build-aarch64
         timeout_in_minutes: 40
         agents:
@@ -1130,6 +1178,18 @@ steps:
               composition: persist
               args: [--node-count=4, --consensus=cockroach, --blob=maelstrom, --time-limit=300, --rate=500, --txn-wal]
 
+      - id: txn-wal-maelstrom-postgres
+        label: "Maelstrom coverage of txn-wal with :postgres: Consensus"
+        depends_on: build-aarch64
+        timeout_in_minutes: 40
+        agents:
+          queue: hetzner-aarch64-8cpu-16gb
+        artifact_paths: [test/persist/maelstrom/**/*.log]
+        plugins:
+          - ./ci/plugins/mzcompose:
+              composition: persist
+              args: [--node-count=4, --consensus=postgres, --blob=maelstrom, --time-limit=300, --rate=500, --txn-wal]
+
       - id: persistence-failpoints
         label: Persistence failpoints
         depends_on: build-aarch64

ci/nightly/pipeline.template.yml

@@ -15,6 +15,7 @@
 class Features:
     AZURITE = "azurite"
     SQL_SERVER = "sql_server"
+    POSTGRES_CONSENSUS = "postgres_consensus"
 
     def __init__(self, features):
         self.features = features
@@ -24,3 +25,6 @@ def azurite_enabled(self) -> bool:
 
     def sql_server_enabled(self) -> bool:
         return self.features and self.SQL_SERVER in self.features
+
+    def postgres_consensus_enabled(self) -> bool:
+        return self.features and self.POSTGRES_CONSENSUS in self.features

misc/python/materialize/checks/features.py

@@ -36,6 +36,7 @@ rust_binary(
     }),
     version = "0.0.0",
     deps = [
+        "//src/dyncfg:mz_dyncfg",
         "//src/http-util:mz_http_util",
         "//src/orchestrator-tracing:mz_orchestrator_tracing",
         "//src/ore:mz_ore",

src/persist-cli/BUILD.bazel

@@ -26,6 +26,7 @@ clap = { version = "4.5.23", features = ["derive", "env"] }
 differential-dataflow = "0.14.2"
 futures = "0.3.31"
 humantime = "2.2.0"
+mz-dyncfg = { path = "../dyncfg" }
 mz-http-util = { path = "../http-util" }
 mz-orchestrator-tracing = { path = "../orchestrator-tracing" }
 mz-ore = { path = "../ore", features = ["bytes", "network", "panic", "tracing", "test"] }

src/persist-cli/Cargo.toml

@@ -19,6 +19,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 use async_trait::async_trait;
 use differential_dataflow::consolidation::consolidate_updates;
+use mz_dyncfg::ConfigUpdates;
 use mz_ore::metrics::MetricsRegistry;
 use mz_ore::now::{NOW_ZERO, SYSTEM_TIME};
 use mz_persist::cfg::{BlobConfig, ConsensusConfig};
@@ -372,6 +373,14 @@ impl Service for TransactorService {
 
         let mut config =
             PersistConfig::new_default_configs(&mz_persist_client::BUILD_INFO, SYSTEM_TIME.clone());
+        {
+            // We only use the Postgres tuned queries when connected to vanilla
+            // Postgres, so we always want to enable them for testing.
+            let mut updates = ConfigUpdates::default();
+            updates.add(&mz_persist::postgres::USE_POSTGRES_TUNED_QUERIES, true);
+            config.apply_from(&updates);
+        }
+
         let metrics_registry = MetricsRegistry::new();
         let metrics = Arc::new(PersistMetrics::new(&config, &metrics_registry));
 
@@ -417,6 +426,7 @@ impl Service for TransactorService {
                     consensus_uri,
                     Box::new(config.clone()),
                     metrics.postgres_consensus.clone(),
+                    Arc::clone(&config.configs),
                 )
                 .expect("consensus_uri should be valid");
                 loop {

src/persist-cli/src/maelstrom/txn_list_append_multi.rs

@@ -17,6 +17,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use async_trait::async_trait;
 use differential_dataflow::consolidation::consolidate_updates;
 use differential_dataflow::lattice::Lattice;
+use mz_dyncfg::ConfigUpdates;
 use mz_ore::metrics::MetricsRegistry;
 use mz_ore::now::SYSTEM_TIME;
 use mz_persist::cfg::{BlobConfig, ConsensusConfig};
@@ -596,6 +597,14 @@ impl Service for TransactorService {
 
         let mut config =
             PersistConfig::new_default_configs(&mz_persist_client::BUILD_INFO, SYSTEM_TIME.clone());
+        {
+            // We only use the Postgres tuned queries when connected to vanilla
+            // Postgres, so we always want to enable them for testing.
+            let mut updates = ConfigUpdates::default();
+            updates.add(&mz_persist::postgres::USE_POSTGRES_TUNED_QUERIES, true);
+            config.apply_from(&updates);
+        }
+
         let metrics = Arc::new(Metrics::new(&config, &MetricsRegistry::new()));
 
         // Construct requested Blob.
@@ -640,6 +649,7 @@ impl Service for TransactorService {
                     consensus_uri,
                     Box::new(config.clone()),
                     metrics.postgres_consensus.clone(),
+                    Arc::clone(&config.configs),
                 )
                 .expect("consensus_uri should be valid");
                 loop {

src/persist-cli/src/maelstrom/txn_list_append_single.rs

@@ -173,6 +173,7 @@ impl PersistClientCache {
                     x.key(),
                     Box::new(self.cfg.clone()),
                     self.metrics.postgres_consensus.clone(),
+                    Arc::clone(&self.cfg().configs),
                 )?;
                 let consensus =
                     retry_external(&self.metrics.retries.external.consensus_open, || {

src/persist-client/src/cache.rs

@@ -130,6 +130,7 @@ pub(super) async fn make_consensus(
         consensus_uri,
         Box::new(cfg.clone()),
         metrics.postgres_consensus.clone(),
+        Arc::clone(&cfg.configs),
     )?;
     let consensus = consensus.clone().open().await?;
     let consensus = if commit {

src/persist-client/src/cli/args.rs

@@ -36,6 +36,7 @@ pub fn all_dyn_configs(configs: ConfigSet) -> ConfigSet {
         .add(&crate::indexed::columnar::arrow::ENABLE_ARROW_LGALLOC_NONCC_SIZES)
         .add(&crate::s3::ENABLE_S3_LGALLOC_CC_SIZES)
         .add(&crate::s3::ENABLE_S3_LGALLOC_NONCC_SIZES)
+        .add(&crate::postgres::USE_POSTGRES_TUNED_QUERIES)
 }
 
 /// Config for an implementation of [Blob].
@@ -233,10 +234,11 @@ impl ConsensusConfig {
         url: &SensitiveUrl,
         knobs: Box<dyn PostgresClientKnobs>,
         metrics: PostgresClientMetrics,
+        dyncfg: Arc<ConfigSet>,
     ) -> Result<Self, ExternalError> {
         let config = match url.scheme() {
             "postgres" | "postgresql" => Ok(ConsensusConfig::Postgres(
-                PostgresConsensusConfig::new(url, knobs, metrics)?,
+                PostgresConsensusConfig::new(url, knobs, metrics, dyncfg)?,
             )),
             "mem" => {
                 if !cfg!(debug_assertions) {