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 7a9c3469bf..05eb5e971f 100644 --- a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto +++ b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto @@ -167,8 +167,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 794f9159d7..c4a330cbf4 100644 --- a/crates/streams/docs/postgres.md +++ b/crates/streams/docs/postgres.md @@ -41,6 +41,32 @@ 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. 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. + +**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. 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 work at a time across processes, also through PgBouncer in transaction mode. A pass trims the topics that are due, at most @@ -69,10 +95,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 +155,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 +172,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 65d600ee3b..1a3a5058ab 100644 --- a/crates/streams/src/lib.rs +++ b/crates/streams/src/lib.rs @@ -33,7 +33,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..6fbb935b2a 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,40 @@ 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 config = self.listen_config(); + let client = match &self.tls { + Some(tls) => { + let (client, connection) = config.connect(tls.clone()).await?; + tokio::spawn(forward(connection, sender)); + client + } + None => { + let (client, connection) = config.connect(NoTls).await?; + tokio::spawn(forward(connection, sender)); + client + } + }; + 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 { @@ -124,6 +161,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. @@ -400,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 [ diff --git a/crates/streams/src/postgres/janitor.rs b/crates/streams/src/postgres/janitor.rs index 93a7a2d514..1444c6cf9a 100644 --- a/crates/streams/src/postgres/janitor.rs +++ b/crates/streams/src/postgres/janitor.rs @@ -2,7 +2,7 @@ //! transaction advisory lock around each step lets one janitor work 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,8 +29,9 @@ 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. A pass - /// stops when another process's janitor holds the lock, since that one does the work. + /// One pass: trims the topics that are due, finishes purges, caps the notification log, and + /// deletes idle chains. A pass stops when another process's janitor holds the lock, since + /// that one does the work. pub(super) async fn janitor_pass(&self) -> StreamResult { let mut client = self.pool.get().await.map_err(unreachable)?; let mut deleted = 0; @@ -61,6 +62,10 @@ impl Inner { break; } } + if self.log { + self.locked(&mut client, &self.sql.cap_log, &[&LOG_CAP]) + .await?; + } self.locked( &mut 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 893fc1fc83..febb039f8f 100644 --- a/crates/streams/src/postgres/mod.rs +++ b/crates/streams/src/postgres/mod.rs @@ -3,12 +3,14 @@ mod conn; mod errors; mod janitor; +mod listen; +mod log; mod migrations; mod wait; use crate::{ - Capabilities, ConditionForms, MAX_BATCH_BYTES, MAX_RECORD_BYTES, StreamError, StreamResult, - StreamStore, + Capabilities, ConditionForms, MAX_BATCH_BYTES, MAX_RECORD_BYTES, NotificationLog, StreamError, + StreamResult, StreamStore, proto::{ ChainId, DeleteOwnerRequest, DeleteOwnerResponse, StoreAppendRequest, StoreAppendResponse, StoreLatestRequest, StoreLatestResponse, StoreReadRequest, StoreReadResponse, StoredClose, @@ -36,6 +38,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)] @@ -63,6 +72,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 { @@ -79,6 +99,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, } } } @@ -86,6 +110,7 @@ impl PostgresStoreOptions { /// A [StreamStore] on PostgreSQL. pub struct PostgresStore { inner: Arc, + log: Option>, tasks: Vec>, } @@ -178,6 +203,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`. @@ -249,6 +278,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)"), } } } @@ -265,6 +313,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 { @@ -276,11 +328,18 @@ 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}}, start with no `pg_` and be no \ + "the schema must match [a-z_][a-z0-9_]{{0,{}}}, start with no `pg_` and be no \ reserved word, 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!( @@ -289,11 +348,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), @@ -306,15 +385,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`. @@ -329,7 +425,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'_') // PostgreSQL keeps `pg_` names for itself, and a reserved word needs quotes in SQL. && !schema.starts_with("pg_") @@ -476,6 +572,7 @@ fn wait_of(request: &StoreReadRequest) -> Duration { struct Look { response: StoreReadResponse, seen: Seen, + topic_id: Option, } impl Inner { @@ -557,6 +654,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 @@ -674,6 +772,10 @@ impl StreamStore for PostgresStore { NAME } + fn notification_log(&self) -> Option> { + self.log.clone().map(|log| log as Arc) + } + fn capabilities(&self) -> Capabilities { Capabilities { producer_state: true, @@ -720,11 +822,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?; @@ -788,14 +890,19 @@ 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 { 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() @@ -1156,8 +1263,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..2108fa42d8 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,30 @@ impl Waiters { } } + /// 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 == 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() { + 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(); @@ -102,10 +131,26 @@ 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(); - 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,11 +163,11 @@ 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()); @@ -132,7 +177,7 @@ mod tests { #[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 6c91072054..ce13bbd602 100644 --- a/crates/streams/src/service.rs +++ b/crates/streams/src/service.rs @@ -448,10 +448,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 6facb4bac5..45152764a8 100644 --- a/crates/streams/tests/conformance.rs +++ b/crates/streams/tests/conformance.rs @@ -861,22 +861,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"); @@ -899,6 +938,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; @@ -923,6 +971,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), @@ -3283,11 +3332,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, ); @@ -3298,11 +3374,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 a5fea22e24..b34e8b1218 100644 --- a/crates/streams/tests/postgres_store.rs +++ b/crates/streams/tests/postgres_store.rs @@ -582,3 +582,103 @@ 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(); + // 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() + .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. 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( + "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); + let waited = waiting.await.unwrap(); + assert!(waited < Duration::from_secs(5), "{waited:?}"); +}