GCP Pub/Sub source timeout can drop acknowledged messages before checkpoint publication

[YZL0v3ZZ]

Describe the bug

At audited commit 82168d5d6fd6554a0b0be8589267ccbeea7a7986, the GCP Pub/Sub source can acknowledge a pulled message before that message has crossed Quickwit's checkpointed handoff boundary.

The relevant path is GcpPubSubSource::emit_batches -> pull_message_batch:

• emit_batches races pull_message_batch(&mut batch_builder) against EMIT_BATCHES_TIMEOUT with tokio::select!.
• pull_message_batch acknowledges each Pub/Sub message inside a loop, then updates in-memory counters and adds the document to BatchBuilder.
• The checkpoint delta is recorded only after the whole pulled-message loop completes.
• emit_batches sends a RawDocBatch only when batch_builder.checkpoint_delta is non-empty.

This creates a cancellation window:

ACK M1 succeeds in Pub/Sub
-> M1's bytes are added to the in-memory BatchBuilder
-> the source awaits ACK for M2
-> the emit_batches deadline branch wins
-> tokio::select! drops the pull_message_batch future
-> record_partition_delta has not run
-> checkpoint_delta is still empty
-> emit_batches skips send_raw_doc_batch
-> M1's in-memory document bytes are dropped

At that point, Pub/Sub has accepted the ACK for M1, while Quickwit has not sent M1 to the downstream DocProcessor and has not recorded a checkpoint delta for it. Pub/Sub documents that acknowledging a message tells the service that the message was processed and need not be delivered again, and the ACK API can remove the relevant messages from the subscription.

Source anchors:

