refactor: funcs

Author: alexcambose

example.reader.config.toml

@@ -3,12 +3,14 @@ use async_nats::jetstream::{self};
 use async_nats::HeaderMap;
 use clap::Parser;
 use config::Config;
+use futures::future::join_all;
 use pyth_sdk_solana::state::{load_price_account, PriceStatus, PythnetPriceAccount};
 use pyth_stream::utils::setup_jetstream;
 use serde::{Deserialize, Deserializer, Serialize};
 use solana_account_decoder::UiAccountEncoding;
 use solana_client::nonblocking::pubsub_client::PubsubClient;
 use solana_client::rpc_config::{RpcAccountInfoConfig, RpcProgramAccountsConfig};
+use solana_client::rpc_response::{Response, RpcKeyedAccount};
 use solana_metrics::datapoint_info;
 use solana_sdk::account::Account;
 use solana_sdk::commitment_config::CommitmentConfig;
@@ -17,6 +19,7 @@ use std::collections::HashMap;
 use std::collections::HashSet;
 use std::path::PathBuf;
 use std::str::FromStr;
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::Instant;
 use tokio::sync::Mutex;
@@ -26,6 +29,7 @@ use tokio_stream::StreamExt;
 use tracing::{debug, error, info, warn};
 use tracing_subscriber::{fmt, EnvFilter};
 use url::Url;
+
 #[derive(Debug, Serialize, Deserialize)]
 struct PriceUpdate {
     #[serde(rename = "type")]
@@ -101,6 +105,111 @@ struct Args {
     config: Option<PathBuf>,
 }
 
+fn get_price_account_from_update(
+    update: &Response<RpcKeyedAccount>,
+) -> Result<PythnetPriceAccount> {
+    let account: Account = match update.value.account.decode() {
+        Some(account) => account,
+        _none => {
+            warn!("Failed to decode account from update");
+            return Err(anyhow::anyhow!("Failed to decode account from update"));
+        }
+    };
+
+    let price_account: PythnetPriceAccount = match load_price_account(&account.data) {
+        Ok(pyth_account) => *pyth_account,
+        Err(_) => {
+            debug!("Not a price account, skipping");
+            return Err(anyhow::anyhow!("Not a price account, skipping"));
+        }
+    };
+
+    Ok(price_account)
+}
+
+async fn publish_price_updates(
+    jetstream: jetstream::Context,
+    price_account: PythnetPriceAccount,
+    update: &Response<RpcKeyedAccount>,
+) {
+    let price_update = PriceUpdate {
+        update_type: "price_update".to_string(),
+        price_feed: PriceFeed {
+            id: update.value.pubkey.to_string(),
+            price: PriceInfo {
+                price: price_account.agg.price.to_string(),
+                conf: price_account.agg.conf.to_string(),
+                expo: price_account.expo,
+                publish_time: price_account.timestamp,
+                slot: update.context.slot,
+            },
+            ema_price: PriceInfo {
+                price: price_account.ema_price.val.to_string(),
+                conf: price_account.ema_conf.val.to_string(),
+                expo: price_account.expo,
+                publish_time: price_account.timestamp,
+                slot: update.context.slot,
+            },
+        },
+    };
+    let price_update_message = serde_json::to_string(&price_update).unwrap();
+
+    // Create a unique message ID
+    let price_update_message_id = format!(
+        "{}:{}",
+        price_update.price_feed.id, price_update.price_feed.price.slot
+    );
+
+    // Create headers with the Nats-Msg-Id
+    let mut price_update_headers = HeaderMap::new();
+    price_update_headers.insert("Nats-Msg-Id", price_update_message_id.as_str());
+
+    match jetstream
+        .publish_with_headers(
+            "pyth.price.updates",
+            price_update_headers,
+            price_update_message.into(),
+        )
+        .await
+    {
+        Ok(_) => debug!(
+            "Published price update to JetStream with ID: {}",
+            price_update_message_id
+        ),
+        Err(e) => warn!("Failed to publish price update to JetStream: {}", e),
+    }
+}
+
+/**
+ * Process the publisher price updates for a given price account and update
+ * @param price_account: The price account
+ * @param update: The update
+ * @param publisher_buffer: The publisher buffer
+ */
+async fn process_publisher_price_updates(
+    price_account: PythnetPriceAccount,
+    update: &Response<RpcKeyedAccount>,
+    publisher_buffer: PublisherBuffer,
+) {
+    for component in price_account.comp {
+        if component.publisher != Pubkey::default() {
+            let publisher_price_update = PublisherPriceUpdate {
+                feed_id: update.value.pubkey.to_string(),
+                publisher: component.publisher.to_string(),
+                price: price_account.agg.price.to_string(),
+                slot: update.context.slot,
+            };
+
+            let key = (
+                publisher_price_update.feed_id.clone(),
+                publisher_price_update.publisher.clone(),
+            );
+            let mut buf = publisher_buffer.lock().await;
+            buf.insert(key, publisher_price_update);
+        }
+    }
+}
+
 async fn fetch_price_updates(jetstream: jetstream::Context, config: &AppConfig) -> Result<()> {
     info!("Starting Pyth reader");
     let publisher_buffer: PublisherBuffer = Arc::new(Mutex::new(HashMap::new()));
@@ -125,10 +234,9 @@ async fn fetch_price_updates(jetstream: jetstream::Context, config: &AppConfig)
         .program_subscribe(&config.pyth.program_key, Some(rpc_config))
         .await?;
 
-    let mut update_count = 0;
-    let mut duration_count = 0;
-    let mut unique_price_feeds = HashSet::new();
-    let mut last_report_time = Instant::now();
+    let update_count = Arc::new(AtomicUsize::new(0));
+    let duration_count = Arc::new(AtomicUsize::new(0));
+    let unique_price_feeds = Arc::new(Mutex::new(HashSet::<Pubkey>::new()));
     let jetstream_clone = jetstream.clone();
     let buffer_clone = publisher_buffer.clone();
     let mut msg_id_counter = 0;
@@ -170,26 +278,40 @@ async fn fetch_price_updates(jetstream: jetstream::Context, config: &AppConfig)
         }
     });
 
