refactor: forester(send_transaction): add timeout checks (#1631)

Author: sergeytimoshin

forester/src/send_transaction.rs

@@ -3,6 +3,7 @@ use std::{
         atomic::{AtomicBool, AtomicUsize, Ordering},
         Arc,
     },
+    time::Duration,
     vec,
 };
 
@@ -144,7 +145,6 @@ pub async fn send_batched_transactions<T: TransactionBuilder, R: RpcConnection>(
         return Ok(0);
     }
 
-    // Get priority fee and blockhash
     let (recent_blockhash, current_block_height) = {
         let mut rpc = pool.get_connection().await?;
         (
@@ -202,7 +202,6 @@ pub async fn send_batched_transactions<T: TransactionBuilder, R: RpcConnection>(
         }
     });
 
-    // Process work items in chunks of `config.build_transaction_batch_config.batch_size`
     let work_items: Vec<WorkItem> = queue_item_data
         .into_iter()
         .map(|data| WorkItem {
@@ -211,10 +210,16 @@ pub async fn send_batched_transactions<T: TransactionBuilder, R: RpcConnection>(
         })
         .collect();
 
+    let buffer_duration = Duration::from_secs(2);
+    let adjusted_timeout = if config.retry_config.timeout > buffer_duration {
+        config.retry_config.timeout - buffer_duration
+    } else {
+        return Ok(0);
+    };
+    let timeout_deadline = start_time + adjusted_timeout;
+
     for work_chunk in work_items.chunks(config.build_transaction_batch_config.batch_size as usize) {
-        if cancel_signal.load(Ordering::SeqCst)
-            || start_time.elapsed() >= config.retry_config.timeout
-        {
+        if cancel_signal.load(Ordering::SeqCst) || Instant::now() >= timeout_deadline {
             break;
         }
 
@@ -230,12 +235,22 @@ pub async fn send_batched_transactions<T: TransactionBuilder, R: RpcConnection>(
             )
             .await?;
 
-        // Spawn transaction senders
+        let now = Instant::now();
+        if now >= timeout_deadline {
+            break;
+        }
+
         for tx in transactions {
             if cancel_signal.load(Ordering::SeqCst) {
                 break;
             }
 
+            let now = Instant::now();
+            if now >= timeout_deadline {
+                warn!("Reached timeout deadline, stopping batch processing");
+                break;
+            }
+
             let tx_sender = tx_sender.clone();
             let pool_clone = pool.clone();
             let config = RpcSendTransactionConfig {
@@ -245,21 +260,25 @@ pub async fn send_batched_transactions<T: TransactionBuilder, R: RpcConnection>(
                 ..Default::default()
             };
 
+            let cancel_signal_clone = cancel_signal.clone();
+            let deadline = timeout_deadline;
+
             tokio::spawn(async move {
+                if cancel_signal_clone.load(Ordering::SeqCst) || Instant::now() >= deadline {
+                    return;
+                }
+
                 if let Ok(mut rpc) = pool_clone.get_connection().await {
                     let result = rpc.process_transaction_with_config(tx, config).await;
-                    let _ = tx_sender.send(result).await;
+                    if !cancel_signal_clone.load(Ordering::SeqCst) {
+                        let _ = tx_sender.send(result).await;
+                    }
                 }
             });
         }
     }
-
-    // Drop sender to allow processor to complete
     drop(tx_sender);
-
-    // Wait for processor to complete
     processor_handle.await?;
-
     Ok(num_sent_transactions.load(Ordering::SeqCst))
 }