feat: enhance subscription management with metrics and reconnection

Cargo.toml

@@ -91,7 +91,7 @@ jsonrpc-pubsub = "18.0.0"
 jsonrpc-ws-server = "18.0.0"
 lazy_static = "1.4.0"
 libc = "0.2.153"
-log = { version = "0.4.20", features = ["release_max_level_info"] }
+log = { version = "0.4.20" }
 lru = "0.16.0"
 macrotest = "1"
 magic-domain-program = { git = "https://github.com/magicblock-labs/magic-domain-program.git", rev = "ea04d46", default-features = false }

magicblock-accounts-db/src/lib.rs

@@ -1,4 +1,4 @@
-use std::{path::Path, sync::Arc};
+use std::{collections::HashSet, path::Path, sync::Arc};
 
 use error::AccountsDbError;
 use index::{
@@ -356,7 +356,7 @@ impl AccountsBank for AccountsDb {
             .iter_all()
             .filter(|(pk, acc)| predicate(pk, acc))
             .map(|(pk, _)| pk)
-            .collect::<Vec<_>>();
+            .collect::<HashSet<_>>();
         let removed = to_remove.len();
         for pk in to_remove {
             self.remove_account(&pk);

magicblock-aperture/src/requests/http/mod.rs

@@ -112,7 +112,7 @@ impl HttpDispatcher {
             .inspect_err(|e| {
                 // There is nothing we can do if fetching the account fails
                 // Log the error and return whatever is in the accounts db
-                warn!("Failed to ensure account {pubkey}: {e}");
+                debug!("Failed to ensure account {pubkey}: {e}");
             });
         self.accountsdb.get_account(pubkey)
     }

magicblock-aperture/src/requests/http/send_transaction.rs

@@ -1,4 +1,4 @@
-use log::{debug, trace};
+use log::*;
 use magicblock_metrics::metrics::{
     TRANSACTION_PROCESSING_TIME, TRANSACTION_SKIP_PREFLIGHT,
 };

magicblock-chainlink/Cargo.toml

@@ -4,6 +4,7 @@ version.workspace = true
 edition.workspace = true
 
 [dependencies]
+arc-swap = "1.7"
 async-trait = { workspace = true }
 bincode = { workspace = true }
 env_logger = { workspace = true }
@@ -12,7 +13,8 @@ log = { workspace = true }
 lru = { workspace = true }
 magicblock-core = { workspace = true }
 magicblock-magic-program-api = { workspace = true }
-magicblock-delegation-program = { workspace = true }
+magicblock-metrics = { workspace = true }
+ magicblock-delegation-program = { workspace = true }
 serde_json = { workspace = true }
 solana-account = { workspace = true }
 solana-account-decoder = { workspace = true }

magicblock-chainlink/src/chainlink/errors.rs

@@ -18,7 +18,7 @@ pub enum ChainlinkError {
     #[error("Cloner error: {0}")]
     ClonerError(#[from] crate::cloner::errors::ClonerError),
 
-    #[error("Delegation could not be decoded: {0} ({1:?})")]
+    #[error("Delegation record could not be decoded: {0} ({1:?})")]
     InvalidDelegationRecord(Pubkey, ProgramError),
 
     #[error("Failed to resolve one or more accounts {0} when getting delegation records")]

magicblock-chainlink/src/chainlink/fetch_cloner.rs

@@ -581,7 +581,7 @@ where
         // For accounts we couldn't find we cannot do anything. We will let code depending
         // on them to be in the bank fail on its own
         if !not_found.is_empty() {
-            debug!(
+            trace!(
                 "Could not find accounts on chain: {:?}",
                 not_found
                     .iter()
@@ -960,24 +960,37 @@ where
                 .lock()
                 .expect("pending_requests lock poisoned");
 
-            for &pubkey in pubkeys {
-                // Check synchronously if account is in bank
-                if self.accounts_bank.get_account(&pubkey).is_some() {
-                    // Account is already in bank, we can skip it as it will be handled
-                    // by the existing fetch_and_clone_accounts logic when needed
-                    continue;
+            for pubkey in pubkeys {
+                // Check synchronously if account is in bank and subscribed when it should be
+                if let Some(account_in_bank) =
+                    self.accounts_bank.get_account(pubkey)
+                {
+                    // NOTE: we defensively correct accounts that we should have been watching but
+                    //       were not for some reason. We fetch them again in that case.
+                    //       This actually would point to a bug in the subscription logic.
+                    // TODO(thlorenz): remove this once we are certain (by perusing logs) that this
+                    //                 does not happen anymore
+                    if account_in_bank.owner().eq(&dlp::id())
+                        || account_in_bank.delegated()
+                        || self.blacklisted_accounts.contains(pubkey)
+                        || self.is_watching(pubkey)
+                    {
+                        continue;
+                    } else if !self.is_watching(pubkey) {
+                        debug!("Account {pubkey} should be watched but wasn't");
+                    }
                 }
 
                 // Check if account fetch is already pending
-                if let Some(requests) = pending.get_mut(&pubkey) {
+                if let Some(requests) = pending.get_mut(pubkey) {
                     let (sender, receiver) = oneshot::channel();
                     requests.push(sender);
-                    await_pending.push((pubkey, receiver));
+                    await_pending.push((*pubkey, receiver));
                     continue;
                 }
 
                 // Account needs to be fetched - add to fetch list
-                fetch_new.push(pubkey);
+                fetch_new.push(*pubkey);
             }
 
             // Create pending entries for accounts we need to fetch
@@ -1024,9 +1037,14 @@ where
 
         // Wait for any pending requests to complete
         let mut joinset = JoinSet::new();
-        for (_, receiver) in await_pending {
+        for (pubkey, receiver) in await_pending {
             joinset.spawn(async move {
-                if let Err(err) = receiver.await {
+                if let Err(err) = receiver
+                    .await
+                    .inspect_err(|err| {
+                        warn!("FetchCloner::clone_accounts - RecvError occurred while awaiting account {}: {err:?}. This indicates the account fetch sender was dropped without sending a value.", pubkey);
+                    })
+                {
                     // The sender was dropped, likely due to an error in the other request
                     error!(
                         "Failed to receive account from pending request: {err}"
@@ -1499,9 +1517,12 @@ mod tests {
                 rpc_client,
                 pubsub_client,
                 forward_tx,
-                &RemoteAccountProviderConfig::default_with_lifecycle_mode(
+                &RemoteAccountProviderConfig::try_new_with_metrics(
+                    1000,
                     LifecycleMode::Ephemeral,
-                ),
+                    false,
+                )
+                .unwrap(),
             )
             .await
             .unwrap(),

magicblock-chainlink/src/chainlink/mod.rs

@@ -1,4 +1,7 @@
-use std::sync::Arc;
+use std::sync::{
+    atomic::{AtomicU64, Ordering},
+    Arc,
+};
 
 use dlp::pda::ephemeral_balance_pda_from_payer;
 use errors::ChainlinkResult;
@@ -136,15 +139,55 @@ impl<T: ChainRpcClient, U: ChainPubsubClient, V: AccountsBank, C: Cloner>
     pub fn reset_accounts_bank(&self) {
         let blacklisted_accounts =
             blacklisted_accounts(&self.validator_id, &self.faucet_id);
+
+        let delegated = AtomicU64::new(0);
+        let dlp_owned_not_delegated = AtomicU64::new(0);
+        let blacklisted = AtomicU64::new(0);
+        let remaining = AtomicU64::new(0);
+        let remaining_empty = AtomicU64::new(0);
+
         let removed = self.accounts_bank.remove_where(|pubkey, account| {
-            (!account.delegated()
-                // This fixes the edge-case of accounts that were in the process of
-                // being undelegated but never completed while the validator was running
-                || account.owner().eq(&dlp::id()))
-                && !blacklisted_accounts.contains(pubkey)
+            if blacklisted_accounts.contains(pubkey) {
+                blacklisted.fetch_add(1, Ordering::Relaxed);
+                return false;
+            }
+            if account.delegated() {
+                delegated.fetch_add(1, Ordering::Relaxed);
+                return false;
+            }
+            if account.owner().eq(&dlp::id()) {
+                dlp_owned_not_delegated.fetch_add(1, Ordering::Relaxed);
+                return true;
+            }
+            trace!(
+                "Removing non-delegated, non-DLP-owned account: {pubkey} {:#?}",
+                account
+            );
+            remaining.fetch_add(1, Ordering::Relaxed);
+            if account.lamports() == 0
+                && account.owner().ne(&solana_sdk::feature::id())
+            {
+                remaining_empty.fetch_add(1, Ordering::Relaxed);
+            }
+            true
         });
 
-        debug!("Removed {removed} non-delegated accounts");
+        let non_empty = remaining
+            .load(Ordering::Relaxed)
+            .saturating_sub(remaining_empty.load(Ordering::Relaxed));
+
+        info!(
+            "Removed {removed} accounts from bank:
+{} DLP-owned non-delegated
+{} non-delegated non-blacklisted, no-feature non-empty.
+{} non-delegated non-blacklisted empty
+Kept: {} delegated, {} blacklisted",
+            dlp_owned_not_delegated.into_inner(),
+            non_empty,
+            remaining_empty.into_inner(),
+            delegated.into_inner(),
+            blacklisted.into_inner()
+        );
     }
 
     fn subscribe_account_removals(
@@ -283,7 +326,15 @@ impl<T: ChainRpcClient, U: ChainPubsubClient, V: AccountsBank, C: Cloner>
                 .map(|p| p.to_string())
                 .collect::<Vec<_>>()
                 .join(", ");
-            trace!("Fetching accounts: {pubkeys_str}");
+            let mark_empty_str = mark_empty_if_not_found
+                .map(|keys| {
+                    keys.iter()
+                        .map(|p| p.to_string())
+                        .collect::<Vec<_>>()
+                        .join(", ")
+                })
+                .unwrap_or_default();
+            trace!("Fetching accounts: {pubkeys_str}, mark_empty_if_not_found: {mark_empty_str}");
         }
         Self::promote_accounts(
             fetch_cloner,