Pass Database into consumers; StreamReader conn

Create a single Database instance for each consumer group and pass it into consumer run functions and ChainConsumerRunner instead of creating new DB clients inside each task. Update ChainConsumerRunner::new to accept a Database and use runner.database in consumers (store, indexer, rewards). Add optional StreamConnection support to StreamReader (from_connection, create_channel, try_connect_shared) so readers can be created from an existing connection and reconnect uses the shared connection when available. Adjust imports and reader creation sites accordingly.

Author: gemcoder21

apps/daemon/src/consumers/indexer/mod.rs

@@ -12,10 +12,10 @@ use cacher::CacherClient;
 use primitives::{Chain, NFTChain};
 use settings::Settings;
 use storage::Database;
-use streamer::{ChainAddressPayload, ConsumerStatusReporter, FetchAssetsPayload, FetchBlocksPayload, QueueName, ShutdownReceiver, run_consumer};
+use streamer::{ChainAddressPayload, ConsumerStatusReporter, FetchAssetsPayload, FetchBlocksPayload, QueueName, ShutdownReceiver, StreamConnection, StreamReader, run_consumer};
 
 use crate::consumers::runner::ChainConsumerRunner;
-use crate::consumers::{chain_providers, consumer_config, reader_for_queue};
+use crate::consumers::{chain_providers, consumer_config, reader_config};
 
 use fetch_address_transactions_consumer::FetchAddressTransactionsConsumer;
 use fetch_assets_consumer::FetchAssetsConsumer;
@@ -25,23 +25,34 @@ use fetch_nft_assets_addresses_consumer::FetchNftAssetsAddressesConsumer;
 use fetch_token_addresses_consumer::FetchTokenAddressesConsumer;
 
 pub async fn run_consumer_indexer(settings: Settings, shutdown_rx: ShutdownReceiver, reporter: Arc<dyn ConsumerStatusReporter>) -> Result<(), Box<dyn Error + Send + Sync>> {
+    let database = Database::new(&settings.postgres.url, settings.postgres.pool);
     let settings = Arc::new(settings);
 
     futures::future::try_join_all(vec![
-        tokio::spawn(run_fetch_blocks(settings.clone(), shutdown_rx.clone(), reporter.clone())),
-        tokio::spawn(run_fetch_assets(settings.clone(), shutdown_rx.clone(), reporter.clone())),
-        tokio::spawn(run_fetch_token_associations(settings.clone(), shutdown_rx.clone(), reporter.clone())),
-        tokio::spawn(run_fetch_coin_associations(settings.clone(), shutdown_rx.clone(), reporter.clone())),
-        tokio::spawn(run_fetch_nft_associations(settings.clone(), shutdown_rx.clone(), reporter.clone())),
-        tokio::spawn(run_fetch_transaction_associations(settings.clone(), shutdown_rx.clone(), reporter.clone())),
+        tokio::spawn(run_fetch_blocks(settings.clone(), database.clone(), shutdown_rx.clone(), reporter.clone())),
+        tokio::spawn(run_fetch_assets(settings.clone(), database.clone(), shutdown_rx.clone(), reporter.clone())),
+        tokio::spawn(run_fetch_token_associations(settings.clone(), database.clone(), shutdown_rx.clone(), reporter.clone())),
+        tokio::spawn(run_fetch_coin_associations(settings.clone(), database.clone(), shutdown_rx.clone(), reporter.clone())),
+        tokio::spawn(run_fetch_nft_associations(settings.clone(), database.clone(), shutdown_rx.clone(), reporter.clone())),
+        tokio::spawn(run_fetch_transaction_associations(
+            settings.clone(),
+            database.clone(),
+            shutdown_rx.clone(),
+            reporter.clone(),
+        )),
     ])
     .await?;
 
     Ok(())
 }
 
