[kv store] add timeout to bigtable ingestion

crates/sui-data-ingestion-core/src/progress_store/mod.rs

@@ -38,6 +38,9 @@ impl<P: ProgressStore> ProgressStore for ProgressStoreWrapper<P> {
         task_name: String,
         checkpoint_number: CheckpointSequenceNumber,
     ) -> Result<()> {
+        if checkpoint_number % 100 != 0 {
+            return Ok(());
+        }
         self.progress_store
             .save(task_name.clone(), checkpoint_number)
             .await?;

crates/sui-data-ingestion-core/src/reducer.rs

@@ -8,6 +8,7 @@ use std::collections::HashMap;
 use sui_types::messages_checkpoint::CheckpointSequenceNumber;
 use tokio::sync::mpsc;
 use tokio_stream::wrappers::ReceiverStream;
+use tracing::info;
 
 pub(crate) async fn reduce<W: Worker>(
     task_name: String,
@@ -39,6 +40,7 @@ pub(crate) async fn reduce<W: Worker>(
             }
             current_checkpoint_number += 1;
         }
+        info!("reducer batch size {}", batch.len());
         match reducer {
             Some(ref reducer) => {
                 if reducer.should_close_batch(&batch, None) {

crates/sui-data-ingestion/src/main.rs

@@ -6,6 +6,7 @@ use prometheus::Registry;
 use serde::{Deserialize, Serialize};
 use std::env;
 use std::path::PathBuf;
+use std::time::Duration;
 use sui_data_ingestion::{
     ArchivalConfig, ArchivalReducer, ArchivalWorker, BlobTaskConfig, BlobWorker,
     DynamoDBProgressStore, KVStoreTaskConfig, KVStoreWorker,
@@ -45,6 +46,7 @@ struct ProgressStoreConfig {
 #[derive(Serialize, Deserialize, Clone, Debug)]
 struct BigTableTaskConfig {
     instance_id: String,
+    timeout_secs: usize,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -154,7 +156,12 @@ async fn main() -> Result<()> {
                 executor.register(worker_pool).await?;
             }
             Task::BigTableKV(kv_config) => {
-                let client = BigTableClient::new_remote(kv_config.instance_id, false, None).await?;
+                let client = BigTableClient::new_remote(
+                    kv_config.instance_id,
+                    false,
+                    Some(Duration::from_secs(kv_config.timeout_secs as u64)),
+                )
+                .await?;
                 let worker_pool = WorkerPool::new(
                     KvWorker { client },
                     task_config.name,