Fix tokio task leak

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pyth-agent"
-version = "2.12.1"
+version = "2.12.2"
 edition = "2021"
 
 [[bin]]

config/config.toml

@@ -73,6 +73,15 @@ key_store.pyth_price_store_program_key = "3m6sv6HGqEbuyLV84mD7rJn4MAC9LhUa1y1AUN
 # takes to fetch all symbols.
 # oracle.max_lookup_batch_size = 100
 
+# Number of workers used to wait for the handle_price_account_update
+# oracle.handle_price_account_update_worker_poll_size = 25
+# Channel size used to wait for the handle_price_account_update
+# oracle.handle_price_account_update_channel_size = 1000
+# Minimum time for a subscriber to run
+# oracle.subscriber_finished_min_time = "30s"
+# Time to sleep if the subscriber do not run for more than the minimum time
+# oracle.subscriber_finished_sleep_time = "1s"
+
 # How often to refresh the cached network state (current slot and blockhash).
 # It is recommended to set this to slightly less than the network's block time,
 # as the slot fetched will be used as the time of the price update.

src/agent/services/oracle.rs

@@ -13,7 +13,10 @@ use {
         },
         state::oracle::Oracle,
     },
-    anyhow::Result,
+    anyhow::{
+        Context,
+        Result,
+    },
     solana_account_decoder::UiAccountEncoding,
     solana_client::{
         nonblocking::{
@@ -67,6 +70,12 @@ where
     )));
 
     if config.oracle.subscriber_enabled {
+        let number_of_workers = config.oracle.handle_price_account_update_worker_poll_size;
+        let (sender, receiver) =
+            tokio::sync::mpsc::channel(config.oracle.handle_price_account_update_channel_size);
+        let min_elapsed_time = config.oracle.subscriber_finished_min_time;
+        let sleep_time = config.oracle.subscriber_finished_sleep_time;
+
         handles.push(tokio::spawn(async move {
             loop {
                 let current_time = Instant::now();
@@ -75,17 +84,34 @@ where
                     network,
                     state.clone(),
                     key_store.pyth_oracle_program_key,
+                    sender.clone(),
                 )
                 .await
                 {
-                    tracing::error!(err = ?err, "Subscriber exited unexpectedly.");
-                    if current_time.elapsed() < Duration::from_secs(30) {
-                        tracing::warn!("Subscriber restarting too quickly. Sleeping for 1 second.");
-                        tokio::time::sleep(Duration::from_secs(1)).await;
+                    tracing::error!(?err, "Subscriber exited unexpectedly");
+                    if current_time.elapsed() < min_elapsed_time {
+                        tracing::warn!(?sleep_time, "Subscriber restarting too quickly. Sleeping");
+                        tokio::time::sleep(sleep_time).await;
                     }
                 }
             }
         }));
+
+        let receiver = Arc::new(tokio::sync::Mutex::new(receiver));
+        for _ in 0..number_of_workers {
+            let receiver = receiver.clone();
+            handles.push(tokio::spawn(async move {
+                loop {
+                    let mut receiver = receiver.lock().await;
+                    if let Some(task) = receiver.recv().await {
+                        drop(receiver);
+                        if let Err(err) = task.await {
+                            tracing::error!(%err, "error running price update");
+                        }
+                    }
+                }
+            }));
+        }
     }
 
     handles
@@ -102,6 +128,7 @@ async fn subscriber<S>(
     network: Network,
     state: Arc<S>,
     program_key: Pubkey,
+    sender: tokio::sync::mpsc::Sender<tokio::task::JoinHandle<()>>,
 ) -> Result<()>
 where
     S: Oracle,