+    let update_count_clone = update_count.clone();
+    let duration_count_clone = duration_count.clone();
+    tokio::spawn(async move {
+        loop {
+            tokio::time::sleep(Duration::from_secs(1)).await;
+            info!(
+                "Total duration: {:?}, Update Count: {:?}",
+                duration_count_clone.load(Ordering::Relaxed) as u128,
+                update_count_clone.load(Ordering::Relaxed),
+            );
+            if update_count_clone.load(Ordering::Relaxed) > 0 {
+                info!(
+                    "Average duration: {:?}",
+                    duration_count_clone.load(Ordering::Relaxed) as u128
+                        / update_count_clone.load(Ordering::Relaxed) as u128
+                );
+            }
+            duration_count_clone.store(0, Ordering::Relaxed);
+            update_count_clone.store(0, Ordering::Relaxed);
+        }
+    });
+
     while let Some(update) = notif.next().await {
-        debug!("Received price update");
         let start_time = Instant::now();
 
-        let account: Account = match update.value.account.decode() {
-            Some(account) => account,
+        debug!("Received price update");
+        let price_account: PythnetPriceAccount = match get_price_account_from_update(&update) {
+            Ok(account) => account,
             _none => {
                 warn!("Failed to decode account from update");
                 continue;
             }
         };
 
-        let price_account: PythnetPriceAccount = match load_price_account(&account.data) {
-            Ok(pyth_account) => *pyth_account,
-            Err(_) => {
-                debug!("Not a price account, skipping");
-                continue;
-            }
-        };
-
         // We want to send price updates whenever the aggregate changes but sometimes the accounts can change without the aggregate changing
         if price_account.agg.status == PriceStatus::Trading
             && (update.context.slot - price_account.agg.pub_slot)
@@ -200,98 +322,27 @@ async fn fetch_price_updates(jetstream: jetstream::Context, config: &AppConfig)
                 price_account.prod
             );
 
-            let price_update = PriceUpdate {
-                update_type: "price_update".to_string(),
-                price_feed: PriceFeed {
-                    id: update.value.pubkey.to_string(),
-                    price: PriceInfo {
-                        price: price_account.agg.price.to_string(),
-                        conf: price_account.agg.conf.to_string(),
-                        expo: price_account.expo,
-                        publish_time: price_account.timestamp,
-                        slot: update.context.slot, // Add this field
-                    },
-                    ema_price: PriceInfo {
-                        price: price_account.ema_price.val.to_string(),
-                        conf: price_account.ema_conf.val.to_string(),
-                        expo: price_account.expo,
-                        publish_time: price_account.timestamp,
-                        slot: update.context.slot, // Add this field
-                    },
-                },
-            };
-
-            let price_update_message = serde_json::to_string(&price_update)?;
-
-            // Create a unique message ID
-            let price_update_message_id = format!(
-                "{}:{}",
-                price_update.price_feed.id, price_update.price_feed.price.slot
-            );
-
-            // Create headers with the Nats-Msg-Id
-            let mut price_update_headers = HeaderMap::new();
-            price_update_headers.insert("Nats-Msg-Id", price_update_message_id.as_str());
-
             let jetstream_clone = jetstream.clone();
+            let publisher_buffer_clone = publisher_buffer.clone();
+            let unique_price_feeds_clone = unique_price_feeds.clone();
+            let update_count_clone = update_count.clone();
+            let duration_count_clone = duration_count.clone();
             task::spawn(async move {
-                match jetstream_clone
-                    .publish_with_headers(
-                        "pyth.price.updates",
-                        price_update_headers,
-                        price_update_message.into(),
-                    )
-                    .await
-                {
-                    Ok(_) => debug!(
-                        "Published price update to JetStream with ID: {}",
-                        price_update_message_id
-                    ),
-                    Err(e) => warn!("Failed to publish price update to JetStream: {}", e),
-                }
+                publish_price_updates(jetstream_clone, price_account, &update).await;
+                process_publisher_price_updates(price_account, &update, publisher_buffer_clone)
+                    .await;
+
+                let mut feeds = unique_price_feeds_clone.lock().await;
+                feeds.insert(price_account.prod);
+                let end_time = Instant::now();
+                let duration = end_time.duration_since(start_time);
+                update_count_clone.fetch_add(1, Ordering::Relaxed);
+                duration_count_clone.fetch_add(duration.as_millis() as usize, Ordering::Relaxed);
+                info!("Processing price update, slot: {}", update.context.slot);
             });
-
-            for component in price_account.comp {
-                let publisher = component.publisher.to_string();
-                let publisher_price_update = PublisherPriceUpdate {
-                    feed_id: update.value.pubkey.to_string(),
-                    publisher: publisher.clone(),
-                    price: price_account.agg.price.to_string(),
-                    slot: update.context.slot, // Add this field
-                };
-
-                let key = (publisher_price_update.feed_id.clone(), publisher);
-                let mut buf = publisher_buffer.lock().await;
-                buf.insert(key, publisher_price_update);
-            }
-
-            let end_time = Instant::now();
-            let duration = end_time.duration_since(start_time);
-            duration_count += duration.as_micros();
-            update_count += 1;
-            unique_price_feeds.insert(price_account.prod);
         } else {
             debug!("Skipping price update due to invalid status or slot difference");
         }
