feat(sv-consumer): Introduce concurrent tasks configuration (#483)

Author: luizstacio

crates/domains/src/infra/db/db_impl.rs

@@ -46,7 +46,7 @@ pub static DB_POOL_SIZE: LazyLock<usize> = LazyLock::new(|| {
     dotenvy::var("DB_POOL_SIZE")
         .ok()
         .and_then(|val| val.parse().ok())
-        .unwrap_or(2000)
+        .unwrap_or(110)
 });
 
 pub static DB_ACQUIRE_TIMEOUT: LazyLock<usize> = LazyLock::new(|| {
@@ -105,6 +105,7 @@ impl Db {
                 .options(Self::connect_opts(opts));
 
         sqlx::postgres::PgPoolOptions::new()
+            .min_connections(2)
             .max_connections(opts.pool_size.unwrap_or_default())
             .acquire_timeout(opts.acquire_timeout.unwrap_or_default())
             .idle_timeout(opts.idle_timeout)

crates/message-broker/src/nats.rs

@@ -59,6 +59,14 @@ impl NatsMessageBroker {
         Self::new(&opts).await
     }
 
+    pub async fn setup_with_opts(
+        opts: &NatsOpts,
+    ) -> Result<Arc<NatsMessageBroker>, MessageBrokerError> {
+        let broker = Self::new(opts).await?;
+        broker.setup_queues().await?;
+        Ok(broker.arc())
+    }
+
     pub async fn setup(
         url: &str,
         namespace: Option<&str>,

services/consumer/src/cli.rs

@@ -37,4 +37,12 @@ pub struct Cli {
         help = "Enable metrics"
     )]
     pub use_metrics: bool,
+    /// Max concurrent tasks
+    #[arg(
+        long,
+        env = "CONCURRENT_TASKS",
+        default_value = "300",
+        help = "Number of concurrent tasks, this is the number of total blocks processed in parallel, this will also be reflected in the number of connections to the database"
+    )]
+    pub concurrent_tasks: usize,
 }

services/consumer/src/executor/block_executor.rs

@@ -54,13 +54,9 @@ use super::{
 };
 use crate::{errors::ConsumerError, metrics::Metrics};
 
-const MAX_CONCURRENT_TASKS: usize = 32;
-const BATCH_SIZE: usize = 100;
-
 #[derive(Debug)]
 enum ProcessResult {
     Store(Result<BlockStats, ConsumerError>),
-    Stream(Result<BlockStats, ConsumerError>),
 }
 
 pub struct BlockExecutor {
@@ -69,6 +65,7 @@ pub struct BlockExecutor {
     fuel_streams: Arc<FuelStreams>,
     semaphore: Arc<Semaphore>,
     telemetry: Arc<Telemetry<Metrics>>,
+    concurrent_tasks: usize,
 }
 
 impl BlockExecutor {
@@ -77,14 +74,16 @@ impl BlockExecutor {
         message_broker: &Arc<NatsMessageBroker>,
         fuel_streams: &Arc<FuelStreams>,
         telemetry: Arc<Telemetry<Metrics>>,
+        concurrent_tasks: usize,
     ) -> Self {
-        let semaphore = Arc::new(Semaphore::new(MAX_CONCURRENT_TASKS));
+        let semaphore = Arc::new(Semaphore::new(concurrent_tasks));
         Self {
             db,
             semaphore,
             message_broker: message_broker.clone(),
             fuel_streams: fuel_streams.clone(),
             telemetry,
+            concurrent_tasks,
         }
     }
 
@@ -94,22 +93,27 @@ impl BlockExecutor {
     ) -> Result<(), ConsumerError> {
         tracing::info!(
             "Starting consumer with max concurrent tasks: {}",
-            MAX_CONCURRENT_TASKS
+            self.concurrent_tasks
         );
 
         let telemetry = self.telemetry.clone();
         let queue = NatsQueue::BlockImporter(self.message_broker.clone());
+        let mut join_set = JoinSet::new();
 
         while !token.is_cancelled() {
-            let mut messages = queue.subscribe(BATCH_SIZE).await?;
-            let mut join_set = JoinSet::new();
-            while let Some(msg) = messages.next().await {
-                let msg = msg?;
-                self.spawn_processing_tasks(msg, &mut join_set).await?;
-            }
-            // Wait for all spawned tasks to complete before processing next message
-            while let Some(result) = join_set.join_next().await {
-                Self::handle_task_result(result, &telemetry).await?;
+            tokio::select! {
+                msg_result = queue.subscribe(self.concurrent_tasks) => {
+                    let mut messages = msg_result?;
+                    while let Some(msg) = messages.next().await {
+                        let _permit = self.semaphore.acquire().await?;
+                        let msg = msg?;
+                        self.spawn_processing_tasks(msg, &mut join_set,)
+                        .await?;
+                    }
+                }
+                Some(result) = join_set.join_next() => {
+                    Self::handle_task_result(result, &telemetry).await?;
+                }
             }
         }
 
@@ -125,44 +129,45 @@ impl BlockExecutor {
         join_set: &mut JoinSet<Result<ProcessResult, ConsumerError>>,
     ) -> Result<(), ConsumerError> {
         let db = self.db.clone();
-        let semaphore = self.semaphore.clone();
         let fuel_streams = self.fuel_streams.clone();
         let payload = msg.payload();
         let msg_payload = MsgPayload::decode_json(&payload)?.arc();
         let packets = Self::build_packets(&msg_payload);
 
         join_set.spawn({
-            let semaphore = semaphore.clone();
             let packets = packets.clone();
             let msg_payload = msg_payload.clone();
             async move {
-                let _permit = semaphore.acquire().await?;
                 let result = handle_stores(&db, &packets, &msg_payload).await;
                 if let Ok(stats) = result {
                     if stats.error.is_none() {
-                        msg.ack().await.map_err(|e| {
-                            tracing::error!("Failed to ack message: {:?}", e);
-                            ConsumerError::MessageBrokerClient(e)
-                        })?;
+                        tokio::spawn(async move {
+                            let _ = msg.ack().await.map_err(|e| {
+                                tracing::error!(
+                                    "Failed to ack message: {:?}",
+                                    e
+                                );
+                                ConsumerError::MessageBrokerClient(e)
+                            });
+                            tracing::info!(
+                                "[#{}] Message acknowledged",
+                                stats.block_height
+                            );
+                        });
                     }
                     return Ok::<_, ConsumerError>(ProcessResult::Store(Ok(
                         stats,
                     )));
                 }
-                Ok::<_, ConsumerError>(ProcessResult::Store(result))
-            }
-        });
-
-        join_set.spawn({
-            let semaphore = semaphore.clone();
-            let packets = packets.clone();
-            let msg_payload = msg_payload.clone();
-            let fuel_streams = fuel_streams.clone();
-            async move {
-                let _permit = semaphore.acquire_owned().await?;
-                let result =
+                let result_streams =
                     handle_streams(&fuel_streams, &packets, &msg_payload).await;
-                Ok(ProcessResult::Stream(result))
+                if let Ok(stream_stats) = result_streams {
+                    match &stream_stats.error {
+                        Some(error) => stream_stats.log_error(error),
+                        None => stream_stats.log_success(),
+                    }
+                }
+                Ok::<_, ConsumerError>(ProcessResult::Store(result))
             }
         });
 
@@ -185,13 +190,6 @@ impl BlockExecutor {
                     None => store_stats.log_success(),
                 }
             }
-            Ok(Ok(ProcessResult::Stream(stream_result))) => {
-                let stream_stats = stream_result?;
-                match &stream_stats.error {
-                    Some(error) => stream_stats.log_error(error),
-                    None => stream_stats.log_success(),
-                }
-            }
             Ok(Err(e)) => tracing::error!("Task error: {}", e),
             Err(e) => tracing::error!("Task panicked: {}", e),
         }

services/consumer/src/main.rs

@@ -29,10 +29,12 @@ async fn main() -> anyhow::Result<()> {
     let cli = Cli::parse();
     let shutdown = Arc::new(ShutdownController::new());
     shutdown.clone().spawn_signal_handler();
-
     // Initialize shared resources
-    let db = setup_db(&cli.db_url).await?;
-    let message_broker = NatsMessageBroker::setup(&cli.nats_url, None).await?;
+    let db = setup_db(&cli.db_url, cli.concurrent_tasks as u32).await?;
+    // 2 minutes before returning to the message broker
+    let opts = fuel_message_broker::NatsOpts::new(cli.nats_url.clone())
+        .with_ack_wait(120);
+    let message_broker = NatsMessageBroker::setup_with_opts(&opts).await?;
     let metrics = Metrics::new(None)?;
     let telemetry = Telemetry::new(Some(metrics)).await?;
     telemetry.start().await?;
@@ -48,6 +50,7 @@ async fn main() -> anyhow::Result<()> {
         &message_broker,
         &fuel_streams,
         Arc::clone(&telemetry),
+        cli.concurrent_tasks,
     );
 
     tokio::select! {
@@ -70,9 +73,13 @@ async fn main() -> anyhow::Result<()> {
     Ok(())
 }
 
-async fn setup_db(db_url: &str) -> Result<Arc<Db>, ConsumerError> {
+async fn setup_db(
+    db_url: &str,
+    concurrent_tasks: u32,
+) -> Result<Arc<Db>, ConsumerError> {
     let db = Db::new(DbConnectionOpts {
         connection_str: db_url.to_string(),
+        pool_size: Some(concurrent_tasks),
         ..Default::default()
     })
     .await?;

services/publisher/src/history.rs

@@ -4,7 +4,10 @@ use fuel_message_broker::NatsMessageBroker;
 use fuel_streams_core::types::*;
 use fuel_streams_domains::infra::Db;
 use fuel_web_utils::{shutdown::ShutdownController, telemetry::Telemetry};
-use tokio::time::{interval, Duration};
+use tokio::{
+    task::JoinSet,
+    time::{interval, Duration},
+};
 use tokio_util::sync::CancellationToken;
 
 use crate::{
@@ -141,9 +144,9 @@ async fn process_blocks(
     'outer: for gap in processed_gaps {
         tracing::info!("Processing gap: {} to {}", gap.start, gap.end);
         let heights = (*gap.start..=*gap.end).collect::<Vec<_>>();
-        let heights_chucks = heights.chunks(50);
+        let heights_chunks = heights.chunks(100);
 
-        for heights in heights_chucks {
+        for heights in heights_chunks {
             tracing::info!(
                 "Processing chunk: {:?} - {:?}",
                 heights.first(),
@@ -154,47 +157,56 @@ async fn process_blocks(
                 break 'outer;
             }
 
+            let mut join_set: JoinSet<Result<Option<u32>, PublishError>> =
+                JoinSet::new();
             for height in heights {
                 tracing::info!("Processing block height: {}", height);
-
-                if token.is_cancelled() {
-                    break 'outer;
-                }
-
                 let message_broker = message_broker.clone();
                 let fuel_core = fuel_core.clone();
                 let telemetry = telemetry.clone();
+                let block_height = *height;
                 let sealed_block =
                     fuel_core.get_sealed_block((*height).into())?;
                 let sealed_block = Arc::new(sealed_block);
+                let token = token.clone();
 
-                tracing::info!("Publishing block height: {}", height);
-                let result = publish_block(
-                    &message_broker,
-                    &fuel_core,
-                    &sealed_block,
-                    &telemetry,
-                    None,
-                )
-                .await;
-
-                match result {
-                    Ok(_) => {
-                        tracing::info!(
-                            "Block {} published successfully",
-                            height
-                        );
+                join_set.spawn(async move {
+                    if token.is_cancelled() {
+                        return Ok(None);
                     }
-                    Err(e) => {
+                    tracing::info!("Publishing block height: {}", block_height);
+                    publish_block(
+                        &message_broker,
+                        &fuel_core,
+                        &sealed_block,
+                        &telemetry,
+                        None,
+                    )
+                    .await
+                    .map(|_| Some(block_height))
+                    .map_err(|e| {
                         tracing::error!(
                             "Error publishing block {}: {:?}",
-                            height,
+                            block_height,
                             e
                         );
-                        return Err(e.into());
-                    }
+                        e
+                    })
+                });
+            }
+            while let Some(result) = join_set.join_next().await {
+                if let Ok(Ok(Some(block_height))) = result {
+                    tracing::info!(
+                        "Block {} published successfully",
+                        block_height
+                    );
                 }
             }
+            tracing::info!(
+                "Finished processing chunk: {:?} - {:?}",
+                heights.first(),
+                heights.last()
+            );
         }
     }
 

tests/tests/services/consumer.rs

@@ -287,8 +287,13 @@ async fn test_consumer_inserting_records() -> anyhow::Result<()> {
         let message_broker = Arc::clone(&message_broker);
         let fuel_streams = Arc::clone(&fuel_streams);
         let telemetry = Telemetry::new(None).await?;
-        let block_executor =
-            BlockExecutor::new(db, &message_broker, &fuel_streams, telemetry);
+        let block_executor = BlockExecutor::new(
+            db,
+            &message_broker,
+            &fuel_streams,
+            telemetry,
+            100,
+        );
         async move { block_executor.start(shutdown.token()).await }
     });