Repository: MystenLabs/sui
File: crates/sui-indexer-alt-framework/src/pipeline/concurrent/collector.rs (Rust)

Code before:
        // Data for checkpoints that are ready to be sent but haven't been written yet.
        let mut pending: BTreeMap<u64, PendingCheckpoint<H>> = BTreeMap::new();
        let mut pending_rows = 0;

        info!(pipeline = H::NAME, "Starting collector");

        // Wait for main_reader_lo to be initialized before processing any checkpoints.
        let reader_lo_atomic = main_reader_lo.wait().await;

        loop {
            // Eager inner loop: drain processor channel while under backpressure limit.
            while pending_rows < H::MAX_PENDING_ROWS {
                match rx.try_recv() {
                    Ok(indexed) => {
                        pending_rows += receive_checkpoint::<H>(
                            indexed,
                            &mut pending,
                            reader_lo_atomic,
                            &metrics,
                        );
                    }
                    Err(_) => break,
                }
            }

            // Eagerly build and send batches while enough data is pending.
            while pending_rows >= H::MIN_EAGER_ROWS {
                let (batch, watermark, batch_len) =
                    build_batch::<H>(&*handler, &mut pending, &checkpoint_lag_reporter, &metrics);
                pending_rows -= batch_len;

                let batched_rows = BatchedRows {
                    batch,
                    batch_len,
                    watermark,
                };

                if tx.send(batched_rows).await.is_err() {
                    info!(
                        pipeline = H::NAME,
                        "Committer closed channel, stopping collector"
                    );
                    return Ok(());
                }

                if batch_len == 0 {
                    break;
                }
            }

            // Fall back to select! for waiting.

Review comment (performance):
If we have a small amount of data left over, I reckon we should just send it through, in a small batch. The pathology we're trying to avoid is sending a stream of small batches, which is solved by checking before we start generating batches at all whether we have `>= MIN_EAGER_ROWS` in total.

Code after:
        // Data for checkpoints that are ready to be sent but haven't been written yet.
        let mut pending: BTreeMap<u64, PendingCheckpoint<H>> = BTreeMap::new();
        let mut pending_rows = 0;

        info!(pipeline = H::NAME, "Starting collector");

        // Wait for main_reader_lo to be initialized before processing any checkpoints.
        let reader_lo_atomic = main_reader_lo.wait().await;

        loop {
            // Eager inner loop: drain processor channel while under backpressure limit.
            while pending_rows < H::MAX_PENDING_ROWS {
                match rx.try_recv() {
                    Ok(indexed) => {
                        pending_rows += receive_checkpoint::<H>(
                            indexed,
                            &mut pending,
                            reader_lo_atomic,
                            &metrics,
                        );
                    }
                    Err(_) => break,
                }
            }

            // Eagerly build and send batches if enough data is pending. Once we start
            // batching, flush everything including any small tail below the threshold.
            if pending_rows >= H::MIN_EAGER_ROWS {
                loop {
                    let (batch, watermark, batch_len) = build_batch::<H>(
                        &*handler,
                        &mut pending,
                        &checkpoint_lag_reporter,
                        &metrics,
                    );
                    pending_rows -= batch_len;

                    let batched_rows = BatchedRows {
                        batch,
                        batch_len,
                        watermark,
                    };

                    if tx.send(batched_rows).await.is_err() {
                        info!(
                            pipeline = H::NAME,
                            "Committer closed channel, stopping collector"
                        );
                        return Ok(());
                    }