From a731e34477f357b76f205a5b2cceb09ee334f3e1 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 15 Oct 2026 19:51:49 -0700 Subject: [PATCH 1/7] Added the notification log to the PostgreSQL store. A driver reads the log as it reads the Redis one, so producers on PostgreSQL send no notifications. A listener connection hears the doorbells, and with notify_reads it wakes waiting reads in other processes at once. --- .../temporal/sdk/core/streams/streams.proto | 9 + crates/sdk-core/CHANGELOG.md | 6 + crates/streams/docs/postgres.md | 48 +++- .../streams/scripts/postgres-test-server.sh | 24 +- crates/streams/src/lib.rs | 2 +- crates/streams/src/postgres/conn.rs | 50 +++- crates/streams/src/postgres/janitor.rs | 8 +- crates/streams/src/postgres/listen.rs | 214 ++++++++++++++++++ crates/streams/src/postgres/log.rs | 121 ++++++++++ crates/streams/src/postgres/migrations.rs | 20 +- .../migrations/0002_notification_log.sql | 126 +++++++++++ crates/streams/src/postgres/mod.rs | 132 +++++++++-- crates/streams/src/postgres/wait.rs | 72 ++++-- crates/streams/src/service.rs | 8 +- crates/streams/tests/conformance.rs | 99 +++++++- crates/streams/tests/driver.rs | 149 +++++++++--- crates/streams/tests/postgres_store.rs | 98 ++++++++ 17 files changed, 1096 insertions(+), 90 deletions(-) create mode 100644 crates/streams/src/postgres/listen.rs create mode 100644 crates/streams/src/postgres/log.rs create mode 100644 crates/streams/src/postgres/migrations/0002_notification_log.sql diff --git a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto index 3d4c6efbb7..8e6db34d99 100644 --- a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto +++ b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto @@ -162,8 +162,17 @@ message PostgresStoreConfig { google.protobuf.Duration floor_lease = 6; // How long a hold lasts after its reader released it. Unset keeps one day. google.protobuf.Duration floor_release_grace = 7; + // With `StreamStoreConfig.notification_log`, how many doorbells the log rings, from 1 to 64. + // Zero uses 16. + uint32 notify_doorbells = 8; + // A URL for the connection that listens and the one that applies migrations, when `url` goes + // through a pooler in transaction mode. Empty uses `url`. + string listen_url = 9; // How often a process looks for new records for its waiting reads. Unset uses 100 ms. google.protobuf.Duration poll_interval = 10; + // Each append also sends NOTIFY, so waiting reads in other processes wake in milliseconds. + // Costs append throughput when many producers write at once. + bool notify_reads = 11; // The store checks the schema version and never changes it. A DBA applies the migrations. bool skip_migrations = 12; // How often the janitor trims records past the retention. Unset uses one minute. diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index 2a1346ed6c..6ca3f115a6 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -83,6 +83,12 @@ relevant information. every `janitor_interval` (one minute). Waiting reads hold no connection and wake within `poll_interval` (100 ms). TLS follows the URL's `sslmode` and `sslrootcert`. The C bridge builds the store. +* Experimental: the PostgreSQL store keeps a notification log with `notification_log`, a + `notifications` table with 64 partitions that the append writes in its transaction. Drivers take + entries with `FOR UPDATE SKIP LOCKED`, and a `NOTIFY` doorbell rings when a partition takes its + first entry. `listen_url` names a direct connection for the listener and the migrations when + `url` goes through PgBouncer in transaction mode. `notify_reads` sends a `NOTIFY` per append, so + waiting reads in other processes wake at once. * Experimental: the new `FlushNotifications` stream call waits up to 10 seconds for the notifications still out. A process calls it before it stops. * Experimental: a server without the stream notifier answers the first notifier call with diff --git a/crates/streams/docs/postgres.md b/crates/streams/docs/postgres.md index 0ccdb114ca..3b4101a5d3 100644 --- a/crates/streams/docs/postgres.md +++ b/crates/streams/docs/postgres.md @@ -41,6 +41,30 @@ store, or for the store's poll. While any read waits, the store runs one stateme another process wakes within about `poll_interval`. The server's stream notifier is the main wake path, and this wait is the path for a reader without it. +**`notify_reads`.** With it, each append also sends `NOTIFY` on `_read_` +with the topic id, and the store's listener wakes the waiters of that topic. A reader in another +process then wakes in a few milliseconds. It is off by default, since a transaction that sent +`NOTIFY` takes a lock on the server's notification queue as it commits, so such commits run one +at a time across the database. With many producers at once, it costs append throughput. + +**Notification log.** With `notification_log` on, the append also writes a row to +`notifications` in its transaction: the topic and the append's last position, in one of 64 +partitions fixed in the schema. Every topic of a chain lands in the partition that the first four +bytes of SHA-256 over the chain's fields give. A repeat of the newest batch writes none. When the +partition held no entry before, the append also sends `NOTIFY` on `_bell_` with the partition, so the doorbell costs nothing while a partition stays busy. Drivers +take entries with `FOR UPDATE SKIP LOCKED`: a consumer's own unacked entries, those another +consumer left for `claim_idle`, and new ones. An ack deletes them. The janitor keeps the newest +10,000 entries per partition, so a log with no driver stays bounded. An entry whose topic an owner +delete removed is acked as broken when a driver takes it. + +**Listener.** One connection per store listens for the doorbells and, with `notify_reads`, the read +channels. It connects with the store, so a store whose listener can't connect fails at once. When +the connection ends, it connects again with backoff, rings every doorbell for every partition and +wakes every waiting read, since what was sent while it was away is gone. The driver's sweep finds a +lost ring too. `LISTEN` keeps state on the server session, so the listener uses `listen_url` when +it is set, and `url` otherwise. Both must name the primary, since a standby refuses `LISTEN`. + **Retention.** PostgreSQL has no key expiry, so each store runs a janitor every `janitor_interval` (one minute). A session advisory lock lets one pass run at a time across processes. A pass trims the topics that are due, at most 10,000 records per topic per @@ -69,10 +93,12 @@ so it reports `REFUSED`. A read reports `STORAGE`. A server below PostgreSQL 14 ## Migrations The store embeds its migrations and applies the missing ones at connect, each in its own -transaction, under a session advisory lock on the schema's name. It refuses to start when an -applied migration differs from the embedded one, or when the database has a migration that needs a -newer build. A function's name carries its version, such as `append_v1`, so a migration that -changes a contract adds a new function and keeps the old one for a release. +transaction, under a session advisory lock on the schema's name. They run on a connection of their +own to `listen_url` when it is set, since a pooler in transaction mode can't keep a session lock. +The store refuses to start when an applied migration differs from the embedded one, or when the +database has a migration that needs a newer build. A function's name carries its version, such +as `append_v1`, so a migration that changes a contract adds a new function and keeps the old one +for a release. With `skip_migrations`, the store only checks that every migration it knows is applied. A DBA then applies the files in `src/postgres/migrations/` in order, after `SET search_path = , @@ -127,8 +153,10 @@ IAM database authentication is not in this release, so the password comes from t **PgBouncer.** Session mode works with no setting. Transaction mode works from PgBouncer 1.21 with `max_prepared_statements` set, since the store prepares its statements. The session settings apply to the server connection that made each pooled connection, so all clients of that -PgBouncer database should use the same `READ COMMITTED` default. Migrations take a session lock, -so run them once through a direct connection, or use `skip_migrations`. +PgBouncer database should use the same `READ COMMITTED` default. Set `listen_url` to a direct +connection, which the listener and the migrations use. The conformance suite runs through +PgBouncer 1.26 in transaction mode with `max_prepared_statements = 200` and a direct +`listen_url`. **Servers that speak the protocol.** CockroachDB, YugabyteDB and other servers that speak the PostgreSQL protocol are not supported. The store needs advisory locks, PL/pgSQL and PostgreSQL's @@ -142,9 +170,11 @@ Redis. ## Tests -`scripts/postgres-test-server.sh` starts PostgreSQL 14 and 17 in Docker on free ports and prints -the URLs to export. The suite reads `STREAMS_POSTGRES_URL`, a role that may create schemas. Each -case runs in a schema of its own. +`scripts/postgres-test-server.sh` starts PostgreSQL 14 and 17 in Docker on free ports, and +PgBouncer in transaction mode in front of 17, and prints the URLs to export. The suite reads +`STREAMS_POSTGRES_URL`, a role that may create schemas, and `STREAMS_POSTGRES_PGBOUNCER_URL`. Each +case runs in a schema of its own. A schema name takes at most 55 bytes, so that the channel names +stay within 63. ```sh eval "$(crates/streams/scripts/postgres-test-server.sh)" diff --git a/crates/streams/scripts/postgres-test-server.sh b/crates/streams/scripts/postgres-test-server.sh index 25b2465744..fdf8c62592 100755 --- a/crates/streams/scripts/postgres-test-server.sh +++ b/crates/streams/scripts/postgres-test-server.sh @@ -1,20 +1,21 @@ #!/usr/bin/env bash # Starts the PostgreSQL servers the stream store tests run against, 14 and 17, in Docker on free -# ports, with TLS on from the image's self-signed certificate. Prints the variables to export. -# `stop` stops them. STREAMS_POSTGRES_NAME moves the container names, so two runs can share a -# host. +# ports, with TLS on from the image's self-signed certificate, and PgBouncer in transaction mode +# in front of 17 when its image pulls. Prints the variables to export. `stop` stops them. +# STREAMS_POSTGRES_NAME moves the container names, so two runs can share a host. # # eval "$(crates/streams/scripts/postgres-test-server.sh)" # STREAMS_POSTGRES_URL=postgres://postgres:streams@localhost:/postgres (17) # STREAMS_POSTGRES_14_URL=postgres://postgres:streams@localhost:/postgres (14) # STREAMS_POSTGRES_TLS_URL=$STREAMS_POSTGRES_URL +# STREAMS_POSTGRES_PGBOUNCER_URL=postgres://postgres:streams@localhost:/postgres set -euo pipefail NAME=${STREAMS_POSTGRES_NAME:-streams-postgres} VERSIONS=(14 17) if [ "${1:-}" = stop ]; then - for version in "${VERSIONS[@]}"; do + for version in "${VERSIONS[@]}" pgbouncer; do docker rm -f "$NAME-$version" >/dev/null 2>&1 || true done exit 0 @@ -42,6 +43,21 @@ for version in "${VERSIONS[@]}"; do eval "URL_$version=postgres://postgres:streams@localhost:$port/postgres" done +# Prepared statements in transaction mode need PgBouncer 1.21 or later. +# PgBouncer resolves names itself, so it reaches 17 by its address on the Docker bridge. +bouncer="$NAME-pgbouncer" +address_17=$(docker inspect -f '{{.NetworkSettings.Networks.bridge.IPAddress}}' "$NAME-17") +if ! docker inspect "$bouncer" >/dev/null 2>&1; then + docker run -d --rm --name "$bouncer" -p "127.0.0.1:$(free_port):5432" \ + -e DATABASE_URL="postgres://postgres:streams@$address_17:5432/postgres" \ + -e POOL_MODE=transaction -e MAX_PREPARED_STATEMENTS=200 -e AUTH_TYPE=scram-sha-256 \ + -e MAX_CLIENT_CONN=1000 -e DEFAULT_POOL_SIZE=40 -e LISTEN_PORT=5432 \ + edoburu/pgbouncer:latest >/dev/null 2>&1 || echo "PgBouncer did not start" >&2 +fi +if docker inspect "$bouncer" >/dev/null 2>&1; then + port=$(docker port "$bouncer" 5432/tcp | head -1 | sed 's/.*://') + echo "export STREAMS_POSTGRES_PGBOUNCER_URL=postgres://postgres:streams@localhost:$port/postgres" +fi echo "export STREAMS_POSTGRES_URL=$URL_17" echo "export STREAMS_POSTGRES_14_URL=$URL_14" echo "export STREAMS_POSTGRES_TLS_URL=$URL_17" diff --git a/crates/streams/src/lib.rs b/crates/streams/src/lib.rs index 4f54c56836..31f04c3332 100644 --- a/crates/streams/src/lib.rs +++ b/crates/streams/src/lib.rs @@ -32,7 +32,7 @@ mod subscribe; mod tiered; #[cfg(feature = "postgres")] -pub use self::postgres::{PostgresStore, PostgresStoreOptions}; +pub use self::postgres::{DEFAULT_POSTGRES_DOORBELLS, PostgresStore, PostgresStoreOptions}; #[cfg(feature = "redis")] pub use self::redis::{DEFAULT_NOTIFY_DOORBELLS, RedisStore, RedisStoreOptions}; #[cfg(feature = "archive")] diff --git a/crates/streams/src/postgres/conn.rs b/crates/streams/src/postgres/conn.rs index 8e5f676587..1704d6a047 100644 --- a/crates/streams/src/postgres/conn.rs +++ b/crates/streams/src/postgres/conn.rs @@ -9,7 +9,10 @@ use rustls::{ pki_types::{CertificateDer, ServerName, UnixTime, pem::PemObject}, }; use std::{sync::Arc, time::Duration}; -use tokio_postgres::{Client, NoTls, config::SslMode}; +use tokio::sync::mpsc; +use tokio_postgres::{ + AsyncMessage, Client, Connection, NoTls, Notification, config::SslMode, tls::TlsStream, +}; use tokio_postgres_rustls::MakeRustlsConnect; /// The `sslmode` values libpq takes. `allow` is left out, since it trusts a plain connection @@ -93,6 +96,27 @@ impl Target { } } + /// One connection of its own whose notifications arrive on the receiver. The receiver + /// closes when the connection ends. + pub(super) async fn listen( + &self, + ) -> Result<(Client, mpsc::UnboundedReceiver), tokio_postgres::Error> { + let (sender, receiver) = mpsc::unbounded_channel(); + let client = match &self.tls { + Some(tls) => { + let (client, connection) = self.config.connect(tls.clone()).await?; + tokio::spawn(forward(connection, sender)); + client + } + None => { + let (client, connection) = self.config.connect(NoTls).await?; + tokio::spawn(forward(connection, sender)); + client + } + }; + Ok((client, receiver)) + } + /// The pool of `size` connections. Each new connection pins the session settings the store /// functions depend on. pub(super) fn pool(&self, size: usize, timeout: Duration) -> StreamResult { @@ -124,6 +148,30 @@ impl Target { } } +/// Drives a connection and hands its notifications on, until it ends. +async fn forward( + mut connection: Connection, + sender: mpsc::UnboundedSender, +) where + S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, + T: TlsStream + Unpin, +{ + while let Some(message) = std::future::poll_fn(|cx| connection.poll_message(cx)).await { + match message { + Ok(AsyncMessage::Notification(notification)) => { + if sender.send(notification).is_err() { + return; + } + } + Ok(_) => {} + Err(error) => { + tracing::debug!("The PostgreSQL store's listener connection ended: {error}"); + return; + } + } + } +} + /// Every call is one statement that depends on `READ COMMITTED`: a statement that waited for a /// row lock reads the newest version of the row. The session pins it, so a database default of /// `SERIALIZABLE` can't change it. diff --git a/crates/streams/src/postgres/janitor.rs b/crates/streams/src/postgres/janitor.rs index 3d332b3a01..68822b464d 100644 --- a/crates/streams/src/postgres/janitor.rs +++ b/crates/streams/src/postgres/janitor.rs @@ -2,7 +2,7 @@ //! session advisory lock lets one pass run at a time across processes. use super::{ - DELETE_BATCH, Inner, + DELETE_BATCH, Inner, LOG_CAP, errors::{Call, mapped, timed_out, unreachable}, migrations::LOCK_SPACE, }; @@ -29,7 +29,8 @@ pub(super) async fn run(inner: Weak, interval: Duration) { } impl Inner { - /// One pass: trims the topics that are due, finishes purges, and deletes idle chains. + /// One pass: trims the topics that are due, finishes purges, caps the notification log, and + /// deletes idle chains. pub(super) async fn janitor_pass(&self) -> StreamResult { let client = self.pool.get().await.map_err(unreachable)?; let key = format!("{}:janitor", self.schema); @@ -89,6 +90,9 @@ impl Inner { break; } } + if self.log { + self.on(client, &self.sql.cap_log, &[&LOG_CAP]).await?; + } self.on( client, &self.sql.idle_chains, diff --git a/crates/streams/src/postgres/listen.rs b/crates/streams/src/postgres/listen.rs new file mode 100644 index 0000000000..b979072a42 --- /dev/null +++ b/crates/streams/src/postgres/listen.rs @@ -0,0 +1,214 @@ +//! The listener: one connection per store that hears the notification log's doorbells and, with +//! `notify_reads`, the appends that wake waiting reads. `LISTEN` keeps state on the server +//! session, so this connection never goes through a pooler in transaction mode. + +use super::{Inner, PARTITIONS, conn::Target, errors::unreachable}; +use crate::StreamResult; +use std::{ + collections::BTreeSet, + sync::{Arc, Mutex, Weak}, + time::Duration, +}; +use tokio::sync::{Notify, mpsc}; +use tokio_postgres::{Client, Notification}; + +const FIRST_BACKOFF: Duration = Duration::from_millis(100); +const LAST_BACKOFF: Duration = Duration::from_secs(5); + +/// The partitions each doorbell rang that no driver took yet. +#[derive(Debug)] +pub(super) struct Rings { + pending: Mutex>>, + bells: Vec, +} + +impl Rings { + pub(super) fn new(doorbells: u32) -> Self { + Self { + pending: Mutex::new(vec![BTreeSet::new(); doorbells as usize]), + bells: (0..doorbells).map(|_| Notify::new()).collect(), + } + } + + fn ring(&self, partition: u32) { + let doorbell = partition as usize % self.bells.len(); + self.pending.lock().unwrap()[doorbell].insert(partition); + self.bells[doorbell].notify_waiters(); + } + + /// Rings every partition, for when rings may have been lost. + fn ring_all(&self) { + for partition in 0..PARTITIONS { + self.ring(partition); + } + } + + fn take(&self, doorbell: usize) -> Vec { + std::mem::take(&mut self.pending.lock().unwrap()[doorbell]) + .into_iter() + .collect() + } + + /// Waits up to `block` for doorbell `doorbell`, and answers the partitions it rang. + pub(super) async fn wait(&self, doorbell: u32, block: Duration) -> Vec { + let Some(bell) = self.bells.get(doorbell as usize) else { + tokio::time::sleep(block).await; + return vec![]; + }; + let notified = bell.notified(); + tokio::pin!(notified); + // Enabled before the look, so a ring between the two still wakes the wait. + notified.as_mut().enable(); + let rung = self.take(doorbell as usize); + if !rung.is_empty() { + return rung; + } + let _ = tokio::time::timeout(block, notified).await; + self.take(doorbell as usize) + } +} + +/// What the listener listens on and whom it wakes. +pub(super) struct Listener { + pub(super) target: Target, + pub(super) schema: String, + /// The `LISTEN` statements, one per channel. + pub(super) listens: String, + pub(super) rings: Option>, +} + +impl Listener { + pub(super) fn new( + target: Target, + schema: &str, + rings: Option>, + notify_reads: bool, + ) -> Self { + let mut channels: Vec = vec![]; + if let Some(rings) = &rings { + channels.extend((0..rings.bells.len()).map(|d| format!("{schema}_bell_{d}"))); + } + if notify_reads { + channels.extend((0..64).map(|n| format!("{schema}_read_{n}"))); + } + let listens = channels + .iter() + .map(|channel| format!("LISTEN {channel}")) + .collect::>() + .join("; "); + Self { + target, + schema: schema.to_string(), + listens, + rings, + } + } + + /// Connects and listens. The store's connect calls this once, so a listener that can't + /// connect fails the store at once. + pub(super) async fn connect( + &self, + ) -> StreamResult<(Client, mpsc::UnboundedReceiver)> { + let (client, notifications) = self.target.listen().await.map_err(unreachable)?; + client + .batch_execute(&self.listens) + .await + .map_err(unreachable)?; + Ok((client, notifications)) + } + + fn deliver(&self, inner: &Weak, notification: &Notification) { + let channel = notification.channel(); + let Some(rest) = channel + .strip_prefix(self.schema.as_str()) + .and_then(|rest| rest.strip_prefix('_')) + else { + return; + }; + if rest.starts_with("bell_") { + if let (Some(rings), Ok(partition)) = (&self.rings, notification.payload().parse()) { + rings.ring(partition); + } + } else if rest.starts_with("read_") + && let (Some(inner), Ok(topic_id)) = (inner.upgrade(), notification.payload().parse()) + { + inner.waiters.wake_topic(topic_id); + } + } + + /// Hands on what the connection hears, and connects again with backoff when it ends. After a + /// reconnect it rings every partition and wakes every waiter, since what was sent while it + /// was away is gone. + pub(super) async fn run( + self, + inner: Weak, + mut client: Client, + mut notifications: mpsc::UnboundedReceiver, + ) { + loop { + while let Some(notification) = notifications.recv().await { + self.deliver(&inner, ¬ification); + } + drop(client); + let mut backoff = FIRST_BACKOFF; + loop { + tokio::time::sleep(backoff).await; + if inner.strong_count() == 0 { + return; + } + match self.connect().await { + Ok(connected) => { + (client, notifications) = connected; + break; + } + Err(error) => { + tracing::debug!("The PostgreSQL store's listener can't reconnect: {error}"); + backoff = (backoff * 2).min(LAST_BACKOFF); + } + } + } + if let Some(rings) = &self.rings { + rings.ring_all(); + } + if let Some(inner) = inner.upgrade() { + inner.waiters.wake_all(); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn a_ring_waits_for_its_doorbell_and_is_taken_once() { + let rings = Rings::new(4); + rings.ring(5); + rings.ring(9); + rings.ring(2); + assert_eq!(rings.wait(1, Duration::ZERO).await, [5, 9]); + assert!(rings.wait(1, Duration::from_millis(10)).await.is_empty()); + assert_eq!(rings.wait(2, Duration::ZERO).await, [2]); + } + + #[tokio::test] + async fn a_ring_wakes_a_waiting_doorbell() { + let rings = Arc::new(Rings::new(16)); + let waiting = { + let rings = rings.clone(); + tokio::spawn(async move { rings.wait(3, Duration::from_secs(10)).await }) + }; + tokio::time::sleep(Duration::from_millis(50)).await; + rings.ring(19); + assert_eq!(waiting.await.unwrap(), [19]); + } + + #[test] + fn ring_all_reaches_every_partition_once() { + let rings = Rings::new(16); + rings.ring_all(); + let total: usize = (0..16).map(|d| rings.take(d).len()).sum(); + assert_eq!(total, PARTITIONS as usize); + } +} diff --git a/crates/streams/src/postgres/log.rs b/crates/streams/src/postgres/log.rs new file mode 100644 index 0000000000..29ff7f5849 --- /dev/null +++ b/crates/streams/src/postgres/log.rs @@ -0,0 +1,121 @@ +//! The notification log on the `notifications` table, for the drivers that read it. + +use super::{Inner, PARTITIONS, errors::Call, listen::Rings}; +use crate::{ + NotificationEntry, NotificationLog, StreamResult, + proto::{ChainId, StreamOwnerKind}, +}; +use std::{sync::Arc, time::Duration}; + +pub(super) struct PostgresLog { + pub(super) inner: Arc, + pub(super) rings: Arc, + pub(super) doorbells: u32, +} + +/// The chain of a `chains` row. A Workflow chain keeps no owner kind, as the Redis store's keys +/// give it back. +fn chain(namespace: String, kind: i16, owner: String, run: String, activity: String) -> ChainId { + match kind { + 2 => ChainId { + namespace, + first_run_id: run, + owner_kind: StreamOwnerKind::Activity as i32, + activity_id: owner, + ..Default::default() + }, + 3 => ChainId { + namespace, + workflow_id: owner, + first_run_id: run, + owner_kind: StreamOwnerKind::WorkflowActivity as i32, + activity_id: activity, + }, + _ => ChainId { + namespace, + workflow_id: owner, + first_run_id: run, + ..Default::default() + }, + } +} + +#[async_trait::async_trait] +impl NotificationLog for PostgresLog { + fn partitions(&self) -> u32 { + PARTITIONS + } + + fn doorbells(&self) -> u32 { + self.doorbells + } + + async fn wait_doorbell( + &self, + doorbell: u32, + _consumer: &str, + block: Duration, + ) -> StreamResult> { + Ok(self.rings.wait(doorbell, block).await) + } + + async fn take( + &self, + partition: u32, + consumer: &str, + claim_idle: Duration, + max: usize, + ) -> StreamResult> { + let inner = &self.inner; + let partition = i16::try_from(partition).unwrap_or(i16::MAX); + let claim_idle_ms = super::millis(claim_idle); + let max = i64::try_from(max.max(1)).unwrap_or(i64::MAX); + let rows = inner + .query( + Call::Write, + &inner.sql.take, + &[&partition, &consumer, &claim_idle_ms, &max], + ) + .await?; + let mut taken = vec![]; + let mut broken = vec![]; + for row in rows { + let id: i64 = row.get(0); + let namespace: Option = row.get(1); + let topic: Option = row.get(6); + match (namespace, topic) { + (Some(namespace), Some(topic)) => taken.push(NotificationEntry { + id: id.to_string(), + chain: chain(namespace, row.get(2), row.get(3), row.get(4), row.get(5)), + topic, + position: row.get::<_, i64>(7).to_string(), + }), + _ => broken.push(id.to_string()), + } + } + // An entry whose topic a delete took can never be delivered. + if !broken.is_empty() { + self.ack(partition as u32, &broken).await?; + } + Ok(taken) + } + + async fn ack(&self, partition: u32, ids: &[String]) -> StreamResult<()> { + let ids: Vec = ids.iter().filter_map(|id| id.parse().ok()).collect(); + if ids.is_empty() { + return Ok(()); + } + let inner = &self.inner; + let partition = i16::try_from(partition).unwrap_or(i16::MAX); + inner + .query(Call::Write, &inner.sql.ack, &[&partition, &ids]) + .await + .map(drop) + } + + async fn busy(&self) -> StreamResult> { + let inner = &self.inner; + let rows = inner.query(Call::Read, &inner.sql.busy, &[]).await?; + Ok(rows.iter().map(|row| row.get::<_, i32>(0) as u32).collect()) + } +} diff --git a/crates/streams/src/postgres/migrations.rs b/crates/streams/src/postgres/migrations.rs index 1d0fff6530..d59442906f 100644 --- a/crates/streams/src/postgres/migrations.rs +++ b/crates/streams/src/postgres/migrations.rs @@ -19,12 +19,20 @@ pub(super) struct Migration { pub(super) compatible_from: i32, } -pub(super) const MIGRATIONS: &[Migration] = &[Migration { - version: 1, - name: "init", - sql: include_str!("migrations/0001_init.sql"), - compatible_from: 1, -}]; +pub(super) const MIGRATIONS: &[Migration] = &[ + Migration { + version: 1, + name: "init", + sql: include_str!("migrations/0001_init.sql"), + compatible_from: 1, + }, + Migration { + version: 2, + name: "notification_log", + sql: include_str!("migrations/0002_notification_log.sql"), + compatible_from: 1, + }, +]; /// The first key of every advisory lock the store takes. pub(super) const LOCK_SPACE: i32 = 1953653099; diff --git a/crates/streams/src/postgres/migrations/0002_notification_log.sql b/crates/streams/src/postgres/migrations/0002_notification_log.sql new file mode 100644 index 0000000000..9dbf25e808 --- /dev/null +++ b/crates/streams/src/postgres/migrations/0002_notification_log.sql @@ -0,0 +1,126 @@ +-- The notification log: one entry per append, in one of 64 partitions fixed here, so a topic's +-- partition stays valid for ever. Drivers take entries with `FOR UPDATE SKIP LOCKED` and ack +-- them with a delete. + +CREATE TABLE notifications ( + partition smallint NOT NULL, + id bigint GENERATED ALWAYS AS IDENTITY, + topic_id bigint NOT NULL, + position bigint NOT NULL, + taken_by text, + taken_at timestamptz, + PRIMARY KEY (partition, id) +) WITH (fillfactor = 70, autovacuum_vacuum_scale_factor = 0.01); + +-- As in migration 1, with the log entry and its doorbell in the append's transaction. +CREATE OR REPLACE FUNCTION append_v1( + p_namespace text, p_kind smallint, p_owner text, p_run text, p_activity text, + p_topic text, p_partition smallint, p_producer text, p_attempt bigint, p_sequence bigint, + p_digest bytea, p_records bytea[], p_keeps_topics boolean, p_log boolean, p_doorbells int, + p_archive_lease_ms bigint, p_retention_ms bigint, p_notify_reads boolean, + OUT status smallint, OUT first_position bigint, OUT last_position bigint, + OUT next_sequence bigint, OUT close bytea) +LANGUAGE plpgsql SET search_path FROM CURRENT AS $$ +DECLARE + v_chain bigint; + v_chain_closed boolean; + v_topic topics%ROWTYPE; + v_held producers%ROWTYPE; + v_count int := cardinality(p_records); + v_now timestamptz; + v_last bigint; +BEGIN + SELECT l.chain_id, l.closed INTO v_chain, v_chain_closed + FROM live_chain_v1(p_namespace, p_kind, p_owner, p_run, p_activity) l; + LOOP + SELECT * INTO v_topic FROM topics t + WHERE t.chain_id = v_chain AND t.topic = p_topic FOR NO KEY UPDATE; + EXIT WHEN FOUND; + -- Checked before the topic exists, so a refused batch leaves no trace of it. + IF v_chain_closed THEN + status := 4; + RETURN; + END IF; + PERFORM topic_id_v1(v_chain, p_topic, p_partition); + END LOOP; + SELECT * INTO v_held FROM producers h + WHERE h.topic_id = v_topic.topic_id AND h.producer_id = p_producer + AND h.attempt = p_attempt; + IF FOUND THEN + IF p_sequence = v_held.sequence THEN + IF p_digest = v_held.digest THEN + status := 1; + first_position := v_held.first_position; + last_position := v_held.first_position + v_held.count - 1; + ELSE + status := 2; + END IF; + RETURN; + END IF; + IF p_sequence < v_held.sequence + v_held.count THEN + status := 3; + next_sequence := v_held.sequence + v_held.count; + RETURN; + END IF; + END IF; + IF v_chain_closed OR v_topic.closed THEN + status := 4; + close := v_topic.close; + RETURN; + END IF; + -- Taken under the row lock, so the times grow with the positions. + v_now := clock_timestamp(); + v_last := coalesce(v_topic.last_position, v_topic.base); + INSERT INTO records (topic_id, position, written_at, record) + SELECT v_topic.topic_id, v_last + u.i, v_now, u.r + FROM unnest(p_records) WITH ORDINALITY AS u(r, i); + INSERT INTO producers VALUES (v_topic.topic_id, p_producer, p_attempt, p_sequence, v_count, + v_last + 1, p_digest, v_now) + ON CONFLICT (topic_id, producer_id, attempt) DO UPDATE + SET sequence = excluded.sequence, count = excluded.count, + first_position = excluded.first_position, digest = excluded.digest, + written_at = excluded.written_at; + UPDATE topics t SET + last_position = v_last + v_count, + written_at = v_now, + written = t.written OR (p_keeps_topics AND p_sequence = 1), + trim_at = coalesce(t.trim_at, v_now + p_retention_ms * interval '1 millisecond') + WHERE t.topic_id = v_topic.topic_id; + IF p_log THEN + -- A busy partition rings no more, so rings stay rare under load. + IF NOT EXISTS (SELECT 1 FROM notifications n WHERE n.partition = v_topic.partition) THEN + PERFORM pg_notify(current_schema() || '_bell_' || (v_topic.partition % p_doorbells), + v_topic.partition::text); + END IF; + INSERT INTO notifications (partition, topic_id, position) + VALUES (v_topic.partition, v_topic.topic_id, v_last + v_count); + END IF; + IF p_notify_reads THEN + PERFORM pg_notify(current_schema() || '_read_' || (v_topic.topic_id % 64), + v_topic.topic_id::text); + END IF; + status := 0; + first_position := v_last + 1; + last_position := v_last + v_count; +END $$; + +-- Keeps the newest `p_keep` entries of each partition, so a log with no driver stays bounded. +CREATE FUNCTION cap_log_v1(p_keep int) RETURNS bigint +LANGUAGE plpgsql SET search_path FROM CURRENT AS $$ +DECLARE + v_partition int; + v_oldest bigint; + v_rows bigint; + deleted bigint := 0; +BEGIN + FOR v_partition IN 0..63 LOOP + SELECT n.id INTO v_oldest FROM notifications n WHERE n.partition = v_partition + ORDER BY n.id DESC OFFSET p_keep LIMIT 1; + IF FOUND THEN + DELETE FROM notifications n WHERE n.partition = v_partition AND n.id <= v_oldest; + GET DIAGNOSTICS v_rows = ROW_COUNT; + deleted := deleted + v_rows; + END IF; + END LOOP; + RETURN deleted; +END $$; diff --git a/crates/streams/src/postgres/mod.rs b/crates/streams/src/postgres/mod.rs index 91f20a2eb0..f3358e2044 100644 --- a/crates/streams/src/postgres/mod.rs +++ b/crates/streams/src/postgres/mod.rs @@ -3,11 +3,13 @@ mod conn; mod errors; mod janitor; +mod listen; +mod log; mod migrations; mod wait; use crate::{ - StreamError, StreamResult, StreamStore, + NotificationLog, StreamError, StreamResult, StreamStore, proto::{ ChainId, DeleteOwnerRequest, DeleteOwnerResponse, StoreAppendRequest, StoreAppendResponse, StoreLatestRequest, StoreLatestResponse, StoreReadRequest, StoreReadResponse, StoredClose, @@ -35,6 +37,13 @@ const DELETE_BATCH: i32 = 10_000; const PARTITIONS: u32 = 64; /// The oldest PostgreSQL the store runs on, as `server_version_num`. const OLDEST_SERVER: i32 = 140_000; +/// How many entries each log partition keeps at most, as the Redis log's `MAXLEN ~ 10000`. +const LOG_CAP: i32 = 10_000; +/// The longest schema name. A channel name, `_read_63`, takes at most 63 bytes. +const MAX_SCHEMA: usize = 55; + +/// How many doorbells a notification log rings when the options don't say. +pub const DEFAULT_POSTGRES_DOORBELLS: u32 = 16; /// How to reach PostgreSQL and how long streams live there. #[derive(Debug, Clone, PartialEq, Eq)] @@ -62,6 +71,17 @@ pub struct PostgresStoreOptions { pub skip_migrations: bool, /// How often the janitor trims records past the retention and deletes old tombstones. pub janitor_interval: Duration, + /// Each append also writes an entry to the notification log in its transaction, which a + /// driver reads. + pub notification_log: bool, + /// How many doorbells the notification log rings, from 1 to 64. + pub notify_doorbells: u32, + /// A URL for the connection that listens and the one that applies migrations, when `url` + /// goes through a pooler in transaction mode. `None` uses `url`. + pub listen_url: Option, + /// Each append also sends `NOTIFY`, so waiting reads in other processes wake in milliseconds. + /// Costs append throughput when many producers write at once. + pub notify_reads: bool, } impl PostgresStoreOptions { @@ -78,6 +98,10 @@ impl PostgresStoreOptions { poll_interval: Duration::from_millis(100), skip_migrations: false, janitor_interval: Duration::from_secs(60), + notification_log: false, + notify_doorbells: DEFAULT_POSTGRES_DOORBELLS, + listen_url: None, + notify_reads: false, } } } @@ -85,6 +109,7 @@ impl PostgresStoreOptions { /// A [StreamStore] on PostgreSQL. pub struct PostgresStore { inner: Arc, + log: Option>, tasks: Vec>, } @@ -176,6 +201,10 @@ struct Sql { due: String, trim: String, idle_chains: String, + take: String, + ack: String, + busy: String, + cap_log: String, } /// The chain key, as `$1` to `$5`. @@ -243,6 +272,25 @@ impl Sql { ), trim: format!("SELECT deleted, done FROM {s}.trim_topic_v1($1, $2, $3, $4)"), idle_chains: format!("SELECT {s}.drop_idle_chains_v1($1, $2, $3)"), + take: format!( + "WITH picked AS (SELECT n.partition, n.id FROM {s}.notifications n \ + WHERE n.partition = $1 AND (n.taken_by IS NULL OR n.taken_by = $2 \ + OR n.taken_at < now() - $3::bigint * interval '1 millisecond') \ + ORDER BY n.id LIMIT $4 FOR UPDATE SKIP LOCKED), \ + taken AS (UPDATE {s}.notifications n SET taken_by = $2, taken_at = now() \ + FROM picked WHERE n.partition = picked.partition AND n.id = picked.id \ + RETURNING n.id, n.topic_id, n.position) \ + SELECT taken.id, c.namespace, c.owner_kind, c.owner_id, c.first_run_id, \ + c.activity_id, t.topic, taken.position \ + FROM taken LEFT JOIN {s}.topics t ON t.topic_id = taken.topic_id \ + LEFT JOIN {s}.chains c ON c.chain_id = t.chain_id ORDER BY taken.id" + ), + ack: format!("DELETE FROM {s}.notifications WHERE partition = $1 AND id = ANY ($2)"), + busy: format!( + "SELECT p FROM generate_series(0, 63) p \ + WHERE EXISTS (SELECT 1 FROM {s}.notifications n WHERE n.partition = p)" + ), + cap_log: format!("SELECT {s}.cap_log_v1($1)"), } } } @@ -259,6 +307,10 @@ struct Inner { /// statement timeout, so the server's own timeout comes first. call_timeout: Duration, waiters: Waiters, + /// Whether appends write the notification log. + log: bool, + doorbells: i32, + notify_reads: bool, } impl PostgresStore { @@ -270,10 +322,17 @@ impl PostgresStore { pub async fn connect(options: PostgresStoreOptions) -> StreamResult { if !valid_schema(&options.schema) { return Err(StreamError::refused(format!( - "the schema must match [a-z_][a-z0-9_]{{0,62}}, not {:?}", + "the schema must match [a-z_][a-z0-9_]{{0,{}}}, not {:?}", + MAX_SCHEMA - 1, options.schema ))); } + if !(1..=PARTITIONS).contains(&options.notify_doorbells) { + return Err(StreamError::refused(format!( + "the notification log rings 1 to {PARTITIONS} doorbells, not {}", + options.notify_doorbells + ))); + } let retention_ms = i64::try_from(options.retention.as_millis()).unwrap_or(i64::MAX); if retention_ms == 0 { return Err(StreamError::refused(format!( @@ -282,11 +341,31 @@ impl PostgresStore { ))); } let target = Target::from_url(&options.url)?; - let mut direct = target.connect().await.map_err(unreachable)?; + // Migrations take a session lock, so they run on the direct connection when there is one. + let direct_target = match &options.listen_url { + Some(url) => Target::from_url(url)?, + None => target.clone(), + }; + let mut direct = direct_target.connect().await.map_err(unreachable)?; check_server(&direct).await?; migrations::run(&mut direct, &options.schema, !options.skip_migrations).await?; drop(direct); let pool = target.pool(options.max_connections, options.response_timeout)?; + let rings = options + .notification_log + .then(|| Arc::new(listen::Rings::new(options.notify_doorbells))); + let listener = (options.notification_log || options.notify_reads).then(|| { + listen::Listener::new( + direct_target, + &options.schema, + rings.clone(), + options.notify_reads, + ) + }); + let listening = match &listener { + Some(listener) => Some(listener.connect().await?), + None => None, + }; let inner = Arc::new(Inner { pool, sql: Sql::new(&options.schema), @@ -299,15 +378,32 @@ impl PostgresStore { floor_release_grace_ms: millis(options.floor_release_grace), call_timeout: options.response_timeout + Duration::from_secs(2), waiters: Waiters::default(), + log: options.notification_log, + doorbells: options.notify_doorbells as i32, + notify_reads: options.notify_reads, }); - let tasks = vec![ + let mut tasks = vec![ tokio::spawn(poll(Arc::downgrade(&inner), options.poll_interval)), tokio::spawn(janitor::run( Arc::downgrade(&inner), options.janitor_interval, )), ]; - Ok(Self { inner, tasks }) + if let (Some(listener), Some((client, notifications))) = (listener, listening) { + tasks.push(tokio::spawn(listener.run( + Arc::downgrade(&inner), + client, + notifications, + ))); + } + let log = rings.map(|rings| { + Arc::new(log::PostgresLog { + inner: inner.clone(), + rings, + doorbells: options.notify_doorbells, + }) + }); + Ok(Self { inner, log, tasks }) } /// Runs one janitor pass now, as the store's janitor task does every `janitor_interval`. @@ -322,7 +418,7 @@ fn valid_schema(schema: &str) -> bool { bytes .next() .is_some_and(|first| first.is_ascii_lowercase() || first == b'_') - && schema.len() <= 63 + && schema.len() <= MAX_SCHEMA && bytes.all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'_') } @@ -384,6 +480,7 @@ fn wait_of(request: &StoreReadRequest) -> Duration { struct Look { response: StoreReadResponse, seen: Seen, + topic_id: Option, } impl Inner { @@ -461,6 +558,7 @@ impl Inner { close: close.clone().filter(|_| topic_closed), }, seen, + topic_id, }; let (Some(_), Some(last)) = (topic_id, last) else { // A topic that never held a record. A closed one answers empty and closed, as a @@ -578,6 +676,10 @@ impl StreamStore for PostgresStore { NAME } + fn notification_log(&self) -> Option> { + self.log.clone().map(|log| log as Arc) + } + async fn append(&self, request: StoreAppendRequest) -> StreamResult { if request.records.is_empty() { return Err(StreamError::refused("an append needs at least one record")); @@ -606,11 +708,11 @@ impl StreamStore for PostgresStore { &request.digest, &request.records, &chain.keeps_topics(), - &false, - &0i32, + &inner.log, + &inner.doorbells, &no_lease, &inner.retention_ms, - &false, + &inner.notify_reads, ], ) .await?; @@ -674,9 +776,11 @@ impl StreamStore for PostgresStore { if wait.is_zero() || !look.response.records.is_empty() || look.response.closed { return Ok(look.response); } - let mut wake = inner - .waiters - .register((chain.clone(), request.topic.clone()), look.seen); + let mut wake = inner.waiters.register( + (chain.clone(), request.topic.clone()), + look.seen, + look.topic_id, + ); loop { look = inner.look(&chain, &request).await?; if !look.response.records.is_empty() || look.response.closed { @@ -1019,8 +1123,8 @@ mod tests { ] { assert_eq!(valid_schema(schema), ok, "{schema:?}"); } - assert!(valid_schema(&"a".repeat(63))); - assert!(!valid_schema(&"a".repeat(64))); + assert!(valid_schema(&"a".repeat(55))); + assert!(!valid_schema(&"a".repeat(56))); } #[test] diff --git a/crates/streams/src/postgres/wait.rs b/crates/streams/src/postgres/wait.rs index c36d532a3c..8ab01009b4 100644 --- a/crates/streams/src/postgres/wait.rs +++ b/crates/streams/src/postgres/wait.rs @@ -1,7 +1,7 @@ //! Waiting reads. A waiting read holds no connection. It waits for a wake from an append or a -//! close through this store, or from one batched poll per process for every stream with a -//! waiter. A wake only makes the read look again, so a lost or a spurious wake costs latency and -//! never a record. +//! close through this store, from one batched poll per process for every stream with a waiter, +//! or from a `NOTIFY` with `notify_reads`. A wake only makes the read look again, so a lost or a +//! spurious wake costs latency and never a record. use super::Chain; use std::{collections::HashMap, sync::Mutex}; @@ -23,22 +23,27 @@ struct Waiting { wake: watch::Sender, /// The state the waiters were woken for last, which the poll compares with. seen: Seen, + /// The topic's row, which a `NOTIFY` names. `None` before the topic has one. + topic_id: Option, } impl Waiters { /// Registers a waiter that saw `seen`. Register before the read looks once more, so a write /// between the two wakes it. - pub(super) fn register(&self, key: StreamKey, seen: Seen) -> watch::Receiver { - self.streams - .lock() - .unwrap() - .entry(key) - .or_insert_with(|| Waiting { - wake: watch::Sender::new(0), - seen, - }) - .wake - .subscribe() + pub(super) fn register( + &self, + key: StreamKey, + seen: Seen, + topic_id: Option, + ) -> watch::Receiver { + let mut streams = self.streams.lock().unwrap(); + let waiting = streams.entry(key).or_insert_with(|| Waiting { + wake: watch::Sender::new(0), + seen, + topic_id, + }); + waiting.topic_id = waiting.topic_id.or(topic_id); + waiting.wake.subscribe() } /// Wakes the waiters of one stream. @@ -57,6 +62,23 @@ impl Waiters { } } + /// Wakes the waiters of the topic a `NOTIFY` named, and those of topics with no row yet, + /// which a `NOTIFY` can't name. + pub(super) fn wake_topic(&self, topic_id: i64) { + for waiting in self.streams.lock().unwrap().values() { + if waiting.topic_id.is_none_or(|id| id == topic_id) { + waiting.wake.send_modify(|n| *n += 1); + } + } + } + + /// Wakes every waiter, for when wakes may have been lost. + pub(super) fn wake_all(&self) { + for waiting in self.streams.lock().unwrap().values() { + waiting.wake.send_modify(|n| *n += 1); + } + } + /// The streams with a waiter. Streams whose waiters all went are dropped here. pub(super) fn streams(&self) -> Vec { let mut streams = self.streams.lock().unwrap(); @@ -105,7 +127,7 @@ mod tests { #[test] fn a_poll_wakes_only_on_a_change() { let waiters = Waiters::default(); - let mut rx = waiters.register(key("out"), (Some(1), false)); + let mut rx = waiters.register(key("out"), (Some(1), false), Some(7)); waiters.polled(&key("out"), (Some(1), false)); assert!(!rx.has_changed().unwrap()); waiters.polled(&key("out"), (Some(2), false)); @@ -118,21 +140,33 @@ mod tests { #[test] fn a_chain_wake_reaches_its_topics_only() { let waiters = Waiters::default(); - let out = waiters.register(key("out"), (None, false)); - let other = waiters.register(key("other"), (None, false)); + let out = waiters.register(key("out"), (None, false), None); + let other = waiters.register(key("other"), (None, false), None); let mut elsewhere = key("out"); elsewhere.0.owner = "wf2".to_string(); - let far = waiters.register(elsewhere, (None, false)); + let far = waiters.register(elsewhere, (None, false), None); waiters.wake_chain(&key("out").0); assert!(out.has_changed().unwrap()); assert!(other.has_changed().unwrap()); assert!(!far.has_changed().unwrap()); } + #[test] + fn a_notify_wakes_its_topic_and_topics_with_no_row() { + let waiters = Waiters::default(); + let seven = waiters.register(key("seven"), (Some(1), false), Some(7)); + let eight = waiters.register(key("eight"), (Some(1), false), Some(8)); + let new = waiters.register(key("new"), (None, false), None); + waiters.wake_topic(7); + assert!(seven.has_changed().unwrap()); + assert!(!eight.has_changed().unwrap()); + assert!(new.has_changed().unwrap()); + } + #[test] fn a_dropped_waiter_leaves_no_stream() { let waiters = Waiters::default(); - let rx = waiters.register(key("out"), (None, false)); + let rx = waiters.register(key("out"), (None, false), None); assert_eq!(waiters.streams().len(), 1); assert_eq!(waiters.len(), 1); drop(rx); diff --git a/crates/streams/src/service.rs b/crates/streams/src/service.rs index 8b36173674..39e1845844 100644 --- a/crates/streams/src/service.rs +++ b/crates/streams/src/service.rs @@ -438,10 +438,16 @@ async fn connect_redis( #[cfg(feature = "postgres")] async fn connect_postgres( config: crate::proto::PostgresStoreConfig, - _notification_log: bool, + notification_log: bool, _archive_lease: Option, ) -> StreamResult<(Arc, Duration)> { let mut options = crate::PostgresStoreOptions::new(config.url); + options.notification_log = notification_log; + if config.notify_doorbells > 0 { + options.notify_doorbells = config.notify_doorbells; + } + options.listen_url = Some(config.listen_url).filter(|url| !url.is_empty()); + options.notify_reads = config.notify_reads; if !config.schema.is_empty() { options.schema = config.schema; } diff --git a/crates/streams/tests/conformance.rs b/crates/streams/tests/conformance.rs index 1fecda04df..75363a5c26 100644 --- a/crates/streams/tests/conformance.rs +++ b/crates/streams/tests/conformance.rs @@ -796,22 +796,61 @@ mod on_postgres { } pub(crate) async fn postgres(variable: &str) -> Option { + postgres_with(variable, |_| {}).await + } + + /// [postgres] with a notification log. + pub(crate) async fn postgres_with_log(variable: &str) -> Option { + postgres_with(variable, |options| { + options.notification_log = true; + options.notify_doorbells = 4; + }) + .await + } + + /// [postgres] with a `NOTIFY` per append for waiting reads. + pub(crate) async fn postgres_notifying_reads(variable: &str) -> Option { + postgres_with(variable, |options| options.notify_reads = true).await + } + + pub(crate) async fn postgres_with( + variable: &str, + change: impl FnOnce(&mut PostgresStoreOptions), + ) -> Option { let Ok(url) = std::env::var(variable) else { eprintln!("set {variable} to run the PostgreSQL store cases"); return None; }; let schema = schema(); + let mut options = options(&url, &schema); + change(&mut options); Some(Postgres { - store: Arc::new( - PostgresStore::connect(options(&url, &schema)) - .await - .unwrap(), - ), + store: Arc::new(PostgresStore::connect(options).await.unwrap()), raw: Arc::new(raw(&url).await), schema, }) } + /// The store through PgBouncer in transaction mode at `STREAMS_POSTGRES_PGBOUNCER_URL`, with + /// the log, and the listener and the migrations on the direct `STREAMS_POSTGRES_URL`. + pub(crate) async fn postgres_through_pgbouncer() -> Option { + let Ok(pooled) = std::env::var("STREAMS_POSTGRES_PGBOUNCER_URL") else { + eprintln!("set STREAMS_POSTGRES_PGBOUNCER_URL to run the PgBouncer cases"); + return None; + }; + let direct = std::env::var("STREAMS_POSTGRES_URL").unwrap(); + let schema = schema(); + let mut options = options(&pooled, &schema); + options.listen_url = Some(direct.clone()); + options.notification_log = true; + options.notify_doorbells = 4; + Some(Postgres { + store: Arc::new(PostgresStore::connect(options).await.unwrap()), + raw: Arc::new(raw(&direct).await), + schema, + }) + } + /// The role lines of the PostgreSQL page: one `CREATE ROLE` and the `GRANT`s after it. pub(crate) fn documented_lines() -> Vec { let page = include_str!("../docs/postgres.md"); @@ -834,6 +873,15 @@ mod on_postgres { /// A store as a role with exactly the documented grants, on a schema an owner migrated. The /// page's schema, role and password are swapped for the test's own, and nothing else changes. pub(crate) async fn postgres_as_documented_role(variable: &str) -> Option { + postgres_as_documented_role_with(variable, false).await + } + + /// [postgres_as_documented_role] with a notification log. + pub(crate) async fn postgres_as_documented_role_with_log(variable: &str) -> Option { + postgres_as_documented_role_with(variable, true).await + } + + async fn postgres_as_documented_role_with(variable: &str, log: bool) -> Option { let Ok(url) = std::env::var(variable) else { eprintln!("set {variable} to run the PostgreSQL store cases"); return None; @@ -858,6 +906,7 @@ mod on_postgres { let role_url = format!("postgres://{role}:{password}@{address}"); let mut options = options(&role_url, &schema); options.skip_migrations = true; + options.notification_log = log; Some(Postgres { store: Arc::new(PostgresStore::connect(options).await.unwrap()), raw: Arc::new(raw), @@ -3030,11 +3079,38 @@ mod postgres { ); } +#[cfg(feature = "postgres")] +mod postgres_with_log { + store_cases!(super::on_postgres::postgres_with_log( + "STREAMS_POSTGRES_URL" + )); + log_cases!(super::on_postgres::postgres_with_log( + "STREAMS_POSTGRES_URL" + )); +} + +#[cfg(feature = "postgres")] +mod postgres_notifying_reads { + store_cases!(super::on_postgres::postgres_notifying_reads( + "STREAMS_POSTGRES_URL" + )); +} + +#[cfg(feature = "postgres")] +mod pgbouncer_postgres { + // Transaction mode, with the listener and the migrations on a direct connection. + store_cases!(super::on_postgres::postgres_through_pgbouncer()); + log_cases!(super::on_postgres::postgres_through_pgbouncer()); +} + #[cfg(feature = "postgres")] mod activity_postgres { store_cases!(super::activity(super::on_postgres::postgres( "STREAMS_POSTGRES_URL" ))); + log_cases!(super::activity(super::on_postgres::postgres_with_log( + "STREAMS_POSTGRES_URL" + ))); conformance!(super::activity(super::on_postgres::postgres("STREAMS_POSTGRES_URL")); deleting_an_owner_drops_every_chain_of_it, ); @@ -3045,11 +3121,24 @@ mod workflow_activity_postgres { store_cases!(super::workflow_activity(super::on_postgres::postgres( "STREAMS_POSTGRES_URL" ))); + log_cases!(super::workflow_activity( + super::on_postgres::postgres_with_log("STREAMS_POSTGRES_URL") + )); conformance!(super::workflow_activity(super::on_postgres::postgres("STREAMS_POSTGRES_URL")); deleting_a_workflow_drops_its_workflow_activities, ); } +#[cfg(feature = "postgres")] +mod documented_role_postgres_with_log { + store_cases!(super::on_postgres::postgres_as_documented_role_with_log( + "STREAMS_POSTGRES_URL" + )); + log_cases!(super::on_postgres::postgres_as_documented_role_with_log( + "STREAMS_POSTGRES_URL" + )); +} + #[cfg(feature = "postgres")] mod notifying_postgres { store_cases!(super::notifying(super::on_postgres::postgres( diff --git a/crates/streams/tests/driver.rs b/crates/streams/tests/driver.rs index 21c731f90c..6a82d9c79b 100644 --- a/crates/streams/tests/driver.rs +++ b/crates/streams/tests/driver.rs @@ -1,10 +1,11 @@ -//! The stream driver over a Redis notification log, on a standalone server and on a cluster. +//! The stream driver over a Redis notification log, on a standalone server and on a cluster, and +//! over a PostgreSQL one. //! -//! Set `STREAMS_REDIS_URL` and `STREAMS_REDIS_CLUSTER_URL` to run these. The cluster cases that -//! move slots and fail a primary over need `STREAMS_REDIS_CLUSTER_ADMIN`, since they change the -//! cluster, and replicas for the failover. +//! Set `STREAMS_REDIS_URL`, `STREAMS_REDIS_CLUSTER_URL` and `STREAMS_POSTGRES_URL` to run these. +//! The cluster cases that move slots and fail a primary over need `STREAMS_REDIS_CLUSTER_ADMIN`, +//! since they change the cluster, and replicas for the failover. -#![cfg(feature = "redis")] +#![cfg(any(feature = "redis", feature = "postgres"))] use std::{ collections::{BTreeMap, HashMap}, @@ -15,14 +16,17 @@ use std::{ time::Duration, }; use temporalio_streams::{ - NotifierClient, NotifyError, RedisStore, RedisStoreOptions, StreamDriver, StreamDriverOptions, - StreamNotification, StreamStore, chain_owner, mint_cursor, progress_counter, + NotifierClient, NotifyError, StreamDriver, StreamDriverOptions, StreamNotification, + StreamStore, chain_owner, mint_cursor, progress_counter, proto::{ChainId, StoreAppendRequest, StreamOwnerKind}, stream_hash, }; +#[cfg(feature = "redis")] +use temporalio_streams::{RedisStore, RedisStoreOptions}; static IDS: AtomicUsize = AtomicUsize::new(0); +#[cfg(feature = "redis")] /// The cases that move a slot or fail a primary over hold it alone, the others share it, since a /// cluster change makes appends fail for a moment. static CLUSTER: tokio::sync::RwLock<()> = tokio::sync::RwLock::const_new(()); @@ -105,6 +109,7 @@ impl Notifiers { } } +#[cfg(feature = "redis")] async fn store(url: &str, cluster: bool, prefix: &str) -> Arc { let mut options = RedisStoreOptions::new(url); options.cluster = cluster; @@ -127,7 +132,7 @@ fn options(consumer: &str) -> StreamDriverOptions { /// Appends one batch to `streams` chains, each in its own slot most likely, and answers each /// stream's newest counter. async fn write( - store: &RedisStore, + store: &dyn StreamStore, streams: usize, sequence: i64, ) -> BTreeMap<(String, String), i64> { @@ -162,7 +167,7 @@ async fn write( /// Appends one batch to the runs of `streams` standalone Activities, `act-`, and answers /// each Activity stream's newest counter and its cursor. async fn write_activities( - store: &RedisStore, + store: &dyn StreamStore, streams: usize, ) -> BTreeMap<(String, String), (i64, String)> { let mut expected = BTreeMap::new(); @@ -200,15 +205,12 @@ async fn write_activities( expected } -async fn activity_streams_reach_their_notifiers(url: &str, cluster: bool) { - let _shared = CLUSTER.read().await; - let prefix = unique("driver-activity"); - let store = store(url, cluster, &prefix).await; +async fn activity_streams_reach_their_notifiers(store: Arc) { let notifiers = Arc::new(Notifiers::default()); let driver = Arc::new(StreamDriver::new(store.clone(), notifiers.clone(), options("d1")).unwrap()); let running = tokio::spawn(driver.run()); - let expected = write_activities(&store, 100).await; + let expected = write_activities(&*store, 100).await; let counters = expected .iter() .map(|(stream, (counter, _))| (stream.clone(), *counter)) @@ -221,20 +223,26 @@ async fn activity_streams_reach_their_notifiers(url: &str, cluster: bool) { } } +#[cfg(feature = "redis")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn activity_streams_on_a_server_reach_their_notifiers() { let Some(url) = url("STREAMS_REDIS_URL") else { return; }; - activity_streams_reach_their_notifiers(&url, false).await; + let _shared = CLUSTER.read().await; + activity_streams_reach_their_notifiers(store(&url, false, &unique("driver-activity")).await) + .await; } +#[cfg(feature = "redis")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn activity_streams_in_a_cluster_reach_their_notifiers() { let Some(url) = url("STREAMS_REDIS_CLUSTER_URL") else { return; }; - activity_streams_reach_their_notifiers(&url, true).await; + let _shared = CLUSTER.read().await; + activity_streams_reach_their_notifiers(store(&url, true, &unique("driver-activity")).await) + .await; } fn url(variable: &str) -> Option { @@ -245,18 +253,15 @@ fn url(variable: &str) -> Option { url } -async fn streams_in_every_slot_reach_their_notifiers(url: &str, cluster: bool) { - let _shared = CLUSTER.read().await; - let prefix = unique("driver-spread"); - let store = store(url, cluster, &prefix).await; +async fn streams_in_every_slot_reach_their_notifiers(store: Arc) { let notifiers = Arc::new(Notifiers::default()); let driver = Arc::new(StreamDriver::new(store.clone(), notifiers.clone(), options("d1")).unwrap()); let running = tokio::spawn(driver.run()); - let expected = write(&store, 300, 1).await; + let expected = write(&*store, 300, 1).await; notifiers.reach(&expected, Duration::from_secs(20)).await; // A second batch moves every counter on. - let expected = write(&store, 300, 2).await; + let expected = write(&*store, 300, 2).await; notifiers.reach(&expected, Duration::from_secs(20)).await; running.abort(); let log = store.notification_log().unwrap(); @@ -266,22 +271,29 @@ async fn streams_in_every_slot_reach_their_notifiers(url: &str, cluster: bool) { ); } +#[cfg(feature = "redis")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn streams_on_a_server_reach_their_notifiers() { let Some(url) = url("STREAMS_REDIS_URL") else { return; }; - streams_in_every_slot_reach_their_notifiers(&url, false).await; + let _shared = CLUSTER.read().await; + streams_in_every_slot_reach_their_notifiers(store(&url, false, &unique("driver-spread")).await) + .await; } +#[cfg(feature = "redis")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn streams_in_every_slot_of_a_cluster_reach_their_notifiers() { let Some(url) = url("STREAMS_REDIS_CLUSTER_URL") else { return; }; - streams_in_every_slot_reach_their_notifiers(&url, true).await; + let _shared = CLUSTER.read().await; + streams_in_every_slot_reach_their_notifiers(store(&url, true, &unique("driver-spread")).await) + .await; } +#[cfg(feature = "redis")] /// The four doorbells of a store under `prefix`, which needs no percent-encoding. fn bell_keys(prefix: &str) -> Vec { let tags: Vec = @@ -292,6 +304,7 @@ fn bell_keys(prefix: &str) -> Vec { .collect() } +#[cfg(feature = "redis")] async fn delete(url: &str, cluster: bool, keys: &[String]) { for key in keys { if cluster { @@ -320,11 +333,12 @@ async fn delete(url: &str, cluster: bool, keys: &[String]) { } } +#[cfg(feature = "redis")] async fn a_lost_ring_is_found_by_the_sweep(url: &str, cluster: bool) { let _shared = CLUSTER.read().await; let prefix = unique("driver-sweep"); let store = store(url, cluster, &prefix).await; - let expected = write(&store, 20, 1).await; + let expected = write(&*store, 20, 1).await; // The rings are gone, as when Core stopped between the append and its ring. delete(url, cluster, &bell_keys(&prefix)).await; let notifiers = Arc::new(Notifiers::default()); @@ -336,6 +350,7 @@ async fn a_lost_ring_is_found_by_the_sweep(url: &str, cluster: bool) { running.abort(); } +#[cfg(feature = "redis")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn a_lost_ring_on_a_server_is_found_by_the_sweep() { let Some(url) = url("STREAMS_REDIS_URL") else { @@ -344,6 +359,7 @@ async fn a_lost_ring_on_a_server_is_found_by_the_sweep() { a_lost_ring_is_found_by_the_sweep(&url, false).await; } +#[cfg(feature = "redis")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn a_lost_ring_on_a_cluster_is_found_by_the_sweep() { let Some(url) = url("STREAMS_REDIS_CLUSTER_URL") else { @@ -352,21 +368,24 @@ async fn a_lost_ring_on_a_cluster_is_found_by_the_sweep() { a_lost_ring_is_found_by_the_sweep(&url, true).await; } +#[cfg(feature = "redis")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn a_stopped_drivers_entries_move_to_another_driver() { let Some(url) = url("STREAMS_REDIS_CLUSTER_URL") else { return; }; let _shared = CLUSTER.read().await; - let prefix = unique("driver-claim"); - let store = store(&url, true, &prefix).await; + a_stopped_drivers_entries_move_on(store(&url, true, &unique("driver-claim")).await).await; +} + +async fn a_stopped_drivers_entries_move_on(store: Arc) { let failing = Arc::new(Notifiers::default()); failing.fail.store(true, Ordering::SeqCst); let stopped = tokio::spawn( Arc::new(StreamDriver::new(store.clone(), failing.clone(), options("stopped")).unwrap()) .run(), ); - let expected = write(&store, 30, 1).await; + let expected = write(&*store, 30, 1).await; // The first driver takes the entries and can't send them, then stops. tokio::time::sleep(Duration::from_millis(500)).await; stopped.abort(); @@ -379,12 +398,14 @@ async fn a_stopped_drivers_entries_move_to_another_driver() { other.abort(); } +#[cfg(feature = "redis")] /// The cluster to move slots on and fail over, when the test may change it. fn admin_url() -> Option { let admin = std::env::var("STREAMS_REDIS_CLUSTER_ADMIN").ok()?; url("STREAMS_REDIS_CLUSTER_URL").filter(|_| admin == "1") } +#[cfg(feature = "redis")] /// One node of the cluster, as `CLUSTER NODES` lists it. #[derive(Debug, Clone)] struct Node { @@ -393,6 +414,7 @@ struct Node { primary_id: Option, } +#[cfg(feature = "redis")] async fn node_list(url: &str) -> Vec { let mut conn = redis::Client::open(url) .unwrap() @@ -417,6 +439,7 @@ async fn node_list(url: &str) -> Vec { .collect() } +#[cfg(feature = "redis")] async fn on(url: &str, cmd: &redis::Cmd) -> redis::Value { let mut conn = redis::Client::open(url) .unwrap() @@ -426,6 +449,7 @@ async fn on(url: &str, cmd: &redis::Cmd) -> redis::Value { cmd.query_async(&mut conn).await.unwrap() } +#[cfg(feature = "redis")] /// A key's Cluster slot, as Redis computes it from the hash tag. fn keyslot(key: &str) -> u16 { let tag = key @@ -435,6 +459,7 @@ fn keyslot(key: &str) -> u16 { crc16::State::::calculate(tag.as_bytes()) % 16384 } +#[cfg(feature = "redis")] /// The primary that serves `slot`. async fn owner_of(url: &str, slot: u16) -> Node { let mut conn = redis::Client::open(url) @@ -480,6 +505,7 @@ async fn owner_of(url: &str, slot: u16) -> Node { .unwrap() } +#[cfg(feature = "redis")] /// Workflow ids whose chains land in `slot`, found by asking the cluster. fn ids_in_slot(prefix: &str, slot: u16, count: usize) -> Vec { let mut found = vec![]; @@ -494,6 +520,7 @@ fn ids_in_slot(prefix: &str, slot: u16, count: usize) -> Vec { found } +#[cfg(feature = "redis")] /// Appends to each of `ids` again and again until `stop`, retrying a batch that a slot move /// refused, and answers each stream's newest counter. async fn traffic( @@ -537,6 +564,7 @@ async fn traffic( expected } +#[cfg(feature = "redis")] /// Moves `slot` from its primary to another one, key by key, as `redis-cli --cluster reshard` /// does. async fn move_slot(url: &str, slot: u16) { @@ -613,6 +641,7 @@ async fn move_slot(url: &str, slot: u16) { } } +#[cfg(feature = "redis")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn a_slot_moved_under_traffic_loses_no_notification() { let Some(url) = admin_url() else { return }; @@ -636,6 +665,7 @@ async fn a_slot_moved_under_traffic_loses_no_notification() { running.abort(); } +#[cfg(feature = "redis")] #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn a_primary_failover_under_traffic_still_notifies_every_stream() { let Some(url) = admin_url() else { return }; @@ -678,3 +708,66 @@ async fn a_primary_failover_under_traffic_still_notifies_every_stream() { notifiers.reach(&expected, Duration::from_secs(30)).await; running.abort(); } + +#[cfg(feature = "postgres")] +mod on_postgres { + use super::*; + use temporalio_streams::{PostgresStore, PostgresStoreOptions}; + + /// A schema no other case uses. + fn schema() -> String { + unique("driver").replace('-', "_") + } + + async fn store(url: &str, schema: &str) -> Arc { + let mut options = PostgresStoreOptions::new(url); + options.schema = schema.to_string(); + options.max_connections = 8; + options.notification_log = true; + options.notify_doorbells = 4; + Arc::new(PostgresStore::connect(options).await.unwrap()) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn streams_on_postgres_reach_their_notifiers() { + let Some(url) = url("STREAMS_POSTGRES_URL") else { + return; + }; + streams_in_every_slot_reach_their_notifiers(store(&url, &schema()).await).await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn activity_streams_on_postgres_reach_their_notifiers() { + let Some(url) = url("STREAMS_POSTGRES_URL") else { + return; + }; + activity_streams_reach_their_notifiers(store(&url, &schema()).await).await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn a_lost_ring_on_postgres_is_found_by_the_sweep() { + let Some(url) = url("STREAMS_POSTGRES_URL") else { + return; + }; + let schema = schema(); + // The rings go to the writer's listener, as when the driver's process was away. + let writer = store(&url, &schema).await; + let expected = write(&*writer, 20, 1).await; + let reader = store(&url, &schema).await; + let notifiers = Arc::new(Notifiers::default()); + let driver = + Arc::new(StreamDriver::new(reader.clone(), notifiers.clone(), options("d1")).unwrap()); + let running = tokio::spawn(driver.run()); + // One sweep interval and the time one sweep takes. + notifiers.reach(&expected, Duration::from_secs(5)).await; + running.abort(); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn a_stopped_drivers_entries_on_postgres_move_to_another_driver() { + let Some(url) = url("STREAMS_POSTGRES_URL") else { + return; + }; + a_stopped_drivers_entries_move_on(store(&url, &schema()).await).await; + } +} diff --git a/crates/streams/tests/postgres_store.rs b/crates/streams/tests/postgres_store.rs index f4b4325106..6dbe92fa83 100644 --- a/crates/streams/tests/postgres_store.rs +++ b/crates/streams/tests/postgres_store.rs @@ -534,3 +534,101 @@ async fn tls_with_require_reaches_a_server_with_tls() { let error = PostgresStore::connect(options).await.unwrap_err(); assert_eq!(error.kind, StreamFailureKind::Storage, "{error}"); } + +/// Waits for one record on `chain` from `reader` in a task of its own, and answers how long the +/// wait took. +fn wait_for_one(reader: Arc, chain: &ChainId) -> tokio::task::JoinHandle { + let chain = chain.clone(); + tokio::spawn(async move { + let started = Instant::now(); + let page = reader + .read(read(&chain, "out", "", Duration::from_secs(10))) + .await + .unwrap(); + assert_eq!(page.records.len(), 1); + started.elapsed() + }) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn notify_reads_wakes_a_reader_in_another_process_at_once() { + let Some(url) = url() else { return }; + let mut options = options(&url); + options.notify_reads = true; + // The poll can't be what wakes the reader. + options.poll_interval = Duration::from_secs(3600); + let reader = Arc::new(PostgresStore::connect(options.clone()).await.unwrap()); + let writer = PostgresStore::connect(options).await.unwrap(); + let chain = chain("notified"); + // The topic exists first, so the waiter knows the row a `NOTIFY` names. + writer + .append(append(&chain, "other", "p", 1, 1)) + .await + .unwrap(); + writer + .append(append(&chain, "out", "q", 1, 1)) + .await + .unwrap(); + let after = writer + .read(read(&chain, "out", "", Duration::ZERO)) + .await + .unwrap() + .records[0] + .position + .clone(); + let waiting = { + let (reader, chain, after) = (reader.clone(), chain.clone(), after.clone()); + tokio::spawn(async move { + let page = reader + .read(read(&chain, "out", &after, Duration::from_secs(10))) + .await + .unwrap(); + (page.records.len(), Instant::now()) + }) + }; + tokio::time::sleep(Duration::from_millis(300)).await; + let sent = Instant::now(); + writer + .append(append(&chain, "out", "q", 2, 1)) + .await + .unwrap(); + let (records, woke) = waiting.await.unwrap(); + assert_eq!(records, 1); + let took = woke - sent; + assert!(took < Duration::from_millis(50), "{took:?}"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_listener_that_lost_its_connection_wakes_every_waiter() { + let Some(url) = url() else { return }; + let mut options = options(&url); + options.notify_reads = true; + options.poll_interval = Duration::from_secs(3600); + let schema = options.schema.clone(); + let reader = Arc::new(PostgresStore::connect(options.clone()).await.unwrap()); + let mut quiet = options.clone(); + quiet.notify_reads = false; + let writer = PostgresStore::connect(quiet).await.unwrap(); + let chain = chain("reconnect"); + let waiting = wait_for_one(reader.clone(), &chain); + tokio::time::sleep(Duration::from_millis(300)).await; + // The writer sends no `NOTIFY` and the poll is an hour away, so only the listener's + // reconnect can wake the waiter. + let admin = raw(&url).await; + let killed: i64 = admin + .query_one( + "SELECT count(pg_terminate_backend(pid)) FROM pg_stat_activity \ + WHERE application_name = 'temporal-streams' AND query LIKE $1", + &[&format!("LISTEN {schema}_%")], + ) + .await + .unwrap() + .get(0); + assert_eq!(killed, 1); + writer + .append(append(&chain, "out", "p", 1, 1)) + .await + .unwrap(); + let waited = waiting.await.unwrap(); + assert!(waited < Duration::from_secs(5), "{waited:?}"); +} From 59b8725c5916e9789f3567eb12040c76562a8787 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 15 Oct 2026 19:51:49 -0700 Subject: [PATCH 2/7] Dropped a pooled PostgreSQL connection that a fatal error ended. The pool's fast recycle could hand a terminated connection out again, and the next call then failed with a closed connection. --- crates/streams/src/postgres/mod.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/crates/streams/src/postgres/mod.rs b/crates/streams/src/postgres/mod.rs index f3358e2044..efb1891fee 100644 --- a/crates/streams/src/postgres/mod.rs +++ b/crates/streams/src/postgres/mod.rs @@ -505,7 +505,11 @@ impl Inner { Ok(Ok(rows)) => return Ok(rows), Ok(Err(error)) if retryable(&error) && tries < 3 => tries += 1, Ok(Err(error)) => { - if error.as_db_error().is_none() { + // A connection the error ended never goes back to the pool. + let fatal = error + .as_db_error() + .is_none_or(|db| matches!(db.severity(), "FATAL" | "PANIC")); + if fatal || client.is_closed() { drop(deadpool_postgres::Object::take(client)); } return Err(mapped(call, &error)); From 09bf3b51e598d4141fb3889488e46ce650f4a20a Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 15 Oct 2026 19:52:11 -0700 Subject: [PATCH 3/7] Counted only TLS connections in the PostgreSQL TLS test. Other tests in the process connect to the same server without TLS. --- crates/streams/tests/postgres_store.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/streams/tests/postgres_store.rs b/crates/streams/tests/postgres_store.rs index 6dbe92fa83..de439f697b 100644 --- a/crates/streams/tests/postgres_store.rs +++ b/crates/streams/tests/postgres_store.rs @@ -521,8 +521,8 @@ async fn tls_with_require_reaches_a_server_with_tls() { let raw = raw(&url).await; let encrypted: bool = raw .query_one( - "SELECT bool_and(s.ssl) FROM pg_stat_ssl s JOIN pg_stat_activity a USING (pid) \ - WHERE a.application_name = 'temporal-streams'", + "SELECT count(*) FILTER (WHERE s.ssl) > 0 FROM pg_stat_ssl s \ + JOIN pg_stat_activity a USING (pid) WHERE a.application_name = 'temporal-streams'", &[], ) .await From 9095044d8a528c23686005aec23afe7d469ac985 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 15 Oct 2026 20:54:39 -0700 Subject: [PATCH 4/7] Warmed the reader's pool before the timed notify_reads wake. Under load the reader's first connection could take longer than the 50 ms the test allows for the wake. --- crates/streams/tests/postgres_store.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/streams/tests/postgres_store.rs b/crates/streams/tests/postgres_store.rs index de439f697b..062c526e9a 100644 --- a/crates/streams/tests/postgres_store.rs +++ b/crates/streams/tests/postgres_store.rs @@ -569,7 +569,8 @@ async fn notify_reads_wakes_a_reader_in_another_process_at_once() { .append(append(&chain, "out", "q", 1, 1)) .await .unwrap(); - let after = writer + // The reader reads first, so its pool holds a connection before the timed part. + let after = reader .read(read(&chain, "out", "", Duration::ZERO)) .await .unwrap() From f89bf9e1e1929bc70a0ed867dfd91fc70597daaa Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Fri, 16 Oct 2026 05:47:02 -0700 Subject: [PATCH 5/7] Woke on a `NOTIFY` only the waiters of the topic it names. Waiters of a topic with no row woke on every `NOTIFY`, about a read per waiter per append. They learn the row from their reads, and the poll wakes them before that. --- crates/streams/docs/postgres.md | 3 ++- crates/streams/src/postgres/mod.rs | 3 +++ crates/streams/src/postgres/wait.rs | 41 ++++++++++++++++++----------- 3 files changed, 31 insertions(+), 16 deletions(-) diff --git a/crates/streams/docs/postgres.md b/crates/streams/docs/postgres.md index 2d54e72ecc..96248707e0 100644 --- a/crates/streams/docs/postgres.md +++ b/crates/streams/docs/postgres.md @@ -43,7 +43,8 @@ path, and this wait is the path for a reader without it. **`notify_reads`.** With it, each append also sends `NOTIFY` on `_read_` with the topic id, and the store's listener wakes the waiters of that topic. A reader in another -process then wakes in a few milliseconds. It is off by default, since a transaction that sent +process then wakes in a few milliseconds. A reader that waits on a topic with no row yet wakes +from the poll, since a `NOTIFY` can't name it. It is off by default, since a transaction that sent `NOTIFY` takes a lock on the server's notification queue as it commits, so such commits run one at a time across the database. With many producers at once, it costs append throughput. diff --git a/crates/streams/src/postgres/mod.rs b/crates/streams/src/postgres/mod.rs index 8837564266..aa54def317 100644 --- a/crates/streams/src/postgres/mod.rs +++ b/crates/streams/src/postgres/mod.rs @@ -876,6 +876,9 @@ impl StreamStore for PostgresStore { if !look.response.records.is_empty() || look.response.closed { return Ok(look.response); } + inner + .waiters + .learn(&(chain.clone(), request.topic.clone()), look.topic_id); if tokio::time::timeout_at(deadline, wake.changed()) .await .is_err() diff --git a/crates/streams/src/postgres/wait.rs b/crates/streams/src/postgres/wait.rs index 8ab01009b4..2108fa42d8 100644 --- a/crates/streams/src/postgres/wait.rs +++ b/crates/streams/src/postgres/wait.rs @@ -62,16 +62,23 @@ impl Waiters { } } - /// Wakes the waiters of the topic a `NOTIFY` named, and those of topics with no row yet, - /// which a `NOTIFY` can't name. + /// Wakes the waiters of the topic a `NOTIFY` named. Waiters of a topic with no row yet are + /// left to the poll, since one `NOTIFY` per append would wake them all. pub(super) fn wake_topic(&self, topic_id: i64) { for waiting in self.streams.lock().unwrap().values() { - if waiting.topic_id.is_none_or(|id| id == topic_id) { + if waiting.topic_id == Some(topic_id) { waiting.wake.send_modify(|n| *n += 1); } } } + /// Notes the row a read found for a stream whose waiters had none. + pub(super) fn learn(&self, key: &StreamKey, topic_id: Option) { + if let Some(waiting) = self.streams.lock().unwrap().get_mut(key) { + waiting.topic_id = waiting.topic_id.or(topic_id); + } + } + /// Wakes every waiter, for when wakes may have been lost. pub(super) fn wake_all(&self) { for waiting in self.streams.lock().unwrap().values() { @@ -124,6 +131,22 @@ mod tests { ) } + #[test] + fn a_notify_wakes_only_the_topic_it_names() { + let waiters = Waiters::default(); + let with_row = waiters.register(key("a"), (Some(1), false), Some(7)); + let no_row = waiters.register(key("b"), (None, false), None); + waiters.wake_topic(9); + assert!(!with_row.has_changed().unwrap()); + assert!(!no_row.has_changed().unwrap()); + waiters.wake_topic(7); + assert!(with_row.has_changed().unwrap()); + assert!(!no_row.has_changed().unwrap()); + waiters.learn(&key("b"), Some(8)); + waiters.wake_topic(8); + assert!(no_row.has_changed().unwrap()); + } + #[test] fn a_poll_wakes_only_on_a_change() { let waiters = Waiters::default(); @@ -151,18 +174,6 @@ mod tests { assert!(!far.has_changed().unwrap()); } - #[test] - fn a_notify_wakes_its_topic_and_topics_with_no_row() { - let waiters = Waiters::default(); - let seven = waiters.register(key("seven"), (Some(1), false), Some(7)); - let eight = waiters.register(key("eight"), (Some(1), false), Some(8)); - let new = waiters.register(key("new"), (None, false), None); - waiters.wake_topic(7); - assert!(seven.has_changed().unwrap()); - assert!(!eight.has_changed().unwrap()); - assert!(new.has_changed().unwrap()); - } - #[test] fn a_dropped_waiter_leaves_no_stream() { let waiters = Waiters::default(); From 9ec35532ab2557029353b3a8037734e0b49942bd Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Fri, 16 Oct 2026 05:47:02 -0700 Subject: [PATCH 6/7] Set TCP keepalive on the PostgreSQL listener. A half-open listener connection showed only after the system's two-hour keepalive, and its wakes stopped meanwhile. --- crates/streams/docs/postgres.md | 5 +++-- crates/streams/src/postgres/conn.rs | 31 +++++++++++++++++++++++++++-- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/crates/streams/docs/postgres.md b/crates/streams/docs/postgres.md index 96248707e0..97ff4a9aad 100644 --- a/crates/streams/docs/postgres.md +++ b/crates/streams/docs/postgres.md @@ -63,8 +63,9 @@ delete removed is acked as broken when a driver takes it. channels. It connects with the store, so a store whose listener can't connect fails at once. When the connection ends, it connects again with backoff, rings every doorbell for every partition and wakes every waiting read, since what was sent while it was away is gone. The driver's sweep finds a -lost ring too. `LISTEN` keeps state on the server session, so the listener uses `listen_url` when -it is set, and `url` otherwise. Both must name the primary, since a standby refuses `LISTEN`. +lost ring too. The listener sets TCP keepalive to notice a dead connection within a minute. +`LISTEN` keeps state on the server session, so the listener uses `listen_url` when it is set, and +`url` otherwise. Both must name the primary, since a standby refuses `LISTEN`. **Retention.** PostgreSQL has no key expiry, so each store runs a janitor every `janitor_interval` (one minute). A transaction advisory lock around each step lets one janitor diff --git a/crates/streams/src/postgres/conn.rs b/crates/streams/src/postgres/conn.rs index 1704d6a047..6fbb935b2a 100644 --- a/crates/streams/src/postgres/conn.rs +++ b/crates/streams/src/postgres/conn.rs @@ -102,14 +102,15 @@ impl Target { &self, ) -> Result<(Client, mpsc::UnboundedReceiver), tokio_postgres::Error> { let (sender, receiver) = mpsc::unbounded_channel(); + let config = self.listen_config(); let client = match &self.tls { Some(tls) => { - let (client, connection) = self.config.connect(tls.clone()).await?; + let (client, connection) = config.connect(tls.clone()).await?; tokio::spawn(forward(connection, sender)); client } None => { - let (client, connection) = self.config.connect(NoTls).await?; + let (client, connection) = config.connect(NoTls).await?; tokio::spawn(forward(connection, sender)); client } @@ -117,6 +118,18 @@ impl Target { Ok((client, receiver)) } + /// The listener only receives, so a half-open connection shows only to TCP keepalive. The + /// system default notices it after about two hours, and the listener's wakes stop meanwhile. + fn listen_config(&self) -> tokio_postgres::Config { + let mut config = self.config.clone(); + config + .keepalives(true) + .keepalives_idle(Duration::from_secs(30)) + .keepalives_interval(Duration::from_secs(10)) + .keepalives_retries(3); + config + } + /// The pool of `size` connections. Each new connection pins the session settings the store /// functions depend on. pub(super) fn pool(&self, size: usize, timeout: Duration) -> StreamResult { @@ -448,6 +461,20 @@ mod tests { assert_eq!(mode, Mode::VerifyCa); } + #[test] + fn the_listener_notices_a_dead_connection_within_a_minute() { + let target = Target { + config: parse_url("postgres://u:p@localhost/db").unwrap().0, + tls: None, + mode: Mode::Disable, + }; + let config = target.listen_config(); + assert!(config.get_keepalives()); + let detect = config.get_keepalives_idle() + + config.get_keepalives_interval().unwrap() * config.get_keepalives_retries().unwrap(); + assert!(detect <= Duration::from_secs(60), "{detect:?}"); + } + #[test] fn another_scheme_or_sslmode_is_refused() { for url in [ From 1b7f23adc7e7cfa376ea121cdb7662b420f21bdc Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Fri, 16 Oct 2026 16:31:53 -0700 Subject: [PATCH 7/7] Appended before the listener's kill in the reconnect test. A reconnect that came before the append woke the waiter to an empty topic, so the test failed under load. --- crates/streams/tests/postgres_store.rs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/crates/streams/tests/postgres_store.rs b/crates/streams/tests/postgres_store.rs index 6593a7bf03..b34e8b1218 100644 --- a/crates/streams/tests/postgres_store.rs +++ b/crates/streams/tests/postgres_store.rs @@ -662,7 +662,12 @@ async fn a_listener_that_lost_its_connection_wakes_every_waiter() { let waiting = wait_for_one(reader.clone(), &chain); tokio::time::sleep(Duration::from_millis(300)).await; // The writer sends no `NOTIFY` and the poll is an hour away, so only the listener's - // reconnect can wake the waiter. + // reconnect can wake the waiter. The record lands before the kill, since a reconnect that + // came before the append would wake the waiter to an empty topic. + writer + .append(append(&chain, "out", "p", 1, 1)) + .await + .unwrap(); let admin = raw(&url).await; let killed: i64 = admin .query_one( @@ -674,10 +679,6 @@ async fn a_listener_that_lost_its_connection_wakes_every_waiter() { .unwrap() .get(0); assert_eq!(killed, 1); - writer - .append(append(&chain, "out", "p", 1, 1)) - .await - .unwrap(); let waited = waiting.await.unwrap(); assert!(waited < Duration::from_secs(5), "{waited:?}"); }