refactor(agent): extract oracle component/service

Author: Reisen

src/agent.rs

@@ -128,19 +128,45 @@ impl Agent {
         // Create the Application State.
         let state = Arc::new(state::State::new(self.config.state.clone()).await);
 
-        // Spawn the primary network
-        jhs.extend(network::spawn_network(
-            self.config.primary_network.clone(),
-            network::Network::Primary,
-            state.clone(),
-        )?);
+        // Spawn the primary network Oracle.
+        {
+            // Publisher permissions updates between oracle and exporter
+            let (publisher_permissions_tx, publisher_permissions_rx) =
+                watch::channel(<_>::default());
+
+            jhs.push(tokio::spawn(services::oracle(
+                self.config.primary_network.clone(),
+                network::Network::Primary,
+                state.clone(),
+                publisher_permissions_tx.clone(),
+            )));
+
+            // Spawn the primary network
+            jhs.extend(network::spawn_network(
+                self.config.primary_network.clone(),
+                network::Network::Primary,
+                state.clone(),
+                publisher_permissions_rx.clone(),
+            )?);
+        }
 
-        // Spawn the secondary network, if needed
+        // Spawn the secondary network Oracle, if needed.
         if let Some(config) = &self.config.secondary_network {
+            let (publisher_permissions_tx, publisher_permissions_rx) =
+                watch::channel(<_>::default());
+
+            jhs.push(tokio::spawn(services::oracle(
+                config.clone(),
+                network::Network::Secondary,
+                state.clone(),
+                publisher_permissions_tx.clone(),
+            )));
+
             jhs.extend(network::spawn_network(
                 config.clone(),
                 network::Network::Secondary,
                 state.clone(),
+                publisher_permissions_rx,
             )?);
         }
 

src/agent/metrics.rs

@@ -1,6 +1,6 @@
 use {
     super::state::local::PriceInfo,
-    crate::agent::solana::oracle::PriceEntry,
+    crate::agent::state::oracle::PriceEntry,
     lazy_static::lazy_static,
     prometheus_client::{
         encoding::{

src/agent/services.rs

@@ -1,7 +1,9 @@
 pub mod keypairs;
 pub mod notifier;
+pub mod oracle;
 
 pub use {
     keypairs::keypairs,
     notifier::notifier,
+    oracle::oracle,
 };

src/agent/services/oracle.rs

@@ -0,0 +1,206 @@
+//! Oracle
+//!
+//! The Oracle service is respoinsible for reacting to all remote/on-chain events.
+
+use {
+    crate::agent::{
+        solana::{
+            key_store::KeyStore,
+            network::{
+                Config,
+                Network,
+            },
+        },
+        state::oracle::{
+            Oracle,
+            PricePublishingMetadata,
+        },
+    },
+    anyhow::Result,
+    solana_account_decoder::UiAccountEncoding,
+    solana_client::{
+        nonblocking::{
+            pubsub_client::PubsubClient,
+            rpc_client::RpcClient,
+        },
+        rpc_config::{
+            RpcAccountInfoConfig,
+            RpcProgramAccountsConfig,
+        },
+        rpc_response::{
+            Response,
+            RpcKeyedAccount,
+        },
+    },
+    solana_sdk::{
+        account::Account,
+        commitment_config::CommitmentConfig,
+        pubkey::Pubkey,
+    },
+    std::{
+        collections::HashMap,
+        sync::Arc,
+    },
+    tokio::{
+        select,
+        sync::watch::Sender,
+    },
+    tokio_stream::StreamExt,
+};
+
+pub async fn oracle<S>(
+    config: Config,
+    network: Network,
+    state: Arc<S>,
+    publisher_permissions_tx: Sender<HashMap<Pubkey, HashMap<Pubkey, PricePublishingMetadata>>>,
+) where
+    S: Oracle,
+    S: Send + Sync + 'static,
+{
+    loop {
+        let _ = run(
+            config.clone(),
+            network,
+            state.clone(),
+            publisher_permissions_tx.clone(),
+        )
+        .await;
+    }
+}
+
+async fn run<S>(
+    config: Config,
+    network: Network,
+    state: Arc<S>,
+    publisher_permissions_tx: Sender<HashMap<Pubkey, HashMap<Pubkey, PricePublishingMetadata>>>,
+) -> Result<()>
+where
+    S: Oracle,
+    S: Send + Sync + 'static,
+{
+    let key_store = KeyStore::new(config.key_store.clone())?;
+    let mut poll_interval = tokio::time::interval(config.oracle.poll_interval_duration);
+
+    // Setup PubsubClient to listen for account changes on the Oracle program.
+    let client = PubsubClient::new(config.wss_url.as_str()).await?;
+    let (mut notifier, _unsub) = {
+        let program_key = key_store.program_key;
+        let commitment = config.oracle.commitment;
+        let config = RpcProgramAccountsConfig {
+            account_config: RpcAccountInfoConfig {
+                commitment: Some(CommitmentConfig { commitment }),
+                encoding: Some(UiAccountEncoding::Base64Zstd),
+                ..Default::default()
+            },
+            filters:        None,
+            with_context:   Some(true),
+        };
+        client.program_subscribe(&program_key, Some(config)).await
+    }?;
+
+    // Setup an RpcClient for manual polling.
+    let client = {
+        Arc::new(RpcClient::new_with_timeout_and_commitment(
+            config.rpc_url,
+            config.rpc_timeout,
+            CommitmentConfig {
+                commitment: config.oracle.commitment,
+            },
+        ))
+    };
+
+    loop {
+        select! {
+            update = notifier.next()      => handle_account_update(state.clone(), network, update).await,
+            _      = poll_interval.tick() => handle_poll_updates(
+                state.clone(),
+                network,
+                key_store.mapping_key,
+                client.clone(),
+                config.oracle.max_lookup_batch_size,
+                config.oracle.subscriber_enabled,
+                publisher_permissions_tx.clone(),
+            ).await,
+        }
+    }
+}
+
+/// When an account RPC Subscription update is receiveed.
+///
+/// We check if the account is one we're aware of and tracking, and if so, spawn
+/// a small background task that handles that update. We only do this for price
+/// accounts, all other accounts are handled below in the poller.
+async fn handle_account_update<S>(
+    state: Arc<S>,
+    network: Network,
+    update: Option<Response<RpcKeyedAccount>>,
+) where
+    S: Oracle,
+    S: Send + Sync + 'static,
+{
+    match update {
+        Some(update) => {
+            match update.value.account.decode::<Account>() {
+                Some(account) => {
+                    tokio::spawn(async move {
+                        if let Err(err) = async move {
+                            let key: Pubkey = update.value.pubkey.as_str().try_into()?;
+                            Oracle::handle_price_account_update(&*state, network, &key, &account)
+                                .await
+                        }
+                        .await
+                        {
+                            tracing::error!(err = ?err, "Failed to handle account update.");
+                        }
+                    });
+                }
+
+                None => {
+                    tracing::error!(
+                        update = ?update,
+                        "Failed to decode account from update.",
+                    );
+                }
+            };
+        }
+        None => {
+            tracing::debug!("subscriber closed connection");
+        }
+    }
+}
+
+/// On poll lookup all Pyth Mapping/Product/Price accounts and sync.
+async fn handle_poll_updates<S>(
+    state: Arc<S>,
+    network: Network,
+    mapping_key: Pubkey,
+    rpc: Arc<RpcClient>,
+    max_lookup_batch_size: usize,
+    enabled: bool,
+    publisher_permissions_tx: Sender<HashMap<Pubkey, HashMap<Pubkey, PricePublishingMetadata>>>,
+) where
+    S: Oracle,
+    S: Send + Sync + 'static,
+{
+    if !enabled {
+        return;
+    }
+
+    tokio::spawn(async move {
+        if let Err(err) = async {
+            Oracle::poll_updates(
+                &*state,
+                mapping_key,
+                &rpc,
+                max_lookup_batch_size,
+                publisher_permissions_tx,
+            )
+            .await?;
+            Oracle::sync_global_store(&*state, network).await
+        }
+        .await
+        {
+            tracing::error!(err = ?err, "Failed to handle poll updates.");
+        }
+    });
+}

src/agent/solana.rs

@@ -1,5 +1,4 @@
 pub mod exporter;
-pub mod oracle;
 
 /// This module encapsulates all the interaction with a single Solana network:
 /// - The Oracle, which reads data from the network
@@ -12,15 +11,22 @@ pub mod network {
                 self,
                 KeyStore,
             },
-            oracle,
         },
-        crate::agent::state::State,
+        crate::agent::state::{
+            oracle::{
+                self,
+                PricePublishingMetadata,
+            },
+            State,
+        },
         anyhow::Result,
         serde::{
             Deserialize,
             Serialize,
         },
+        solana_sdk::pubkey::Pubkey,
         std::{
+            collections::HashMap,
             sync::Arc,
             time::Duration,
         },
@@ -70,25 +76,51 @@ pub mod network {
         pub exporter:    exporter::Config,
     }
 
+    /// Spawn an Oracle, in-progress porting this to State.
+    ///
+    /// Behaviour:
+    /// - Spawns Oracle: (Obsolete, now Extracted to state/oracle.rs)
+    ///   - Spawns a Subscriber:
+    ///     o Subscribes to the Oracle program key.
+    ///     o Decodes account events related to the Oracle.
+    ///     o Sends update.
+    ///   - Spawns a Poller:
+    ///     o Fetches Mapping Accounts
+    ///     o Iterates Product+Price Accounts
+    ///     o Sends update.
+    ///   - Oracle then Listens for Updates from Subscriber
+    ///     o Filters for Price Account Updates.
+    ///     o Stores its own copy of the Price Account.
+    ///     o Updates the Global Store for that Price Account.
+    ///   - Oracle also Listens for Updates from Poller
+    ///     o Tracks if any new Mapping Accounts were found.
+    ///     o Update Local Data
+    ///     o Updates entire Global Store View.
+    /// - Spawns Exporter:
+    ///   - Spawns NetworkQuerier
+    ///     - Queries BlockHash in a timer.
+    ///     - Sends BlockHash + Slot
+    /// - Spawns Transaction Monitor:
+    ///   - Listens for for Transactions
+    ///   - Adds to tracked Transactions
+    ///   - Responds to queries about Tx status.
+    /// - Spawns Exporter
+    ///   - On Publish tick: pushes updates to the network as a batch.
+    ///   - On Compute Unit Price Tick: calculates new median price fee from recent
+    ///
+    /// Plan:
+    ///  - Subscriber & Poller Can Be Spawnable Tasks
+    ///  - Oracle becomes a State API
+    ///  -
     pub fn spawn_network(
         config: Config,
         network: Network,
         state: Arc<State>,
+        publisher_permissions_rx: watch::Receiver<
+            HashMap<Pubkey, HashMap<Pubkey, PricePublishingMetadata>>,
+        >,
     ) -> Result<Vec<JoinHandle<()>>> {
-        // Publisher permissions updates between oracle and exporter
-        let (publisher_permissions_tx, publisher_permissions_rx) = watch::channel(<_>::default());
-
-        // Spawn the Oracle
-        let mut jhs = oracle::spawn_oracle(
-            config.oracle.clone(),
-            network,
-            &config.rpc_url,
-            &config.wss_url,
-            config.rpc_timeout,
-            publisher_permissions_tx,
-            KeyStore::new(config.key_store.clone())?,
-            state.clone(),
-        );
+        let mut jhs = vec![];
 
         // Spawn the Exporter
         let exporter_jhs = exporter::spawn_exporter(
@@ -108,7 +140,7 @@ pub mod network {
 }
 
 /// The key_store module is responsible for parsing the pythd key store.
-mod key_store {
+pub mod key_store {
     use {
         anyhow::Result,
         serde::{

src/agent/solana/exporter.rs

@@ -3,7 +3,6 @@ use {
     super::{
         key_store,
         network::Network,
-        oracle::PricePublishingMetadata,
     },
     crate::agent::state::{
         global::GlobalStore,
@@ -12,6 +11,7 @@ use {
             LocalStore,
             PriceInfo,
         },
+        oracle::PricePublishingMetadata,
         State,
     },
     anyhow::{