@@ -129,14 +156,17 @@ where
             Some(account) => {
                 let pubkey: Pubkey = update.value.pubkey.as_str().try_into()?;
                 let state = state.clone();
-                tokio::spawn(async move {
-                    if let Err(err) =
-                        Oracle::handle_price_account_update(&*state, network, &pubkey, &account)
-                            .await
-                    {
-                        tracing::error!(err = ?err, "Failed to handle account update.");
-                    }
-                });
+                sender
+                    .send(tokio::spawn(async move {
+                        if let Err(err) =
+                            Oracle::handle_price_account_update(&*state, network, &pubkey, &account)
+                                .await
+                        {
+                            tracing::error!(?err, "Failed to handle account update");
+                        }
+                    }))
+                    .await
+                    .context("sending handle_price_account_update task to worker")?;
             }
 
             None => {

src/agent/state/oracle.rs

@@ -138,6 +138,22 @@ pub struct Data {
     pub publisher_buffer_key:  Option<Pubkey>,
 }
 
+fn default_handle_price_account_update_channel_size() -> usize {
+    1000
+}
+
+fn default_handle_price_account_update_worker_poll_size() -> usize {
+    25
+}
+
+fn default_subscriber_finished_min_time() -> Duration {
+    Duration::from_secs(30)
+}
+
+fn default_subscriber_finished_sleep_time() -> Duration {
+    Duration::from_secs(1)
+}
+
 #[derive(Clone, Serialize, Deserialize, Debug)]
 #[serde(default)]
 pub struct Config {
@@ -159,17 +175,36 @@ pub struct Config {
     /// socket count at bay, the batches are looked up sequentially,
     /// trading off overall time it takes to fetch all symbols.
     pub max_lookup_batch_size: usize,
+
+    /// Number of workers used to wait for the handle_price_account_update
+    #[serde(default = "default_handle_price_account_update_worker_poll_size")]
+    pub handle_price_account_update_worker_poll_size: usize,
+    /// Channel size used to wait for the handle_price_account_update
+    #[serde(default = "default_handle_price_account_update_channel_size")]
+    pub handle_price_account_update_channel_size:     usize,
+    /// Minimum time for a subscriber to run
+    #[serde(default = "default_subscriber_finished_min_time")]
+    pub subscriber_finished_min_time:                 Duration,
+    /// Time to sleep if the subscriber do not run for more than the minimum time
+    #[serde(default = "default_subscriber_finished_sleep_time")]
+    pub subscriber_finished_sleep_time:               Duration,
 }
 
 impl Default for Config {
     fn default() -> Self {
         Self {
-            commitment:               CommitmentLevel::Confirmed,
-            poll_interval_duration:   Duration::from_secs(5),
-            subscriber_enabled:       true,
-            updates_channel_capacity: 10000,
-            data_channel_capacity:    10000,
-            max_lookup_batch_size:    100,
+            commitment:                                   CommitmentLevel::Confirmed,
+            poll_interval_duration:                       Duration::from_secs(5),
+            subscriber_enabled:                           true,
+            updates_channel_capacity:                     10000,
+            data_channel_capacity:                        10000,
+            max_lookup_batch_size:                        100,
+            handle_price_account_update_worker_poll_size:
+                default_handle_price_account_update_worker_poll_size(),
+            handle_price_account_update_channel_size:
+                default_handle_price_account_update_channel_size(),
+            subscriber_finished_min_time:                 default_subscriber_finished_min_time(),
+            subscriber_finished_sleep_time:               default_subscriber_finished_sleep_time(),
         }
     }
 }
@@ -241,6 +276,7 @@ where
         );
 
         data.price_accounts.insert(*account_key, price_entry.into());
+        drop(data);
 
         Prices::update_global_price(
             self,
@@ -333,13 +369,16 @@ where
         let mut data = self.into().data.write().await;
         log_data_diff(&data, &new_data);
         *data = new_data;
+        let data_publisher_permissions = data.publisher_permissions.clone();
+        let data_publisher_buffer_key = data.publisher_buffer_key;
+        drop(data);
 
         Exporter::update_on_chain_state(
             self,
             network,
             publish_keypair,
-            data.publisher_permissions.clone(),
-            data.publisher_buffer_key,
+            data_publisher_permissions,
+            data_publisher_buffer_key,
         )
         .await?;