refactor: use ordered publisher list

This change brings down the average cost of a price update without
aggregation from 5k to 2k because binary search requires much less
lookups. There are mechanims implemented in the code to make sure
that upon the upgrade we sort the array without adding overhead to
the happy path (working binary search). Also add_publisher now does
an insertion sort to keep the list sorted and del_publisher won't
change the order of the list.

Author: ali-behjati

program/rust/src/processor/add_publisher.rs

@@ -73,13 +73,22 @@ pub fn add_publisher(
         }
     }
 
-    let current_index: usize = try_convert(price_data.num_)?;
+    let mut current_index: usize = try_convert(price_data.num_)?;
     sol_memset(
         bytes_of_mut(&mut price_data.comp_[current_index]),
         0,
         size_of::<PriceComponent>(),
     );
     price_data.comp_[current_index].pub_ = cmd_args.publisher;
+
+    // Shift the element back to keep the publishers components sorted.
+    while current_index > 0
+        && price_data.comp_[current_index].pub_ < price_data.comp_[current_index - 1].pub_
+    {
+        price_data.comp_.swap(current_index, current_index - 1);
+        current_index -= 1;
+    }
+
     price_data.num_ += 1;
     price_data.header.size = try_convert::<_, u32>(PriceAccount::INITIAL_SIZE)?;
     Ok(())

program/rust/src/processor/upd_price.rs

