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

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: https://github.com/quickwit-oss/quickwit/blob/82168d5d6fd6554a0b0be8589267ccbeea7a7986/quickwit/quickwit-indexing/src/source/gcp_pubsub_source.rs#L158-L205
• GCP Pub/Sub ACK-before-checkpoint loop: https://github.com/quickwit-oss/quickwit/blob/82168d5d6fd6554a0b0be8589267ccbeea7a7986/quickwit/quickwit-indexing/src/source/gcp_pubsub_source.rs#L237-L276
• Source checkpoint/exactly-once model: https://github.com/quickwit-oss/quickwit/blob/82168d5d6fd6554a0b0be8589267ccbeea7a7986/quickwit/quickwit-indexing/src/source/mod.rs#L26-L46
• EMIT_BATCHES_TIMEOUT: https://github.com/quickwit-oss/quickwit/blob/82168d5d6fd6554a0b0be8589267ccbeea7a7986/quickwit/quickwit-indexing/src/source/mod.rs#L143-L155
• BatchBuilder::add_doc / BatchBuilder::build: https://github.com/quickwit-oss/quickwit/blob/82168d5d6fd6554a0b0be8589267ccbeea7a7986/quickwit/quickwit-indexing/src/source/mod.rs#L556-L568
• RawDocBatch carries both docs and checkpoint delta: https://github.com/quickwit-oss/quickwit/blob/82168d5d6fd6554a0b0be8589267ccbeea7a7986/quickwit/quickwit-indexing/src/models/raw_doc_batch.rs#L22-L43
• 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.

#[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