refactor: parallelize compression transactions

Author: ananas-block

forester/src/compressible/compressor.rs

@@ -31,6 +31,16 @@ pub struct Compressor<R: Rpc> {
     payer_keypair: Keypair,
 }
 
+impl<R: Rpc> Clone for Compressor<R> {
+    fn clone(&self) -> Self {
+        Self {
+            rpc_pool: Arc::clone(&self.rpc_pool),
+            tracker: Arc::clone(&self.tracker),
+            payer_keypair: self.payer_keypair.insecure_clone(),
+        }
+    }
+}
+
 impl<R: Rpc> Compressor<R> {
     pub fn new(
         rpc_pool: Arc<SolanaRpcPool<R>>,

forester/src/compressible/config.rs

@@ -7,17 +7,25 @@ pub struct CompressibleConfig {
     /// Batch size for compression operations
     #[serde(default = "default_batch_size")]
     pub batch_size: usize,
+    /// Maximum number of concurrent compression batches
+    #[serde(default = "default_max_concurrent_batches")]
+    pub max_concurrent_batches: usize,
 }
 
 fn default_batch_size() -> usize {
     10
 }
 
+fn default_max_concurrent_batches() -> usize {
+    10
+}
+
 impl CompressibleConfig {
     pub fn new(ws_url: String) -> Self {
         Self {
             ws_url,
             batch_size: default_batch_size(),
+            max_concurrent_batches: default_max_concurrent_batches(),
         }
     }
 }

forester/src/epoch_manager.rs

@@ -1684,34 +1684,62 @@ impl<R: Rpc> EpochManager<R> {
                 current_epoch,
             );
 
-        let mut total_compressed = 0;
-        for (batch_idx, batch) in accounts.chunks(config.batch_size).enumerate() {
-            debug!(
-                "Processing compression batch {}/{} with {} accounts",
-                batch_idx + 1,
-                num_batches,
-                batch.len()
-            );
+        // Create parallel compression futures
+        use futures::stream::StreamExt;
+
+        // Collect chunks into owned vectors to avoid lifetime issues
+        let batches: Vec<(usize, Vec<_>)> = accounts
+            .chunks(config.batch_size)
+            .enumerate()
+            .map(|(idx, chunk)| (idx, chunk.to_vec()))
+            .collect();
+
+        let compression_futures = batches.into_iter().map(|(batch_idx, batch)| {
+            let compressor = compressor.clone();
+            async move {
+                debug!(
+                    "Processing compression batch {}/{} with {} accounts",
+                    batch_idx + 1,
+                    num_batches,
+                    batch.len()
+                );
 
-            match compressor
-                .compress_batch(batch, registered_forester_pda)
-                .await
-            {
-                Ok(sig) => {
+                match compressor
+                    .compress_batch(&batch, registered_forester_pda)
+                    .await
+                {
+                    Ok(sig) => Ok((batch_idx, batch.len(), sig)),
+                    Err(e) => Err((batch_idx, batch.len(), e)),
+                }
+            }
+        });
+
+        // Execute batches in parallel with concurrency limit
+        let results = futures::stream::iter(compression_futures)
+            .buffer_unordered(config.max_concurrent_batches)
+            .collect::<Vec<_>>()
+            .await;
+
+        // Aggregate results
+        let mut total_compressed = 0;
+        for result in results {
+            match result {
+                Ok((batch_idx, count, sig)) => {
                     info!(
                         "Successfully compressed {} accounts in batch {}/{}: {}",
-                        batch.len(),
+                        count,
                         batch_idx + 1,
                         num_batches,
                         sig
                     );
-                    total_compressed += batch.len();
+                    total_compressed += count;
                 }
-                Err(e) => {
+                Err((batch_idx, count, e)) => {
                     error!(
-                        "Compression batch {}/{} failed: {:?}",
+                        "Compression batch {}/{} ({} accounts) failed: {:?}",
                         batch_idx + 1,
                         num_batches,
+                        count,
                         e
                     );
                 }

forester/tests/e2e_test.rs

@@ -254,6 +254,7 @@ async fn e2e_test() {
         compressible_config: Some(forester::compressible::CompressibleConfig {
             ws_url: get_ws_rpc_url(),
             batch_size: 10,
+            max_concurrent_batches: 10,
         }),
     };
     let test_mode = TestMode::from_env();