-
-        // Report aggregate information every minute and emit metrics
-        if last_report_time.elapsed() >= Duration::from_secs(1) {
-            info!(
-                "Processed {} updates from {} unique price feeds in the last 1 secs",
-                update_count,
-                unique_price_feeds.len()
-            );
-            info!(
-                "last slot: {}, comp: {}",
-                update.context.slot,
-                price_account.comp.len()
-            );
-            info!("Average duration: {:?}", duration_count / update_count);
-            duration_count = 0;
-            update_count = 0;
-            unique_price_feeds.clear();
-            last_report_time = Instant::now();
-        }
     }
 
     // If we exit the loop, it means the stream has ended

example.websocket.config.toml

@@ -9,7 +9,7 @@ pub async fn setup_jetstream(nats_client: &async_nats::Client) -> Result<jetstre
     let price_updates_stream_config = stream::Config {
         name: "PYTH_PRICE_UPDATES".to_string(),
         subjects: vec!["pyth.price.updates".to_string()],
-        max_bytes: 1024 * 1024 * 4000,
+        max_bytes: 1024 * 1024 * 40000,
         duplicate_window: Duration::from_secs(60),
         discard: stream::DiscardPolicy::Old,
         allow_direct: true,
@@ -20,7 +20,8 @@ pub async fn setup_jetstream(nats_client: &async_nats::Client) -> Result<jetstre
     let publisher_updates_stream_config = stream::Config {
         name: "PYTH_PUBLISHER_UPDATES".to_string(),
         subjects: vec!["pyth.publisher.updates".to_string()],
-        max_bytes: 1024 * 1024 * 4000,
+        max_bytes: 1024 * 1024 * 40000,
+        max_messages: 1000000,
         duplicate_window: Duration::from_secs(60),
         discard: stream::DiscardPolicy::Old,
         allow_direct: true,