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

Code before:
    handler: Arc<H>,
    config: CommitterConfig,
    rx: mpsc::Receiver<BatchedRows<H>>,
    tx: mpsc::Sender<Vec<WatermarkPart>>,
    db: H::Store,
    metrics: Arc<IndexerMetrics>,
) -> Service {
    Service::new().spawn_aborting(async move {
        info!(pipeline = H::NAME, "Starting committer");
        let checkpoint_lag_reporter = CheckpointLagMetricReporter::new_for_pipeline::<H>(
            &metrics.partially_committed_checkpoint_timestamp_lag,
            &metrics.latest_partially_committed_checkpoint_timestamp_lag_ms,
            &metrics.latest_partially_committed_checkpoint,
        );

        match ReceiverStream::new(rx)
            .try_for_each_spawned(
                config.write_concurrency,
                |BatchedRows {
                     batch,
                     batch_len,
                     watermark,
                 }| {
                    let batch = Arc::new(batch);
                    let handler = handler.clone();
                    let tx = tx.clone();
                    let db = db.clone();
                    let metrics = metrics.clone();
                    let checkpoint_lag_reporter = checkpoint_lag_reporter.clone();

                    // Repeatedly try to get a connection to the DB and write the batch. Use an
                    // exponential backoff in case the failure is due to contention over the DB
                    // connection pool.
                    let backoff = ExponentialBackoff {
                        initial_interval: INITIAL_RETRY_INTERVAL,
                        current_interval: INITIAL_RETRY_INTERVAL,
                        max_interval: MAX_RETRY_INTERVAL,
                        max_elapsed_time: None,
                        ..Default::default()
                    };

                    let highest_checkpoint = watermark.iter().map(|w| w.checkpoint()).max();
                    let highest_checkpoint_timestamp =
                        watermark.iter().map(|w| w.timestamp_ms()).max();

                    use backoff::Error as BE;
                    let commit = move || {
                        let batch = batch.clone();
                        let handler = handler.clone();
                        let db = db.clone();
                        let metrics = metrics.clone();

Review comment (suggestion):
It seems like it's completely possible for someone to pass a `0` value and cause this line to panic, which is not ideal. Couple of options:

- Use `NonZeroU32` as the type all the way up in the config layer. `serde` and `clap` both support validating for zero when parsing `NonZero*` types (preferred).
- Use `0` as the sentinel value meaning "no rate limit" (less ideal).

Code after:
/// batches on `rx` and writes them out to the `db` concurrently (`config.write_concurrency`
/// controls the degree of fan-out).
///
/// The writing of each batch will be repeatedly retried on an exponential back-off until it
/// succeeds. Once the write succeeds, the [WatermarkPart]s for that batch are sent on `tx` to the
/// watermark task.
///
/// This task will shutdown if its receiver or sender channels are closed.
pub(super) fn committer<H: Handler + 'static>(
    handler: Arc<H>,
    config: CommitterConfig,
    rx: mpsc::Receiver<BatchedRows<H>>,
    tx: mpsc::Sender<Vec<WatermarkPart>>,
    db: H::Store,
    metrics: Arc<IndexerMetrics>,
) -> Service {
    Service::new().spawn_aborting(async move {
        info!(pipeline = H::NAME, "Starting committer");
        let checkpoint_lag_reporter = CheckpointLagMetricReporter::new_for_pipeline::<H>(
            &metrics.partially_committed_checkpoint_timestamp_lag,
            &metrics.latest_partially_committed_checkpoint_timestamp_lag_ms,
            &metrics.latest_partially_committed_checkpoint,
        );

        let rate_limiter: Option<Arc<SharedRateLimiter>> = config.max_rows_per_second.map(|rps| {
            let burst = NonZeroU32::new(rps as u32).expect("max_rows_per_second must be > 0");
            let quota = Quota::per_second(burst);
            Arc::new(SharedRateLimiter {
                limiter: RateLimiter::direct(quota),
                burst,
            })
        });

        match ReceiverStream::new(rx)
            .try_for_each_spawned(
                config.write_concurrency,
                |BatchedRows {
                     batch,
                     batch_len,
                     watermark,
                 }| {
                    let batch = Arc::new(batch);
                    let handler = handler.clone();
                    let tx = tx.clone();
                    let db = db.clone();
                    let metrics = metrics.clone();
                    let checkpoint_lag_reporter = checkpoint_lag_reporter.clone();
                    let rate_limiter = rate_limiter.clone();

                    // Repeatedly try to get a connection to the DB and write the batch. Use an
                    // exponential backoff in case the failure is due to contention over the DB