storage: use read_uncommitted as the progress client's isolation level

See comments within for how using this isolation level is required to
avoid  missing the latest progress record for a sink.

Fix #18628.

Author: Sean Loiselle

src/storage-client/src/sink.rs

@@ -257,7 +257,13 @@ pub async fn publish_kafka_schemas(
     Ok((key_schema_id, value_schema_id))
 }
 
-/// Ensures that the Kafka sink's data and consistency collateral exist.
+/// Ensures that the Kafka sink's data and consistency collateral exist, as well
+/// as returning the last complete timestamp that the last incarnation of the
+/// sink committed to Kafka.
+///
+/// IMPORTANT: to achieve exactly once guarantees, the producer that will resume
+/// production at the returned timestamp *must* have called `init_transactions`
+/// prior to calling this method.
 ///
 /// # Errors
 /// - If the [`KafkaSinkConnection`]'s consistency collateral exists and
@@ -267,95 +273,84 @@ pub async fn build_kafka(
     connection: &KafkaSinkConnection,
     connection_cx: &ConnectionContext,
 ) -> Result<Option<Timestamp>, ContextCreationError> {
-    let client: AdminClient<_> = connection
-        .connection
-        .create_with_context(connection_cx, MzClientContext::default(), &BTreeMap::new())
-        .await
-        .add_context("creating admin client failed")?;
-
-    // Check for existence of progress topic; if it exists and contains data for
-    // this sink, we expect the data topic to exist, as well. Note that we don't
-    // expect the converse to be true because we don't want to prevent users
-    // from creating topics before setting up their sinks.
-    let meta = client
-        .inner()
-        .fetch_metadata(None, Duration::from_secs(10))
-        .check_ssh_status(client.inner().context())
-        .add_context("fetching metadata")?;
-
-    // Check if the broker's metadata already contains the progress topic.
+    // Fetch the progress of the last incarnation of the sink, if any.
     let progress_topic = match &connection.consistency_config {
-        KafkaConsistencyConfig::Progress { topic } => {
-            meta.topics().iter().find(|t| t.name() == topic)
-        }
+        KafkaConsistencyConfig::Progress { topic } => topic,
     };
-
-    // If the consistency topic exists, check to see if it contains this sink's
-    // data.
-    let latest_ts = if let Some(progress_topic) = progress_topic {
-        let progress_client: BaseConsumer<_> = connection
+    // For details about the two clients constructed here, see
+    // `determine_latest_progress_record`.
+    let make_progress_client = |isolation_level: &'static str| async {
+        connection
             .connection
             .create_with_context(
                 connection_cx,
                 MzClientContext::default(),
                 &btreemap! {
                     "group.id" => SinkGroupId::new(sink_id),
-                    "isolation.level" => "read_committed".into(),
+                    "isolation.level" => isolation_level.into(),
                     "enable.auto.commit" => "false".into(),
                     "auto.offset.reset" => "earliest".into(),
                     "enable.partition.eof" => "true".into(),
                 },
             )
-            .await?;
-
-        let progress_client = Arc::new(progress_client);
-        let latest_ts = determine_latest_progress_record(
-            format!("build_kafka_{}", sink_id),
-            progress_topic.name().to_string(),
-            ProgressKey::new(sink_id),
-            Arc::clone(&progress_client),
-        )
+            .await
+    };
+    let progress_client_read_committed = Arc::new(make_progress_client("read_committed").await?);
+    let progress_client_read_uncommitted =
+        Arc::new(make_progress_client("read_uncommitted").await?);
+    let latest_ts = determine_latest_progress_record(
+        format!("build_kafka_{}", sink_id),
+        Arc::clone(&progress_client_read_committed),
+        progress_client_read_uncommitted,
+        progress_topic.to_string(),
+        ProgressKey::new(sink_id),
+    )
+    .await
+    .check_ssh_status(progress_client_read_committed.client().context())?;
+
+    let admin_client: AdminClient<_> = connection
+        .connection
+        .create_with_context(connection_cx, MzClientContext::default(), &BTreeMap::new())
         .await
-        .check_ssh_status(progress_client.client().context())?;
-
-        // If we have progress data, we should have the topic listed in the
-        // broker's metadata. If we don't, error.
-        if latest_ts.is_some() && !meta.topics().iter().any(|t| t.name() == connection.topic) {
-            Err(anyhow::anyhow!(
-                "sink progress data exists, but sink data topic is missing"
-            ))?
-        }
+        .add_context("creating admin client failed")?;
 
-        latest_ts
-    } else {
-        None
-    };
+    // If the progress topic existed and contained data for this sink, we expect
+    // the data topic to exist, as well. Note that we don't expect the converse
+    // to be true because we don't want to prevent users from creating topics
+    // before setting up their sinks.
+    let meta = admin_client
+        .inner()
+        .fetch_metadata(None, Duration::from_secs(10))
+        .check_ssh_status(admin_client.inner().context())
+        .add_context("fetching metadata")?;
+    if latest_ts.is_some() && !meta.topics().iter().any(|t| t.name() == connection.topic) {
+        Err(anyhow::anyhow!(
+            "sink progress data exists, but sink data topic is missing"
+        ))?;
+    }
 