@@ -127,27 +127,74 @@ pub fn upd_price(
     // Check clock
     let clock = Clock::from_account_info(clock_account)?;
 
-    let mut publisher_index: usize = 0;
+    let publisher_index: usize;
     let latest_aggregate_price: PriceInfo;
 
     // The price_data borrow happens in a scope because it must be
     // dropped before we borrow again as raw data pointer for the C
     // aggregation logic.
     {
         // Verify that symbol account is initialized
-        let price_data = load_checked::<PriceAccount>(price_account, cmd_args.header.version)?;
+        let mut price_data = load_checked::<PriceAccount>(price_account, cmd_args.header.version)?;
 
-        // Verify that publisher is authorized
-        while publisher_index < try_convert::<u32, usize>(price_data.num_)? {
-            if price_data.comp_[publisher_index].pub_ == *funding_account.key {
-                break;
+        // Verify that publisher is authorized by initially binary searching
+        // for the publisher's component in the price account. The binary
+        // search might not work if the publisher list is not sorted; therefore
+        // we fall back to a linear search and during the search we try to bring
+        // the publishers list closer to a sorted state. (We do not sort them
+        // entirely to avoid exceeding the compute budget.)
+        let mut binary_search_result = None;
+
+        {
+            // Binary search to find the publisher key. Rust std binary search is not used because
+            // they guarantee valid outcome only if the array is sorted whereas we want to rely on
+            // a Equal match if it is a result on an unsorted array. Currently the rust
+            // implementation behaves the same but we do not want to rely on api internals.
+            let mut left = 0;
+            let mut right = try_convert::<u32, usize>(price_data.num_)?;
+            while left < right {
+                let mid = left + (right - left) / 2;
+                match price_data.comp_[mid].pub_.cmp(funding_account.key) {
+                    std::cmp::Ordering::Less => {
+                        left = mid + 1;
+                    }
+                    std::cmp::Ordering::Greater => {
+                        right = mid;
+                    }
+                    std::cmp::Ordering::Equal => {
+                        binary_search_result = Some(mid);
+                        break;
+                    }
+                }
             }
-            publisher_index += 1;
         }
-        pyth_assert(
-            publisher_index < try_convert::<u32, usize>(price_data.num_)?,
-            OracleError::PermissionViolation.into(),
-        )?;
+
+        publisher_index = match binary_search_result {
+            Some(index) => index,
+            None => {
+                // One pass through the publishers list to make it closer
+                // to being sorted.
+                for i in 1..try_convert::<u32, usize>(price_data.num_)? {
+                    if price_data.comp_[i].pub_ < price_data.comp_[i - 1].pub_ {
+                        price_data.comp_.swap(i, i - 1);
+                    }
+                }
+
+                let mut index = 0;
+                while index < try_convert::<u32, usize>(price_data.num_)? {
+                    if price_data.comp_[index].pub_ == *funding_account.key {
+                        break;
+                    }
+                    index += 1;
+                }
+                pyth_assert(
+                    index < try_convert::<u32, usize>(price_data.num_)?,
+                    OracleError::PermissionViolation.into(),
+                )?;
+                index
+            }
+        };
+
 
         latest_aggregate_price = price_data.agg_;
         let latest_publisher_price = price_data.comp_[publisher_index].latest_;

program/rust/src/tests/test_upd_price_v2.rs

@@ -444,6 +444,90 @@ fn test_upd_price_v2() -> Result<(), Box<dyn std::error::Error>> {
     Ok(())
 }
 
+#[test]
+fn test_upd_works_with_unordered_publisher_set() -> Result<(), Box<dyn std::error::Error>> {
+    let mut instruction_data = [0u8; size_of::<UpdPriceArgs>()];
+
+    let program_id = Pubkey::new_unique();
+
+    let mut price_setup = AccountSetup::new::<PriceAccount>(&program_id);
+    let mut price_account = price_setup.as_account_info();
+    price_account.is_signer = false;
+    PriceAccount::initialize(&price_account, PC_VERSION).unwrap();
+
+    let mut publishers_setup: Vec<_> = (0..20).map(|_| AccountSetup::new_funding()).collect();
+    let mut publishers: Vec<_> = publishers_setup
+        .iter_mut()
+        .map(|s| s.as_account_info())
+        .collect();
+
+    publishers.sort_by_key(|x| x.key);
+
+    {
+        let mut price_data = load_checked::<PriceAccount>(&price_account, PC_VERSION).unwrap();
+        price_data.num_ = 20;
+        // Store the publishers in reverse order
+        publishers
+            .iter()
+            .rev()
+            .enumerate()
+            .for_each(|(i, account)| {
+                println!("{:} {:?}", i, account.key);
+                price_data.comp_[i].pub_ = *account.key;
+            });
+    }
+
+    let mut clock_setup = AccountSetup::new_clock();
+    let mut clock_account = clock_setup.as_account_info();
+    clock_account.is_signer = false;
+    clock_account.is_writable = false;
+
+    update_clock_slot(&mut clock_account, 1);
+
+    // repeat 10 times to also make sure that price updates
+    // make the publisher list sorted
+    for slot in 1..10 {
+        for (i, publisher) in publishers.iter().enumerate() {
+            println!("{:} {:?}", i, publisher.key);
+            populate_instruction(&mut instruction_data, (i + 100) as i64, 10, slot);
+            process_instruction(
+                &program_id,
+                &[
+                    publisher.clone(),
+                    price_account.clone(),
+                    clock_account.clone(),
+                ],
+                &instruction_data,
+            )?;
+        }
+
+        update_clock_slot(&mut clock_account, slot + 1);
+    }
+
+    {
+        let price_data = load_checked::<PriceAccount>(&price_account, PC_VERSION).unwrap();
+
+        // Check publishers are sorted
+        let price_data_pubs: Vec<_> = price_data.comp_[..20].iter().map(|c| c.pub_).collect();
+        assert_eq!(
+            price_data_pubs,
+            publishers.iter().map(|p| *p.key).collect::<Vec<_>>()
+        );
+
+        assert_eq!(price_data.comp_[0].latest_.price_, 100);
+        assert_eq!(price_data.comp_[0].latest_.conf_, 10);
+        assert_eq!(price_data.comp_[0].latest_.pub_slot_, 9);
+        assert_eq!(price_data.comp_[0].latest_.status_, PC_STATUS_TRADING);
+        // aggregate is calculated at slot 9 for up to slot 8
+        assert_eq!(price_data.valid_slot_, 8);
+        assert_eq!(price_data.agg_.pub_slot_, 9);
+        assert_eq!(price_data.agg_.price_, 109);
+        assert_eq!(price_data.agg_.conf_, 8);
+        assert_eq!(price_data.agg_.status_, PC_STATUS_TRADING);
+    }
+    Ok(())
+}
+
 // Create an upd_price instruction with the provided parameters
 fn populate_instruction(instruction_data: &mut [u8], price: i64, conf: u64, pub_slot: u64) {
     let mut cmd = load_mut::<UpdPriceArgs>(instruction_data).unwrap();