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