Преглед изворни кода

Drozdziak1/p2w client use mapping crawl (#294)

* p2w-client: Add attestation config merging logic

This change adds code to append new symbols to an existing attestation
config. In conjunction with the mapping crawling logic, this will aid
in hot reloading attestation configs in runtime

* pyth2wormhole-client: Use the crawled mapping symbols
Stanisław Drozd пре 3 година
родитељ
комит
fdafbe32a4

+ 114 - 4
solana/pyth2wormhole/client/src/attestation_cfg.rs

@@ -1,5 +1,9 @@
 use std::{
-    collections::HashMap,
+    collections::{
+        HashMap,
+        HashSet,
+    },
+    iter,
     str::FromStr,
 };
 
@@ -13,7 +17,7 @@ use serde::{
 use solana_program::pubkey::Pubkey;
 
 /// Pyth2wormhole config specific to attestation requests
-#[derive(Debug, Deserialize, Serialize, PartialEq)]
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
 pub struct AttestationConfig {
     #[serde(default = "default_min_msg_reuse_interval_ms")]
     pub min_msg_reuse_interval_ms: u64,
@@ -28,7 +32,61 @@ pub struct AttestationConfig {
     pub symbol_groups: Vec<SymbolGroup>,
 }
 
-#[derive(Debug, Deserialize, Serialize, PartialEq)]
+impl AttestationConfig {
+    /// Merges new symbols into the attestation config. Pre-existing
+    /// new symbols are ignored. The new_group_name group can already
+    /// exist - symbols will be appended to `symbols` field.
+    pub fn add_symbols(
+        &mut self,
+        mut new_symbols: HashMap<Pubkey, HashSet<Pubkey>>,
+        group_name: String, // Which group is extended by the new symbols
+    ) {
+        // Remove pre-existing symbols from the new symbols collection
+        for existing_group in &self.symbol_groups {
+            for existing_sym in &existing_group.symbols {
+                // Check if new symbols mention this product
+                if let Some(mut prices) = new_symbols.get_mut(&existing_sym.product_addr) {
+                    // Prune the price if exists
+                    prices.remove(&existing_sym.price_addr);
+                }
+            }
+        }
+
+        // Turn the pruned symbols into P2WSymbol structs
+        let mut new_symbols_vec = new_symbols
+            .drain() // Makes us own the elements and lets us move them
+            .map(|(prod, prices)| iter::zip(iter::repeat(prod), prices)) // Convert to iterator over flat (prod, price) tuples
+            .flatten() // Flatten the tuple iterators
+            .map(|(prod, price)| P2WSymbol {
+                name: None,
+                product_addr: prod,
+                price_addr: price,
+            })
+            .collect::<Vec<P2WSymbol>>();
+
+        // Find and extend OR create the group of specified name
+        match self
+            .symbol_groups
+            .iter_mut()
+            .find(|g| g.group_name == group_name) // Advances the iterator and returns Some(item) on first hit
+        {
+            Some(mut existing_group) => existing_group.symbols.append(&mut new_symbols_vec),
+            None if new_symbols_vec.len() != 0 => {
+                // Group does not exist, assume defaults
+                let new_group = SymbolGroup {
+                    group_name,
+                    conditions: Default::default(),
+                    symbols: new_symbols_vec,
+                };
+
+                self.symbol_groups.push(new_group);
+            }
+            None => {}
+        }
+    }
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
 pub struct SymbolGroup {
     pub group_name: String,
     /// Attestation conditions applied to all symbols in this group
@@ -56,7 +114,7 @@ pub const fn default_max_batch_jobs() -> usize {
 /// of the active conditions is met. Option<> fields can be
 /// de-activated with None. All conditions are inactive by default,
 /// except for the non-Option ones.
-#[derive(Clone, Default, Debug, Deserialize, Serialize, PartialEq)]
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
 pub struct AttestationConditions {
     /// Baseline, unconditional attestation interval. Attestation is triggered if the specified interval elapsed since last attestation.
     #[serde(default = "default_min_interval_secs")]
@@ -78,6 +136,17 @@ pub struct AttestationConditions {
     pub publish_time_min_delta_secs: Option<u64>,
 }
 
+impl Default for AttestationConditions {
+    fn default() -> Self {
+        Self {
+            min_interval_secs: default_min_interval_secs(),
+            max_batch_jobs: default_max_batch_jobs(),
+            price_changed_pct: None,
+            publish_time_min_delta_secs: None,
+        }
+    }
+}
+
 /// Config entry for a Pyth product + price pair
 #[derive(Clone, Default, Debug, Deserialize, Serialize, PartialEq, Eq)]
 pub struct P2WSymbol {
@@ -200,4 +269,45 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_add_symbols_works() -> Result<(), ErrBox> {
+        let empty_config = AttestationConfig {
+            min_msg_reuse_interval_ms: 1000,
+            max_msg_accounts: 100,
+            mapping_addr: None,
+            symbol_groups: vec![],
+        };
+
+        let mock_new_symbols = (0..255)
+            .map(|sym_idx| {
+                let mut mock_prod_bytes = [0u8; 32];
+                mock_prod_bytes[31] = sym_idx;
+
+                let mut mock_prices = HashSet::new();
+                for px_idx in 1..=5 {
+                    let mut mock_price_bytes = [0u8; 32];
+                    mock_price_bytes[31] = sym_idx;
+                    mock_prices.insert(Pubkey::new_from_array(mock_price_bytes));
+                }
+
+                (Pubkey::new_from_array(mock_prod_bytes), mock_prices)
+            })
+            .collect::<HashMap<Pubkey, HashSet<Pubkey>>>();
+
+        let mut config1 = empty_config.clone();
+
+        config1.add_symbols(mock_new_symbols.clone(), "default".to_owned());
+
+        let mut config2 = config1.clone();
+
+        // Should not be created because there's no new symbols to add
+        // (we're adding identical mock_new_symbols again)
+        config2.add_symbols(mock_new_symbols.clone(), "default2".to_owned());
+
+        assert_ne!(config1, empty_config); // Check that config grows from empty
+        assert_eq!(config1, config2); // Check that no changes are made if all symbols are already in there
+
+        Ok(())
+    }
 }

+ 2 - 1
solana/pyth2wormhole/client/src/main.rs

@@ -173,12 +173,13 @@ async fn main() -> Result<(), ErrBox> {
             daemon,
         } => {
             // Load the attestation config yaml
-            let attestation_cfg: AttestationConfig =
+            let mut attestation_cfg: AttestationConfig =
                 serde_yaml::from_reader(File::open(attestation_cfg)?)?;
 
             if let Some(mapping_addr) = attestation_cfg.mapping_addr.as_ref() {
                 let additional_accounts = crawl_pyth_mapping(&rpc_client, mapping_addr).await?;
                 info!("Additional mapping accounts:\n{:#?}", additional_accounts);
+                attestation_cfg.add_symbols(additional_accounts, "mapping".to_owned());
             }
 
             handle_attest(

+ 1 - 1
solana/pyth2wormhole/program/src/message.rs

@@ -29,7 +29,7 @@ pub struct P2WMessageDrvData {
     /// The key owning this message account
     pub message_owner: Pubkey,
     /// Size of the batch. It is important that all messages have the same size
-    /// 
+    ///
     /// NOTE: 2022-09-05
     /// Currently wormhole does not resize accounts if they have different
     /// payload sizes; this (along with versioning the seed literal below) is

+ 11 - 3
third_party/pyth/p2w_autoattest.py

@@ -198,9 +198,9 @@ symbol_groups:
     # integer-divide the symbols in ~half for two test
     # groups. Assumes arr[:idx] is exclusive, and arr[idx:] is
     # inclusive
-    half_len = len(pyth_accounts) // 2;
+    third_len = len(pyth_accounts) // 3;
 
-    for thing in pyth_accounts[:half_len]:
+    for thing in pyth_accounts[:third_len]:
         name = thing["name"]
         price = thing["price"]
         product = thing["product"]
@@ -218,7 +218,7 @@ symbol_groups:
     symbols:
 """
 
-    for stuff in pyth_accounts[half_len:]:
+    for stuff in pyth_accounts[third_len:-third_len]:
         name = stuff["name"]
         price = stuff["price"]
         product = stuff["product"]
@@ -228,6 +228,14 @@ symbol_groups:
         price_addr: {price}
         product_addr: {product}"""
 
+    cfg_yaml += f"""
+  - group_name: mapping
+    conditions:
+      min_interval_secs: 30
+      price_changed_pct: 5 
+    symbols: []
+"""
+
     with open(P2W_ATTESTATION_CFG, "w") as f:
         f.write(cfg_yaml)
         f.flush()