-async fn run_fetch_blocks(settings: Arc<Settings>, shutdown_rx: ShutdownReceiver, reporter: Arc<dyn ConsumerStatusReporter>) -> Result<(), Box<dyn Error + Send + Sync>> {
-    ChainConsumerRunner::new((*settings).clone(), QueueName::FetchBlocks, shutdown_rx, reporter)
+async fn run_fetch_blocks(
+    settings: Arc<Settings>,
+    database: Database,
+    shutdown_rx: ShutdownReceiver,
+    reporter: Arc<dyn ConsumerStatusReporter>,
+) -> Result<(), Box<dyn Error + Send + Sync>> {
+    ChainConsumerRunner::new((*settings).clone(), database, QueueName::FetchBlocks, shutdown_rx, reporter)
         .await?
         .run(|runner, chain| async move {
             let queue = QueueName::FetchBlocks;
@@ -64,10 +75,17 @@ async fn run_fetch_blocks(settings: Arc<Settings>, shutdown_rx: ShutdownReceiver
         .await
 }
 
-async fn run_fetch_assets(settings: Arc<Settings>, shutdown_rx: ShutdownReceiver, reporter: Arc<dyn ConsumerStatusReporter>) -> Result<(), Box<dyn Error + Send + Sync>> {
-    let database = Database::new(&settings.postgres.url, settings.postgres.pool);
+async fn run_fetch_assets(
+    settings: Arc<Settings>,
+    database: Database,
+    shutdown_rx: ShutdownReceiver,
+    reporter: Arc<dyn ConsumerStatusReporter>,
+) -> Result<(), Box<dyn Error + Send + Sync>> {
     let queue = QueueName::FetchAssets;
-    let (name, stream_reader) = reader_for_queue(&settings, &queue).await?;
+    let name = queue.to_string();
+    let connection = StreamConnection::new(&settings.rabbitmq.url, name.clone()).await?;
+    let config = reader_config(&settings.rabbitmq, name.clone());
+    let stream_reader = StreamReader::from_connection(&connection, config).await?;
     let cacher = CacherClient::new(&settings.redis.url).await;
     let consumer = FetchAssetsConsumer {
         providers: chain_providers(&settings, &name),
@@ -79,10 +97,11 @@ async fn run_fetch_assets(settings: Arc<Settings>, shutdown_rx: ShutdownReceiver
 
 async fn run_fetch_token_associations(
     settings: Arc<Settings>,
+    database: Database,
     shutdown_rx: ShutdownReceiver,
     reporter: Arc<dyn ConsumerStatusReporter>,
 ) -> Result<(), Box<dyn Error + Send + Sync>> {
-    ChainConsumerRunner::new((*settings).clone(), QueueName::FetchTokenAssociations, shutdown_rx, reporter)
+    ChainConsumerRunner::new((*settings).clone(), database, QueueName::FetchTokenAssociations, shutdown_rx, reporter)
         .await?
         .run(|runner, chain| async move {
             let queue = QueueName::FetchTokenAssociations;
@@ -107,10 +126,11 @@ async fn run_fetch_token_associations(
 
 async fn run_fetch_coin_associations(
     settings: Arc<Settings>,
+    database: Database,
     shutdown_rx: ShutdownReceiver,
     reporter: Arc<dyn ConsumerStatusReporter>,
 ) -> Result<(), Box<dyn Error + Send + Sync>> {
-    ChainConsumerRunner::new((*settings).clone(), QueueName::FetchCoinAssociations, shutdown_rx, reporter)
+    ChainConsumerRunner::new((*settings).clone(), database, QueueName::FetchCoinAssociations, shutdown_rx, reporter)
         .await?
         .run(|runner, chain| async move {
             let queue = QueueName::FetchCoinAssociations;
@@ -132,9 +152,14 @@ async fn run_fetch_coin_associations(
         .await
 }
 
-async fn run_fetch_nft_associations(settings: Arc<Settings>, shutdown_rx: ShutdownReceiver, reporter: Arc<dyn ConsumerStatusReporter>) -> Result<(), Box<dyn Error + Send + Sync>> {
+async fn run_fetch_nft_associations(
+    settings: Arc<Settings>,
+    database: Database,
+    shutdown_rx: ShutdownReceiver,
+    reporter: Arc<dyn ConsumerStatusReporter>,
+) -> Result<(), Box<dyn Error + Send + Sync>> {
     let chains: Vec<Chain> = NFTChain::all().into_iter().map(Into::into).collect();
-    ChainConsumerRunner::new((*settings).clone(), QueueName::FetchNftAssociations, shutdown_rx, reporter)
+    ChainConsumerRunner::new((*settings).clone(), database, QueueName::FetchNftAssociations, shutdown_rx, reporter)
         .await?
         .run_for_chains(chains, |runner, chain| async move {
             FetchNftAssetsAddressesConsumer::run(
@@ -154,10 +179,11 @@ async fn run_fetch_nft_associations(settings: Arc<Settings>, shutdown_rx: Shutdo
 
 async fn run_fetch_transaction_associations(
     settings: Arc<Settings>,
+    database: Database,
     shutdown_rx: ShutdownReceiver,
     reporter: Arc<dyn ConsumerStatusReporter>,
 ) -> Result<(), Box<dyn Error + Send + Sync>> {
-    ChainConsumerRunner::new((*settings).clone(), QueueName::FetchAddressTransactions, shutdown_rx, reporter)
+    ChainConsumerRunner::new((*settings).clone(), database, QueueName::FetchAddressTransactions, shutdown_rx, reporter)
         .await?
         .run(|runner, chain| async move {
             let queue = QueueName::FetchAddressTransactions;

apps/daemon/src/consumers/rewards/mod.rs

@@ -20,19 +20,24 @@ use streamer::{ConsumerStatusReporter, QueueName, RewardsNotificationPayload, Re
 use crate::consumers::{consumer_config, producer_for_queue, reader_for_queue};
 
 pub async fn run_consumer_rewards(settings: Settings, shutdown_rx: ShutdownReceiver, reporter: Arc<dyn ConsumerStatusReporter>) -> Result<(), Box<dyn Error + Send + Sync>> {
+    let database = Database::new(&settings.postgres.url, settings.postgres.pool);
     let settings = Arc::new(settings);
 
     futures::future::try_join_all(vec![
-        tokio::spawn(run_rewards_events(settings.clone(), shutdown_rx.clone(), reporter.clone())),
-        tokio::spawn(run_rewards_redemptions(settings.clone(), shutdown_rx.clone(), reporter.clone())),
+        tokio::spawn(run_rewards_events(settings.clone(), database.clone(), shutdown_rx.clone(), reporter.clone())),
+        tokio::spawn(run_rewards_redemptions(settings.clone(), database.clone(), shutdown_rx.clone(), reporter.clone())),
     ])
     .await?;
 
     Ok(())
 }
 
-async fn run_rewards_events(settings: Arc<Settings>, shutdown_rx: ShutdownReceiver, reporter: Arc<dyn ConsumerStatusReporter>) -> Result<(), Box<dyn Error + Send + Sync>> {
-    let database = Database::new(&settings.postgres.url, settings.postgres.pool);
+async fn run_rewards_events(
+    settings: Arc<Settings>,
+    database: Database,
+    shutdown_rx: ShutdownReceiver,
+    reporter: Arc<dyn ConsumerStatusReporter>,
+) -> Result<(), Box<dyn Error + Send + Sync>> {
     let queue = QueueName::RewardsEvents;
     let (name, stream_reader) = reader_for_queue(&settings, &queue).await?;
     let stream_producer = producer_for_queue(&settings, &name).await?;
@@ -41,8 +46,12 @@ async fn run_rewards_events(settings: Arc<Settings>, shutdown_rx: ShutdownReceiv
     run_consumer::<RewardsNotificationPayload, rewards_consumer::RewardsConsumer, usize>(&name, stream_reader, queue, None, consumer, consumer_config, shutdown_rx, reporter).await
 }
 
-async fn run_rewards_redemptions(settings: Arc<Settings>, shutdown_rx: ShutdownReceiver, reporter: Arc<dyn ConsumerStatusReporter>) -> Result<(), Box<dyn Error + Send + Sync>> {
-    let database = Database::new(&settings.postgres.url, settings.postgres.pool);
+async fn run_rewards_redemptions(
+    settings: Arc<Settings>,
+    database: Database,
+    shutdown_rx: ShutdownReceiver,
+    reporter: Arc<dyn ConsumerStatusReporter>,
+) -> Result<(), Box<dyn Error + Send + Sync>> {
     let config = ConfigCacher::new(database.clone());
     let retry_config = rewards_redemption_consumer::RedemptionRetryConfig {
         max_retries: config.get_i64(ConfigKey::RedemptionRetryMaxRetries)? as u32,

apps/daemon/src/consumers/runner.rs

@@ -25,11 +25,11 @@ pub struct ChainConsumerRunner {
 impl ChainConsumerRunner {
     pub async fn new(
         settings: Settings,
+        database: Database,
         queue: streamer::QueueName,
         shutdown_rx: ShutdownReceiver,
         reporter: Arc<dyn ConsumerStatusReporter>,
     ) -> Result<Self, Box<dyn Error + Send + Sync>> {
-        let database = Database::new(&settings.postgres.url, settings.postgres.pool);
         let connection = StreamConnection::new(&settings.rabbitmq.url, queue.to_string()).await?;
         let cacher = CacherClient::new(&settings.redis.url).await;
         let config = consumer_config(&settings.consumer);

apps/daemon/src/consumers/store/mod.rs

@@ -26,33 +26,38 @@ use store_charts_consumer::StoreChartsConsumer;
 use store_prices_consumer::StorePricesConsumer;
 
 pub async fn run_consumer_store(settings: Settings, shutdown_rx: ShutdownReceiver, reporter: Arc<dyn ConsumerStatusReporter>) -> Result<(), Box<dyn Error + Send + Sync>> {
+    let database = Database::new(&settings.postgres.url, settings.postgres.pool);
     let settings = Arc::new(settings);
 
     futures::future::try_join_all(vec![
-        tokio::spawn(run_store_transactions(settings.clone(), shutdown_rx.clone(), reporter.clone())),
-        tokio::spawn(run_store_prices(settings.clone(), shutdown_rx.clone(), reporter.clone())),
-        tokio::spawn(run_store_charts(settings.clone(), shutdown_rx.clone(), reporter.clone())),
+        tokio::spawn(run_store_transactions(settings.clone(), database.clone(), shutdown_rx.clone(), reporter.clone())),
+        tokio::spawn(run_store_prices(settings.clone(), database.clone(), shutdown_rx.clone(), reporter.clone())),
+        tokio::spawn(run_store_charts(settings.clone(), database.clone(), shutdown_rx.clone(), reporter.clone())),
         tokio::spawn(run_device_stream(settings.clone(), shutdown_rx.clone(), reporter.clone())),
     ])
     .await?;
 
     Ok(())
 }
 
-async fn run_store_transactions(settings: Arc<Settings>, shutdown_rx: ShutdownReceiver, reporter: Arc<dyn ConsumerStatusReporter>) -> Result<(), Box<dyn Error + Send + Sync>> {
-    ChainConsumerRunner::new((*settings).clone(), QueueName::StoreTransactions, shutdown_rx, reporter)
+async fn run_store_transactions(
+    settings: Arc<Settings>,
+    database: Database,
+    shutdown_rx: ShutdownReceiver,
+    reporter: Arc<dyn ConsumerStatusReporter>,
+) -> Result<(), Box<dyn Error + Send + Sync>> {
+    ChainConsumerRunner::new((*settings).clone(), database, QueueName::StoreTransactions, shutdown_rx, reporter)
         .await?
         .run(|runner, chain| async move {
             let queue = QueueName::StoreTransactions;
             let name = format!("{}.{}", queue, chain.as_ref());
             let stream_reader = runner.stream_reader().await?;
             let stream_producer = runner.stream_producer().await?;
-            let database = Database::new(&runner.settings.postgres.url, runner.settings.postgres.pool);
             let consumer = StoreTransactionsConsumer {
-                database: database.clone(),
-                config_cacher: ConfigCacher::new(database.clone()),
+                database: runner.database.clone(),
+                config_cacher: ConfigCacher::new(runner.database.clone()),
                 stream_producer,
-                pusher: Pusher::new(database.clone()),
+                pusher: Pusher::new(runner.database.clone()),
                 config: StoreTransactionsConsumerConfig {},
             };
             run_consumer::<TransactionsPayload, StoreTransactionsConsumer, usize>(
@@ -70,8 +75,12 @@ async fn run_store_transactions(settings: Arc<Settings>, shutdown_rx: ShutdownRe
         .await
 }
 
-async fn run_store_prices(settings: Arc<Settings>, shutdown_rx: ShutdownReceiver, reporter: Arc<dyn ConsumerStatusReporter>) -> Result<(), Box<dyn Error + Send + Sync>> {
-    let database = Database::new(&settings.postgres.url, settings.postgres.pool);
+async fn run_store_prices(
+    settings: Arc<Settings>,
+    database: Database,
+    shutdown_rx: ShutdownReceiver,
+    reporter: Arc<dyn ConsumerStatusReporter>,
+) -> Result<(), Box<dyn Error + Send + Sync>> {
     let queue = QueueName::StorePrices;
     let (name, stream_reader) = reader_for_queue(&settings, &queue).await?;
     let cacher_client = CacherClient::new(&settings.redis.url).await;
@@ -82,8 +91,12 @@ async fn run_store_prices(settings: Arc<Settings>, shutdown_rx: ShutdownReceiver
     run_consumer::<PricesPayload, StorePricesConsumer, usize>(&name, stream_reader, queue, None, consumer, consumer_config(&settings.consumer), shutdown_rx, reporter).await
 }
 
-async fn run_store_charts(settings: Arc<Settings>, shutdown_rx: ShutdownReceiver, reporter: Arc<dyn ConsumerStatusReporter>) -> Result<(), Box<dyn Error + Send + Sync>> {
-    let database = Database::new(&settings.postgres.url, settings.postgres.pool);
+async fn run_store_charts(
+    settings: Arc<Settings>,
+    database: Database,
+    shutdown_rx: ShutdownReceiver,
+    reporter: Arc<dyn ConsumerStatusReporter>,
+) -> Result<(), Box<dyn Error + Send + Sync>> {
     let queue = QueueName::StoreCharts;
     let (name, stream_reader) = reader_for_queue(&settings, &queue).await?;
     let cacher_client = CacherClient::new(&settings.redis.url).await;

crates/streamer/src/stream_reader.rs

@@ -27,12 +27,17 @@ impl StreamReaderConfig {
 pub struct StreamReader {
     config: StreamReaderConfig,
     channel: Channel,
+    connection: Option<StreamConnection>,
 }
 
 impl StreamReader {
     pub async fn new(config: StreamReaderConfig) -> Result<Self, Box<dyn Error + Send + Sync>> {
         let channel = with_retry(&config.retry, &config.name, || Self::try_connect(&config)).await?;
-        Ok(Self { config, channel })
+        Ok(Self {
+            config,
+            channel,
+            connection: None,
+        })
     }
 
     pub async fn from_connection(connection: &StreamConnection, config: StreamReaderConfig) -> Result<Self, Box<dyn Error + Send + Sync>> {
@@ -41,8 +46,12 @@ impl StreamReader {
             name: connection.name().to_string(),
             ..config
         };
-        let channel = with_retry(&config.retry, &config.name, || Self::try_connect(&config)).await?;
-        Ok(Self { config, channel })
+        let channel = Self::create_channel(connection, config.prefetch).await?;
+        Ok(Self {
+            config,
+            channel,
+            connection: Some(connection.clone()),
+        })
     }
 
     async fn try_connect(config: &StreamReaderConfig) -> Result<Channel, Box<dyn Error + Send + Sync>> {
@@ -58,6 +67,19 @@ impl StreamReader {
         Ok(())
     }
 
+    async fn create_channel(connection: &StreamConnection, prefetch: u16) -> Result<Channel, Box<dyn Error + Send + Sync>> {
+        let channel = connection.create_channel().await?;
+        channel.basic_qos(prefetch, BasicQosOptions { global: false }).await?;
+        Ok(channel)
+    }
+
+    async fn try_connect_shared(&self) -> Result<Channel, Box<dyn Error + Send + Sync>> {
+        if let Some(ref conn) = self.connection {
+            return Self::create_channel(conn, self.config.prefetch).await;
+        }
+        Self::try_connect(&self.config).await
+    }
+
     async fn reconnect(&mut self, shutdown_rx: &ShutdownReceiver) -> Result<bool, Box<dyn Error + Send + Sync>> {
         let mut delay = self.config.retry.delay;
         let mut attempt: u32 = 0;
@@ -66,7 +88,7 @@ impl StreamReader {
                 return Ok(false);
             }
             attempt += 1;
-            match Self::try_connect(&self.config).await {
+            match self.try_connect_shared().await {
                 Ok(channel) => {
                     self.channel = channel;
                     info_with_fields!("rabbitmq reconnected", connection = self.config.name.as_str(), attempt = attempt);