diff --git a/nisshi-storage-sql/src/pg.rs b/nisshi-storage-sql/src/pg.rs index 6a28c717..9d36d0b4 100644 --- a/nisshi-storage-sql/src/pg.rs +++ b/nisshi-storage-sql/src/pg.rs @@ -72,10 +72,10 @@ use opentelemetry::{KeyValue, metrics::Counter}; use rand::{prelude::*, rng}; use serde_json::Value; use tokio_postgres::{ - Config, Row, RowStream, + Config, Row, binary_copy::BinaryCopyInWriter, error::SqlState, - types::{BorrowToSql, ToSql, Type}, + types::{ToSql, Type}, }; use tracing::{debug, error, instrument}; use url::Url; @@ -804,28 +804,6 @@ impl Postgres { .map_err(Into::into) } - #[instrument(skip(self, tx, params))] - async fn tx_prepare_query_raw( - &self, - tx: &Transaction<'_>, - sql: &str, - params: I, - ) -> Result - where - P: BorrowToSql, - I: IntoIterator, - I::IntoIter: ExactSizeIterator, - { - let sql = self.sql_lookup(sql)?; - - let prepared = tx - .prepare_cached(sql) - .await - .inspect_err(|err| error!(?err))?; - - tx.query_raw(&prepared, params).await.map_err(Into::into) - } - #[instrument(skip_all)] async fn produce_in_tx( &self, @@ -1916,30 +1894,31 @@ impl Storage for Postgres { debug!(?topic_uuid, cluster = self.cluster, ?topic); - _ = future::try_join_all( - (0..topic.num_partitions) - .map(|partition| { - let cluster = Box::new(self.cluster.clone()) as Box; - let name = Box::new(topic.name.clone()) as Box; - let partition = Box::new(partition) as Box; - [cluster, name, partition] - }) - .map(|parameters| self.tx_prepare_query_raw(&tx, "topition_insert.sql", parameters)) - .chain( - (0..topic.num_partitions) - .map(|partition| { - let cluster = - Box::new(self.cluster.clone()) as Box; - let name = Box::new(topic.name.clone()) as Box; - let partition = Box::new(partition) as Box; - [cluster, name, partition] - }) - .map(|parameters| { - self.tx_prepare_query_raw(&tx, "watermark_insert.sql", parameters) - }), - ), - ) - .await?; + // The watermark insert selects the topition rows that the topition insert creates, so + // every topition insert must finish before the first watermark insert is sent. A + // statement that is not yet prepared on this connection takes a round trip first, + // which would otherwise let a watermark insert overtake its topition and match no + // rows, leaving a topic that rejects every produce. + for sql in ["topition_insert.sql", "watermark_insert.sql"] { + let inserted = future::try_join_all((0..topic.num_partitions).map(|partition| { + let tx = &tx; + let cluster = &self.cluster; + let name = &topic.name; + + async move { + self.tx_prepare_execute(tx, sql, &[cluster, name, &partition]) + .await + } + })) + .await?; + + if inserted.iter().any(|rows| *rows != 1) { + return Err(Error::Message(format!( + "{sql} did not insert one row per partition of {} in {}", + topic.name, self.cluster + ))); + } + } if let Some(configs) = topic.configs { for config in configs { @@ -6565,6 +6544,71 @@ mod tests { Ok(()) } + /// The watermark insert selects the topition row that the topition insert creates, so + /// `create_topic` must finish the second before sending the first. When it sent both + /// together, a connection holding only the watermark statement in its statement cache + /// sent the watermark insert first; it matched no rows without an error, and the topic + /// had no watermark, so every produce to it failed with `UnknownTopicOrPartition`. No + /// current caller leaves the cache in that state; this test sets it up directly. + #[tokio::test] + async fn create_topic_writes_the_watermark_when_only_its_statement_is_cached() -> Result<()> { + let cluster = alphanumeric_string(15); + + let storage = Postgres::builder(CONNECTION)? + .cluster(cluster.as_str()) + .node(rng().random_range(0..i32::MAX)) + .build(); + + if let Err(err) = storage.connection().await { + eprintln!( + "skipping create_topic_writes_the_watermark_when_only_its_statement_is_cached: {err:?}" + ); + return Ok(()); + } + + storage + .register_broker(BrokerRegistrationRequest { + broker_id: 111, + cluster_id: cluster.clone(), + incarnation_id: Uuid::now_v7(), + rack: None, + }) + .await?; + + { + let c = storage.connection().await?; + _ = c + .prepare_cached(storage.sql_lookup("watermark_insert.sql")?) + .await?; + } + + assert_eq!( + 1, + storage.pool.status().size, + "one connection holds the cache" + ); + + let topic_name = alphanumeric_string(15); + + _ = storage + .create_topic( + CreatableTopic::default() + .name(topic_name.clone()) + .num_partitions(1) + .replication_factor(0) + .assignments(Some([].into())) + .configs(Some([].into())), + false, + ) + .await?; + + let stage = storage.offset_stage(&Topition::new(topic_name, 0)).await?; + + assert_eq!(0, stage.high_watermark); + + Ok(()) + } + /// Deleting one topic must leave every other topic in the cluster untouched. /// `topic_delete_by.sql` narrowed the delete on `topic.cluster` instead of /// `topic.id`, so the subquery only established that the named topic existed