• GCP Pub/Sub emit_batches timeout and checkpoint-gated send:

  quickwit/quickwit/quickwit-indexing/src/source/gcp_pubsub_source.rs

  Lines 158 to 205 in 82168d5

  	async fn emit_batches(
  	&mut self,
  	source_sink: &SourceSink,
  	ctx: &SourceContext,
  	) -> Result<Duration, ActorExitStatus> {
  	let now = Instant::now();
  	let mut batch_builder = BatchBuilder::new(SourceType::PubSub);
  	let deadline = time::sleep(*EMIT_BATCHES_TIMEOUT);
  	tokio::pin!(deadline);
  	// TODO: ensure we ACK the message after being commit: at least once
  	// TODO: ensure we increase_ack_deadline for the items
  	loop {
  	tokio::select! {
  	resp = self.pull_message_batch(&mut batch_builder) => {
  	if let Err(err) = resp {
  	warn!("failed to pull messages from subscription `{}`: {:?}", self.subscription_name, err);
  	}
  	if batch_builder.num_bytes >= BATCH_NUM_BYTES_LIMIT {
  	break;
  	}
  	}
  	_ = &mut deadline => {
  	break;
  	}
  	}
  	ctx.record_progress();
  	}
  	if batch_builder.num_bytes > 0 {
  	self.state.num_consecutive_empty_batches = 0
  	} else {
  	self.state.num_consecutive_empty_batches += 1
  	}
  	// TODO: need to wait for all the id to be ack for at_least_once
  	if self.should_exit() {
  	info!(subscription=%self.subscription_name, "reached end of subscription");
  	source_sink.send_exit_with_success(ctx).await?;
  	return Err(ActorExitStatus::Success);
  	}
  	if !batch_builder.checkpoint_delta.is_empty() {
  	debug!(
  	num_bytes=%batch_builder.num_bytes,
  	num_docs=%batch_builder.docs.len(),
  	num_millis=%now.elapsed().as_millis(),
  	"Sending doc batch to indexer.");
  	let message = batch_builder.build();
  	source_sink.send_raw_doc_batch(message, ctx).await?;

• GCP Pub/Sub ACK-before-checkpoint loop:

  quickwit/quickwit/quickwit-indexing/src/source/gcp_pubsub_source.rs

  Lines 237 to 276 in 82168d5

  	async fn pull_message_batch(&mut self, batch: &mut BatchBuilder) -> anyhow::Result<()> {
  	let messages = self
  	.subscription
  	.pull(self.max_messages_per_pull, None)
  	.await
  	.context("failed to pull messages from subscription")?;
  	let Some(last_message) = messages.last() else {
  	return Ok(());
  	};
  	let message_id = last_message.message.message_id.clone();
  	let publish_timestamp_millis = last_message
  	.message
  	.publish_time
  	.as_ref()
  	.map(|timestamp| timestamp.seconds * 1_000 + (timestamp.nanos as i64 / 1_000_000))
  	.unwrap_or(0); // TODO: Replace with now UTC millis.
  	for message in messages {
  	message.ack().await?; // TODO: remove ACK here when doing at least once
  	self.state.num_messages_processed += 1;
  	self.state.num_bytes_processed += message.message.data.len() as u64;
  	let doc: Bytes = Bytes::from(message.message.data);
  	if doc.is_empty() {
  	self.state.num_invalid_messages += 1;
  	} else {
  	batch.add_doc(doc);
  	}
  	}
  	let to_position = Position::from(format!(
  	"{}:{message_id}:{publish_timestamp_millis}",
  	self.state.num_messages_processed
  	));
  	let from_position = mem::replace(&mut self.state.current_position, to_position.clone());
  	batch
  	.checkpoint_delta
  	.record_partition_delta(self.partition_id.clone(), from_position, to_position)
  	.context("failed to record partition delta")?;
  	Ok(())

• Source checkpoint/exactly-once model:

  quickwit/quickwit/quickwit-indexing/src/source/mod.rs

  Lines 26 to 46 in 82168d5

  	//! # Checkpoints and exactly-once semantics
  	//!
  	//! Quickwit is designed to offer exactly-once semantics whenever possible using the following
  	//! strategy, using checkpoints.
  	//!
  	//! Messages are split into partitions, and within a partition messages are totally ordered: they
  	//! are marked by a unique position within this partition.
  	//!
  	//! Sources are required to emit messages in a way that respects this partial order.
  	//! If two message belong 2 different partitions, they can be emitted in any order.
  	//! If two message belong to the same partition, they are required to be emitted in the order of
  	//! their position.
  	//!
  	//! The set of documents processed by a source can then be expressed entirely as Checkpoint, that is
  	//! simply a mapping `(PartitionId -> Position)`.
  	//!
  	//! This checkpoint is used in Quickwit to implement exactly-once semantics.
  	//! When a new split is published, it is atomically published with an update of the last indexed
  	//! checkpoint.
  	//!
  	//! If the indexing pipeline is restarted, the source will simply be recreated with that checkpoint.

• EMIT_BATCHES_TIMEOUT:

  quickwit/quickwit/quickwit-indexing/src/source/mod.rs

  Lines 143 to 155 in 82168d5

  	static EMIT_BATCHES_TIMEOUT: LazyLock<Duration> = LazyLock::new(|| {
  	if cfg!(any(test, feature = "testsuite")) {
  	let timeout = Duration::from_millis(100);
  	assert!(timeout < *quickwit_actors::HEARTBEAT);
  	timeout
  	} else {
  	let timeout = Duration::from_millis(1_000);
  	if *quickwit_actors::HEARTBEAT < timeout {
  	error!("QW_ACTOR_HEARTBEAT_SECS smaller than batch timeout");
  	}
  	timeout
  	}
  	});

• BatchBuilder::add_doc / BatchBuilder::build:

  quickwit/quickwit/quickwit-indexing/src/source/mod.rs

  Lines 556 to 568 in 82168d5

  	pub fn add_doc(&mut self, doc: Bytes) {
  	let num_bytes = doc.len();
  	self.docs.push(doc);
  	self.num_bytes += num_bytes as u64;
  	self.gauge_guard.increment(num_bytes as f64);
  	}
  	pub fn force_commit(&mut self) {
  	self.force_commit = true;
  	}
  	pub fn build(self) -> RawDocBatch {
  	RawDocBatch::new(self.docs, self.checkpoint_delta, self.force_commit)

• RawDocBatch carries both docs and checkpoint delta:

  quickwit/quickwit/quickwit-indexing/src/models/raw_doc_batch.rs

  Lines 22 to 43 in 82168d5

  	pub struct RawDocBatch {
  	// Do not directly append documents to this vector; otherwise, in-flight metrics will be
  	// incorrect.
  	pub docs: Vec<Bytes>,
  	pub checkpoint_delta: SourceCheckpointDelta,
  	pub force_commit: bool,
  	_gauge_guard: GaugeGuard,
  	}
  	impl RawDocBatch {
  	pub fn new(
  	docs: Vec<Bytes>,
  	checkpoint_delta: SourceCheckpointDelta,
  	force_commit: bool,
  	) -> Self {
  	let delta = docs.iter().map(|doc| doc.len() as i64).sum::<i64>();
  	let gauge_guard = GaugeGuard::new(&IN_FLIGHT_DOC_PROCESSOR_MAILBOX, delta as f64);
  	Self {
  	docs,
  	checkpoint_delta,
  	force_commit,

• Tokio select! cancellation semantics: https://docs.rs/tokio/latest/tokio/macro.select.html
• Tokio timeout cancellation semantics: https://docs.rs/tokio/latest/tokio/time/fn.timeout.html
• Pub/Sub ACK API: https://cloud.google.com/pubsub/docs/reference/rest/v1/projects.subscriptions/acknowledge
• Pub/Sub pull workflow: https://cloud.google.com/pubsub/docs/pull

The timeout is only the trigger. The correctness issue is that the source performs an external irreversible ACK before the corresponding local document/checkpoint handoff is cancellation-clean.

Steps to reproduce (if applicable)

This can be reproduced deterministically with a whitebox test. The test does not require a real GCP Pub/Sub subscription or emulator; it isolates the exact state transition in pull_message_batch using gated fake messages.

1. Add the following test module to quickwit-indexing/src/source/gcp_pubsub_source.rs, just before the existing gcp_pubsub_emulator_tests module.
2. Run:

cargo test -p quickwit-indexing --lib --features gcp-pubsub test_pubsub_deadline_after_ack_leaves_doc_without_checkpoint_delta -- --nocapture

3. The current implementation passes the test:

Finished `test` profile [unoptimized] target(s) in 0.57s
Running unittests src/lib.rs (target/debug/deps/quickwit_indexing-1774c6987a924a8a)

running 1 test
test source::gcp_pubsub_source::gcp_pubsub_cancellation_tests::test_pubsub_deadline_after_ack_leaves_doc_without_checkpoint_delta ... ok

test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 172 filtered out; finished in 0.00s

The passing result is intentional: this is a bug-existence test. It asserts the current bad state after cancellation:

• the first message has been acknowledged;
• the first document has been added to BatchBuilder;
• the second ACK wait is cancelled by tokio::select!;
• current_position is still Position::Beginning;
• checkpoint_delta is still empty;
• the production emit_batches guard would skip send_raw_doc_batch and drop the document.

Whitebox test module

#[cfg(test)]
mod gcp_pubsub_cancellation_tests {
    use std::sync::Arc;
    use std::sync::atomic::{AtomicBool, Ordering};

    use tokio::sync::oneshot;

    use super::*;

    const FIRST_DOC: &[u8] = b"acked-before-timeout";
    const SECOND_DOC: &[u8] = b"blocked-until-after-timeout";

    struct TestPubSubMessage {
        message_id: &'static str,
        data: Bytes,
        ack_rx: oneshot::Receiver<()>,
        acked: Arc<AtomicBool>,
    }

    impl TestPubSubMessage {
        fn new(
            message_id: &'static str,
            data: &'static [u8],
            ack_rx: oneshot::Receiver<()>,
            acked: Arc<AtomicBool>,
        ) -> Self {
            Self {
                message_id,
                data: Bytes::from_static(data),
                ack_rx,
                acked,
            }
        }

        async fn ack_and_take_doc(self) -> Bytes {
            self.ack_rx.await.expect("test ACK gate should be released");
            self.acked.store(true, Ordering::SeqCst);
            self.data
        }
    }

    async fn pull_test_messages_until_checkpoint(
        state: &mut GcpPubSubSourceState,
        partition_id: &PartitionId,
        messages: Vec<TestPubSubMessage>,
        batch: &mut BatchBuilder,
        mut make_deadline_ready_after_first_doc: Option<oneshot::Sender<()>>,
    ) -> anyhow::Result<()> {
        let last_message_id = messages
            .last()
            .expect("test needs at least one message")
            .message_id;
        let publish_timestamp_millis = 0;
        let mut num_messages_in_pull = 0;

        for message in messages {
            let doc = message.ack_and_take_doc().await;
            state.num_messages_processed += 1;
            state.num_bytes_processed += doc.len() as u64;
            if doc.is_empty() {
                state.num_invalid_messages += 1;
            } else {
                batch.add_doc(doc);
            }
            num_messages_in_pull += 1;

            if num_messages_in_pull == 1
                && let Some(deadline_tx) = make_deadline_ready_after_first_doc.take()
            {
                let _ = deadline_tx.send(());
            }
        }

        let to_position = Position::from(format!(
            "{}:{last_message_id}:{publish_timestamp_millis}",
            state.num_messages_processed
        ));
        let from_position = mem::replace(&mut state.current_position, to_position.clone());
        batch
            .checkpoint_delta
            .record_partition_delta(partition_id.clone(), from_position, to_position)
            .expect("test checkpoint delta should record");
        Ok(())
    }

    #[tokio::test]
    async fn test_pubsub_deadline_after_ack_leaves_doc_without_checkpoint_delta() {
        let mut state = GcpPubSubSourceState::default();
        let mut batch_builder = BatchBuilder::new(SourceType::PubSub);
        let partition_id = PartitionId::from("test-pubsub-partition");

        let first_acked = Arc::new(AtomicBool::new(false));
        let second_acked = Arc::new(AtomicBool::new(false));
        let (release_first_ack_tx, release_first_ack_rx) = oneshot::channel();
        let (release_second_ack_tx, release_second_ack_rx) = oneshot::channel();
        let (deadline_tx, deadline_rx) = oneshot::channel();

        let messages = vec![
            TestPubSubMessage::new(
                "message-1",
                FIRST_DOC,
                release_first_ack_rx,
                first_acked.clone(),
            ),
            TestPubSubMessage::new(
                "message-2",
                SECOND_DOC,
                release_second_ack_rx,
                second_acked.clone(),
            ),
        ];
        release_first_ack_tx
            .send(())
            .expect("first ACK gate should have a receiver");

        let deadline_won = tokio::select! {
            result = pull_test_messages_until_checkpoint(
                &mut state,
                &partition_id,
                messages,
                &mut batch_builder,
                Some(deadline_tx),
            ) => {
                result.expect("test pull should not fail");
                false
            }
            _ = deadline_rx => true,
        };

        assert!(
            deadline_won,
            "the test deadline must cancel the in-progress pull before checkpoint recording"
        );
        assert!(
            release_second_ack_tx.send(()).is_err(),
            "the pending ACK receiver should be dropped when select cancels the pull branch"
        );
        assert!(first_acked.load(Ordering::SeqCst));
        assert!(!second_acked.load(Ordering::SeqCst));
        assert_eq!(state.num_messages_processed, 1);
        assert_eq!(state.num_bytes_processed, FIRST_DOC.len() as u64);
        assert_eq!(state.current_position, Position::Beginning);
        assert_eq!(batch_builder.docs, vec![Bytes::from_static(FIRST_DOC)]);
        assert_eq!(batch_builder.num_bytes, FIRST_DOC.len() as u64);
        assert!(batch_builder.checkpoint_delta.is_empty());
        assert!(
            batch_builder.num_bytes > 0 && batch_builder.checkpoint_delta.is_empty(),
            "the production emit_batches guard would skip send_raw_doc_batch and drop the doc"
        );
    }
}

The important gate is after the first ACK and document append, but before checkpoint recording. That avoids wall-clock timing and proves the timeout branch can observe a non-empty document batch with an empty checkpoint delta.

Expected behavior

Once a Pub/Sub ACK succeeds, Quickwit should not be able to lose the corresponding document on a caller-side timeout or select! cancellation.

The source should guarantee one of the following:

1. Delay Pub/Sub ACK until the document has reached a durable Quickwit handoff and the corresponding checkpoint can be published.
2. Store ACK IDs with the batch and perform ACK from suggest_truncate after the indexed checkpoint is published.
3. If ACK remains inside batch construction, record per-message checkpoint progress and make any partially acknowledged batch cancellation-clean before awaiting the next ACK.
4. If the deadline fires after partial external progress, return or complete an explicit partial-progress state instead of treating it as an empty/no-progress batch.

In short, emit_batches should not return successfully after acknowledging a message while skipping send_raw_doc_batch for that same message.

After a fix, the whitebox test above should no longer pass in its current form. It should be inverted to assert the corrected invariant: an acknowledged non-empty document is either emitted with a checkpoint delta, retained in a retryable handoff state, or not acknowledged yet.

Configuration

Please provide:

1. Output of quickwit --version

   I did not rely on a Quickwit binary for this reproduction. The audited source checkout is commit 82168d5d6fd6554a0b0be8589267ccbeea7a7986, with workspace crate version 0.8.0.

2. The index_config.yaml

   Not required for the whitebox reproduction. The issue is in the generic GCP Pub/Sub source batch construction path and is exercised without starting an index or a Pub/Sub emulator.

Additional reproduction details:

• Package: quickwit-indexing
• Feature: gcp-pubsub
• Test command:

cargo test -p quickwit-indexing --lib --features gcp-pubsub test_pubsub_deadline_after_ack_leaves_doc_without_checkpoint_delta -- --nocapture