-    // Create Kafka topic.
+    // Create Kafka topics.
+    ensure_kafka_topic(
+        &admin_client,
+        progress_topic,
+        1,
+        connection.replication_factor,
+        KafkaSinkConnectionRetention::default(),
+    )
+    .await
+    .check_ssh_status(admin_client.inner().context())
+    .add_context("error registering kafka progress topic for sink")?;
     ensure_kafka_topic(
-        &client,
+        &admin_client,
         &connection.topic,
         connection.partition_count,
         connection.replication_factor,
         connection.retention,
     )
     .await
-    .check_ssh_status(client.inner().context())
+    .check_ssh_status(admin_client.inner().context())
     .add_context("error registering kafka topic for sink")?;
 
-    match &connection.consistency_config {
-        KafkaConsistencyConfig::Progress { topic } => ensure_kafka_topic(
-            &client,
-            topic,
-            1,
-            connection.replication_factor,
-            KafkaSinkConnectionRetention::default(),
-        )
-        .await
-        .check_ssh_status(client.inner().context())
-        .add_context("error registering kafka consistency topic for sink")?,
-    };
-
     Ok(latest_ts)
 }
 
@@ -376,17 +371,18 @@ pub struct ProgressRecord {
 /// Determines the latest progress record from the specified topic for the given
 /// progress key.
 ///
-/// IMPORTANT: to achieve exactly once guarantees, the producer that will act on
-/// this information *must* have called `init_transactions` prior to calling
-/// this method.
+/// IMPORTANT: to achieve exactly once guarantees, the producer that will resume
+/// production at the returned timestamp *must* have called `init_transactions`
+/// prior to calling this method.
 ///
 /// IMPORTANT: the `progress_client` must have `enable.partition.eof` set to
 /// `true`.
 async fn determine_latest_progress_record(
     name: String,
+    progress_client_read_committed: Arc<BaseConsumer<TunnelingClientContext<MzClientContext>>>,
+    progress_client_read_uncommitted: Arc<BaseConsumer<TunnelingClientContext<MzClientContext>>>,
     progress_topic: String,
     progress_key: ProgressKey,
-    progress_client: Arc<BaseConsumer<TunnelingClientContext<MzClientContext>>>,
 ) -> Result<Option<Timestamp>, anyhow::Error> {
     // ****************************** WARNING ******************************
     // Be VERY careful when editing the code in this function. It is very easy
@@ -403,9 +399,10 @@ async fn determine_latest_progress_record(
     ///
     /// Blocking so should always be called on background thread.
     fn get_latest_ts<C>(
+        progress_client_read_committed: &BaseConsumer<C>,
+        progress_client_read_uncommitted: &BaseConsumer<C>,
         progress_topic: &str,
         progress_key: &ProgressKey,
-        progress_client: &BaseConsumer<C>,
     ) -> Result<Option<Timestamp>, anyhow::Error>
     where
         C: ConsumerContext,
@@ -414,7 +411,7 @@ async fn determine_latest_progress_record(
         // guarantees ordering within a single partition, and we need a strict
         // order on the progress messages we read and write.
         let partitions = match mz_kafka_util::client::get_partitions(
-            progress_client.client(),
+            progress_client_read_committed.client(),
             progress_topic,
             DEFAULT_FETCH_METADATA_TIMEOUT,
         ) {
@@ -452,20 +449,27 @@ async fn determine_latest_progress_record(
         // before we can conclude that we've seen the latest progress record for
         // the specified `progress_key`. A safety argument:
         //
-        //   * The producer has initialized transactions before calling this
-        //     method. At this point, any writes from a previous version of this
-        //     sink should be visible, so the offset we fetch should be larger
-        //     than any previously-written progress record and thus picked up by
-        //     our scan.
+        //   * Our caller has initialized transactions before calling this
+        //     method, which prevents the prior incarnation of this sink from
+        //     committing any further progress records.
         //
-        //     TODO: this is only true when fetching offsets with an isolation
-        //     level of "read uncomitted"! Fix this.
+        //   * We use `read_uncommitted` isolation to ensure that we fetch the
+        //     true high water mark for the topic, even if there are pending
+        //     transactions in the topic. If we used the `read_committed`
+        //     isolation level, we'd instead get the "last stable offset" (LSO),
+        //     which is the offset of the first message in an open transaction,
+        //     which might not include the last progress message committed for
+        //     this sink! (While the caller of this function has fenced out
+        //     older producers for this sink, *other* sinks writing using the
+        //     same progress topic might have long-running transactions that
+        //     hold back the LSO.)
         //
-        //   * If another sink spins up and fences the producer out, we may not
-        //     see the latest progress record... but since the producer has been
-        //     fenced out, it will be unable to act on our stale information.
+        //   * If another sink spins up and fences out the producer for this
+        //     incarnation of the sink, we may not see the latest progress
+        //     record... but since the producer has been fenced out, it will be
+        //     unable to act on our stale information.
         //
-        let (lo, hi) = progress_client
+        let (lo, hi) = progress_client_read_uncommitted
             .fetch_watermarks(progress_topic, partition, DEFAULT_FETCH_METADATA_TIMEOUT)
             .map_err(|e| {
                 anyhow!(
@@ -478,16 +482,18 @@ async fn determine_latest_progress_record(
         let mut tps = TopicPartitionList::new();
         tps.add_partition(progress_topic, partition);
         tps.set_partition_offset(progress_topic, partition, Offset::Beginning)?;
-        progress_client.assign(&tps).with_context(|| {
-            format!(
-                "Error seeking in progress topic {}:{}",
-                progress_topic, partition
-            )
-        })?;
+        progress_client_read_committed
+            .assign(&tps)
+            .with_context(|| {
+                format!(
+                    "Error seeking in progress topic {}:{}",
+                    progress_topic, partition
+                )
+            })?;
 
         // Helper to get the progress consumer's current position.
         let get_position = || {
-            let position = progress_client
+            let position = progress_client_read_committed
                 .position()?
                 .find_partition(progress_topic, partition)
                 .ok_or_else(|| {
@@ -516,13 +522,29 @@ async fn determine_latest_progress_record(
         // Read messages until the consumer is positioned at or beyond the high
         // water mark.
         //
+        // We use `read_committed` isolation to ensure we don't see progress
+        // records for transactions that did not commit. This means we have to
+        // wait for the LSO to progress to the high water mark `hi`, which means
+        // waiting for any open transactions for other sinks using the same
+        // progress topic to complete. We set a short transaction timeout (10s)
+        // to ensure we never need to wait more than 10s.
+        //
+        // Note that the stall time on the progress topic is not a function of
+        // transaction size. We've designed our transactions so that the
+        // progress record is always written last, after all the data has been
+        // written, and so the window of time in which the progress topic has an
+        // open transaction is quite small. The only vulnerability is if another
+        // sink using the same progress topic crashes in that small window
+        // between writing the progress record and committing the transaction,
+        // in which case we have to wait out the transaction timeout.
+        //
         // Important invariant: we only exit this loop successfully (i.e., not
         // returning an error) if we have positive proof of a position at or
         // beyond the high water mark. To make this invariant easy to check, do
         // not use `break` in the body of the loop.
         let mut last_timestamp = None;
         while get_position()? < hi {
-            let message = match progress_client.poll(PROGRESS_RECORD_FETCH_TIMEOUT) {
+            let message = match progress_client_read_committed.poll(PROGRESS_RECORD_FETCH_TIMEOUT) {
                 Some(Ok(message)) => message,
                 Some(Err(KafkaError::PartitionEOF(_))) => {
                     // No message, but the consumer's position may have advanced
@@ -566,7 +588,14 @@ async fn determine_latest_progress_record(
 
     task::spawn_blocking(
         || format!("get_latest_ts:{name}"),
-        move || get_latest_ts(&progress_topic, &progress_key, &progress_client),
+        move || {
+            get_latest_ts(
+                &progress_client_read_committed,
+                &progress_client_read_uncommitted,
+                &progress_topic,
+                &progress_key,
+            )
+        },
     )
     .await?
 }