From c27465dc12afa1388937fffde3de8dce0e7960c2 Mon Sep 17 00:00:00 2001 From: Ali Behjati Date: Sun, 5 May 2024 00:12:44 +0200 Subject: [PATCH 1/3] 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. --- program/rust/src/accounts/price.rs | 30 ++++ program/rust/src/processor/add_publisher.rs | 153 ++++++++++++++++++- program/rust/src/processor/upd_price.rs | 110 +++++++++++-- program/rust/src/tests/test_add_publisher.rs | 43 +++++- program/rust/src/tests/test_upd_price_v2.rs | 81 ++++++++++ 5 files changed, 401 insertions(+), 16 deletions(-) diff --git a/program/rust/src/accounts/price.rs b/program/rust/src/accounts/price.rs index 421cc827b..82f2d62cd 100644 --- a/program/rust/src/accounts/price.rs +++ b/program/rust/src/accounts/price.rs @@ -1,4 +1,6 @@ pub use price_pythnet::*; +#[cfg(test)] +use quickcheck::Arbitrary; use { super::{ AccountHeader, @@ -188,6 +190,7 @@ mod price_pythnet { } #[repr(C)] +#[cfg_attr(test, derive(Debug, PartialEq))] #[derive(Copy, Clone, Pod, Zeroable)] pub struct PriceComponent { pub pub_: Pubkey, @@ -195,7 +198,21 @@ pub struct PriceComponent { pub latest_: PriceInfo, } +#[cfg(test)] +impl Arbitrary for PriceComponent { + fn arbitrary(g: &mut quickcheck::Gen) -> Self { + let mut key = [0u8; 32]; + key.iter_mut().for_each(|item| *item = u8::arbitrary(g)); + PriceComponent { + pub_: Pubkey::new_from_array(key), + agg_: PriceInfo::arbitrary(g), + latest_: PriceInfo::arbitrary(g), + } + } +} + #[repr(C)] +#[cfg_attr(test, derive(Debug, PartialEq))] #[derive(Copy, Clone, Pod, Zeroable)] pub struct PriceInfo { pub price_: i64, @@ -205,6 +222,19 @@ pub struct PriceInfo { pub pub_slot_: u64, } +#[cfg(test)] +impl Arbitrary for PriceInfo { + fn arbitrary(g: &mut quickcheck::Gen) -> Self { + PriceInfo { + price_: i64::arbitrary(g), + conf_: u64::arbitrary(g), + status_: u32::arbitrary(g), + corp_act_status_: u32::arbitrary(g), + pub_slot_: u64::arbitrary(g), + } + } +} + #[repr(C)] #[derive(Copy, Clone, Pod, Zeroable)] pub struct PriceEma { diff --git a/program/rust/src/processor/add_publisher.rs b/program/rust/src/processor/add_publisher.rs index 2e6a081e9..015884ff2 100644 --- a/program/rust/src/processor/add_publisher.rs +++ b/program/rust/src/processor/add_publisher.rs @@ -24,7 +24,10 @@ use { account_info::AccountInfo, entrypoint::ProgramResult, program_error::ProgramError, - program_memory::sol_memset, + program_memory::{ + sol_memcmp, + sol_memset, + }, pubkey::Pubkey, }, std::mem::size_of, @@ -41,8 +44,7 @@ pub fn add_publisher( let cmd_args = load::(instruction_data)?; pyth_assert( - instruction_data.len() == size_of::() - && cmd_args.publisher != Pubkey::default(), + instruction_data.len() == size_of::(), ProgramError::InvalidArgument, )?; @@ -63,6 +65,14 @@ pub fn add_publisher( let mut price_data = load_checked::(price_account, cmd_args.header.version)?; + // Use the call with the default pubkey (000..) as a trigger to sort the publishers as a + // migration step from unsorted list to sorted list. + if cmd_args.publisher == Pubkey::default() { + let num_comps = try_convert::(price_data.num_)?; + sort_price_comps(&mut price_data.comp_, num_comps)?; + return Ok(()); + } + if price_data.num_ >= PC_NUM_COMP { return Err(ProgramError::InvalidArgument); } @@ -73,14 +83,149 @@ 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::(), ); 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(()) } + +/// A copy of rust slice/sort.rs heapsort implementation which is small and fast. We couldn't use +/// the sort directly because it was only accessible behind a unstable feature flag at the time of +/// writing this code. +#[inline(always)] +fn heapsort(v: &mut [(Pubkey, usize)]) { + // This binary heap respects the invariant `parent >= child`. + let sift_down = |v: &mut [(Pubkey, usize)], mut node: usize| { + loop { + // Children of `node`. + let mut child = 2 * node + 1; + if child >= v.len() { + break; + } + + // Choose the greater child. + if child + 1 < v.len() + && sol_memcmp(v[child].0.as_ref(), v[child + 1].0.as_ref(), 32) < 0 + { + child += 1; + } + + // Stop if the invariant holds at `node`. + if sol_memcmp(v[node].0.as_ref(), v[child].0.as_ref(), 32) >= 0 { + break; + } + + // Swap `node` with the greater child, move one step down, and continue sifting. + v.swap(node, child); + node = child; + } + }; + + // Build the heap in linear time. + for i in (0..v.len() / 2).rev() { + sift_down(v, i); + } + + // Pop maximal elements from the heap. + for i in (1..v.len()).rev() { + v.swap(0, i); + sift_down(&mut v[..i], 0); + } +} + +/// Sort the publishers price component list in place by performing minimal swaps. +/// This code is inspired by the sort_by_cached_key implementation in the Rust stdlib. +/// The rust stdlib implementation is not used because it uses a fast sort variant that has +/// a large code size. +/// +/// num_publishers is the number of publishers in the list that should be sorted. It is explicitly +/// passed to avoid callers mistake of passing the full slice which may contain uninitialized values. +#[inline(always)] +fn sort_price_comps(comps: &mut [PriceComponent], num_comps: usize) -> Result<(), ProgramError> { + let comps = comps + .get_mut(..num_comps) + .ok_or(ProgramError::InvalidArgument)?; + + let mut keys = comps + .iter() + .enumerate() + .map(|(i, x)| (x.pub_, i)) + .collect::>(); + + heapsort(&mut keys); + + for i in 0..num_comps { + // We know that the publisher with key[i].0 should be at index i in the sorted array and + // want to swap them in-place in O(n). Normally, the publisher at key[i].0 should be at + // key[i].1 but if it is swapped, we need to find the correct index by following the chain + // of swaps. + let mut index = keys[i].1; + + while index < i { + index = keys[index].1; + } + // Setting the final index here is important to make the code linear as we won't + // loop over from i to index again when we reach i again. + keys[i].1 = index; + comps.swap(i, index); + } + + Ok(()) +} + +#[cfg(test)] +mod test { + use { + super::*, + quickcheck_macros::quickcheck, + }; + + #[quickcheck] + pub fn test_sort_price_comps(mut comps: Vec) { + let mut rust_std_sorted_comps = comps.clone(); + rust_std_sorted_comps.sort_by_key(|x| x.pub_); + + let num_comps = comps.len(); + assert_eq!( + sort_price_comps(&mut comps, num_comps + 1), + Err(ProgramError::InvalidArgument) + ); + + assert_eq!(sort_price_comps(&mut comps, num_comps), Ok(())); + assert_eq!(comps, rust_std_sorted_comps); + } + + #[quickcheck] + pub fn test_sort_price_comps_smaller_slice( + mut comps: Vec, + mut num_comps: usize, + ) { + num_comps = if comps.is_empty() { + 0 + } else { + num_comps % comps.len() + }; + + let mut rust_std_sorted_comps = comps.get(..num_comps).unwrap().to_vec(); + rust_std_sorted_comps.sort_by_key(|x| x.pub_); + + + assert_eq!(sort_price_comps(&mut comps, num_comps), Ok(())); + assert_eq!(comps.get(..num_comps).unwrap(), rust_std_sorted_comps); + } +} diff --git a/program/rust/src/processor/upd_price.rs b/program/rust/src/processor/upd_price.rs index 5e62f241d..f3a3378c4 100644 --- a/program/rust/src/processor/upd_price.rs +++ b/program/rust/src/processor/upd_price.rs @@ -2,6 +2,7 @@ use { crate::{ accounts::{ PriceAccount, + PriceComponent, PriceInfo, PythOracleSerialize, UPD_PRICE_WRITE_SEED, @@ -31,6 +32,7 @@ use { }, program::invoke_signed, program_error::ProgramError, + program_memory::sol_memcmp, pubkey::Pubkey, sysvar::Sysvar, }, @@ -127,7 +129,7 @@ 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 @@ -137,17 +139,15 @@ pub fn upd_price( // Verify that symbol account is initialized let price_data = load_checked::(price_account, cmd_args.header.version)?; - // Verify that publisher is authorized - while publisher_index < try_convert::(price_data.num_)? { - if price_data.comp_[publisher_index].pub_ == *funding_account.key { - break; + publisher_index = match find_publisher_index( + &price_data.comp_[..try_convert::(price_data.num_)?], + funding_account.key, + ) { + Some(index) => index, + None => { + return Err(OracleError::PermissionViolation.into()); } - publisher_index += 1; - } - pyth_assert( - publisher_index < try_convert::(price_data.num_)?, - OracleError::PermissionViolation.into(), - )?; + }; latest_aggregate_price = price_data.agg_; let latest_publisher_price = price_data.comp_[publisher_index].latest_; @@ -281,6 +281,62 @@ pub fn upd_price( Ok(()) } +/// Find the index of the publisher in the list of components. +/// +/// This method first tries to binary search for the publisher's key in the list of components +/// to get the result faster if the list is sorted. If the list is not sorted, it falls back to +/// a linear search. +#[inline(always)] +fn find_publisher_index(comps: &[PriceComponent], key: &Pubkey) -> Option { + // 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. + 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 = comps.len(); + while left < right { + let mid = left + (right - left) / 2; + match sol_memcmp(comps[mid].pub_.as_ref(), key.as_ref(), 32) { + i if i < 0 => { + left = mid + 1; + } + i if i > 0 => { + right = mid; + } + _ => { + binary_search_result = Some(mid); + break; + } + } + } + } + + match binary_search_result { + Some(index) => Some(index), + None => { + let mut index = 0; + while index < comps.len() { + if sol_memcmp(comps[index].pub_.as_ref(), key.as_ref(), 32) == 0 { + break; + } + index += 1; + } + if index == comps.len() { + None + } else { + Some(index) + } + } + } +} + #[allow(dead_code)] // Wrapper struct for the accounts required to add data to the accumulator program. struct MessageBufferAccounts<'a, 'b: 'a> { @@ -289,3 +345,35 @@ struct MessageBufferAccounts<'a, 'b: 'a> { oracle_auth_pda: &'a AccountInfo<'b>, message_buffer_data: &'a AccountInfo<'b>, } + +#[cfg(test)] +mod test { + use { + super::*, + crate::accounts::PriceComponent, + quickcheck_macros::quickcheck, + solana_program::pubkey::Pubkey, + }; + + /// Test the find_publisher_index method works with an unordered list of components. + #[quickcheck] + pub fn test_find_publisher_index_unordered_comp(comps: Vec) { + comps.iter().enumerate().for_each(|(idx, comp)| { + assert_eq!(find_publisher_index(&comps, &comp.pub_), Some(idx)); + }); + + assert_eq!(find_publisher_index(&comps, &Pubkey::new_unique()), None); + } + + /// Test the find_publisher_index method works with a sorted list of components. + #[quickcheck] + pub fn test_find_publisher_index_ordered_comp(mut comps: Vec) { + comps.sort_by_key(|comp| comp.pub_); + + comps.iter().enumerate().for_each(|(idx, comp)| { + assert_eq!(find_publisher_index(&comps, &comp.pub_), Some(idx)); + }); + + assert_eq!(find_publisher_index(&comps, &Pubkey::new_unique()), None); + } +} diff --git a/program/rust/src/tests/test_add_publisher.rs b/program/rust/src/tests/test_add_publisher.rs index cc7b3c90e..6511acdfb 100644 --- a/program/rust/src/tests/test_add_publisher.rs +++ b/program/rust/src/tests/test_add_publisher.rs @@ -144,7 +144,7 @@ fn test_add_publisher() { { let price_data = load_checked::(&price_account, PC_VERSION).unwrap(); assert_eq!(price_data.num_, i + 1); - assert!(price_data.comp_[i as usize].pub_ == cmd.publisher); + assert!(price_data.comp_.iter().any(|c| c.pub_ == cmd.publisher)); assert_eq!(price_data.header.size, PriceAccount::INITIAL_SIZE); } } @@ -163,4 +163,45 @@ fn test_add_publisher() { ), Err(ProgramError::InvalidArgument) ); + + // Make sure that publishers are sorted + { + let price_data = load_checked::(&price_account, PC_VERSION).unwrap(); + for i in 1..PC_NUM_COMP { + assert!(price_data.comp_[i as usize].pub_ > price_data.comp_[(i - 1) as usize].pub_); + } + } + + // Test sorting by reordering the publishers to be in reverse order + // and then adding the default (000...) publisher to trigger the sorting. + { + let mut price_data = load_checked::(&price_account, PC_VERSION).unwrap(); + price_data + .comp_ + .get_mut(..PC_NUM_COMP as usize) + .unwrap() + .reverse(); + } + + cmd.publisher = Pubkey::default(); + instruction_data = bytes_of::(&cmd); + assert!(process_instruction( + &program_id, + &[ + funding_account.clone(), + price_account.clone(), + permissions_account.clone(), + ], + instruction_data + ) + .is_ok()); + + // Make sure that publishers get sorted after adding the default publisher + { + let price_data = load_checked::(&price_account, PC_VERSION).unwrap(); + println!("{:?}", price_data.comp_.map(|x| x.pub_)); + for i in 1..PC_NUM_COMP { + assert!(price_data.comp_[i as usize].pub_ > price_data.comp_[(i - 1) as usize].pub_); + } + } } diff --git a/program/rust/src/tests/test_upd_price_v2.rs b/program/rust/src/tests/test_upd_price_v2.rs index 0e0c1f4b8..63785819b 100644 --- a/program/rust/src/tests/test_upd_price_v2.rs +++ b/program/rust/src/tests/test_upd_price_v2.rs @@ -444,6 +444,87 @@ fn test_upd_price_v2() -> Result<(), Box> { Ok(()) } +#[test] +fn test_upd_works_with_unordered_publisher_set() -> Result<(), Box> { + let mut instruction_data = [0u8; size_of::()]; + + let program_id = Pubkey::new_unique(); + + let mut price_setup = AccountSetup::new::(&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::(&price_account, PC_VERSION).unwrap(); + price_data.num_ = 20; + // Store the publishers in reverse order + publishers + .iter() + .rev() + .enumerate() + .for_each(|(i, account)| { + 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); + + for (i, publisher) in publishers.iter().enumerate() { + populate_instruction(&mut instruction_data, (i + 100) as i64, 10, 1); + process_instruction( + &program_id, + &[ + publisher.clone(), + price_account.clone(), + clock_account.clone(), + ], + &instruction_data, + )?; + } + + update_clock_slot(&mut clock_account, 2); + + // Trigger the aggregate calculation by sending another price + // update + populate_instruction(&mut instruction_data, 100, 10, 2); + process_instruction( + &program_id, + &[ + publishers[0].clone(), + price_account.clone(), + clock_account.clone(), + ], + &instruction_data, + )?; + + + let price_data = load_checked::(&price_account, PC_VERSION).unwrap(); + + // The result will be the following only if all the + // publishers prices are included in the aggregate. + assert_eq!(price_data.valid_slot_, 1); + assert_eq!(price_data.agg_.pub_slot_, 2); + 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::(instruction_data).unwrap(); From 37e99b88a59b7c2f33d74fd2821727dc267ce239 Mon Sep 17 00:00:00 2001 From: Ali Behjati Date: Mon, 13 May 2024 16:49:43 +0200 Subject: [PATCH 2/3] refactor: use heapsort for each publisher addition too --- program/rust/src/processor/add_publisher.rs | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/program/rust/src/processor/add_publisher.rs b/program/rust/src/processor/add_publisher.rs index 015884ff2..841de3b96 100644 --- a/program/rust/src/processor/add_publisher.rs +++ b/program/rust/src/processor/add_publisher.rs @@ -83,23 +83,21 @@ pub fn add_publisher( } } - let mut current_index: usize = try_convert(price_data.num_)?; + let current_index: usize = try_convert(price_data.num_)?; sol_memset( bytes_of_mut(&mut price_data.comp_[current_index]), 0, size_of::(), ); price_data.comp_[current_index].pub_ = cmd_args.publisher; + price_data.num_ += 1; - // 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_ + // Sort the publishers in the list { - price_data.comp_.swap(current_index, current_index - 1); - current_index -= 1; + let num_comps = try_convert::(price_data.num_)?; + sort_price_comps(&mut price_data.comp_, num_comps)?; } - price_data.num_ += 1; price_data.header.size = try_convert::<_, u32>(PriceAccount::INITIAL_SIZE)?; Ok(()) } @@ -107,7 +105,6 @@ pub fn add_publisher( /// A copy of rust slice/sort.rs heapsort implementation which is small and fast. We couldn't use /// the sort directly because it was only accessible behind a unstable feature flag at the time of /// writing this code. -#[inline(always)] fn heapsort(v: &mut [(Pubkey, usize)]) { // This binary heap respects the invariant `parent >= child`. let sift_down = |v: &mut [(Pubkey, usize)], mut node: usize| { @@ -155,16 +152,19 @@ fn heapsort(v: &mut [(Pubkey, usize)]) { /// /// num_publishers is the number of publishers in the list that should be sorted. It is explicitly /// passed to avoid callers mistake of passing the full slice which may contain uninitialized values. -#[inline(always)] fn sort_price_comps(comps: &mut [PriceComponent], num_comps: usize) -> Result<(), ProgramError> { let comps = comps .get_mut(..num_comps) .ok_or(ProgramError::InvalidArgument)?; + // Publishers are likely sorted in ascending order but + // heapsorts creates a max-heap so we reverse the order + // of the keys to make the heapify step faster. let mut keys = comps .iter() .enumerate() .map(|(i, x)| (x.pub_, i)) + .rev() .collect::>(); heapsort(&mut keys); From b991cbac89c806281c51f467fe9579e9ef11a12b Mon Sep 17 00:00:00 2001 From: Ali Behjati Date: Mon, 13 May 2024 20:25:14 +0200 Subject: [PATCH 3/3] fix: address comments --- program/rust/src/processor/upd_price.rs | 16 ++++++++++++++-- program/rust/src/tests/test_add_publisher.rs | 2 +- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/program/rust/src/processor/upd_price.rs b/program/rust/src/processor/upd_price.rs index f3a3378c4..278673b4f 100644 --- a/program/rust/src/processor/upd_price.rs +++ b/program/rust/src/processor/upd_price.rs @@ -303,6 +303,8 @@ fn find_publisher_index(comps: &[PriceComponent], key: &Pubkey) -> Option let mut right = comps.len(); while left < right { let mid = left + (right - left) / 2; + // sol_memcmp is much faster than rust default comparison of Pubkey. It costs + // 10CU whereas rust default comparison costs a few times more. match sol_memcmp(comps[mid].pub_.as_ref(), key.as_ref(), 32) { i if i < 0 => { left = mid + 1; @@ -362,7 +364,12 @@ mod test { assert_eq!(find_publisher_index(&comps, &comp.pub_), Some(idx)); }); - assert_eq!(find_publisher_index(&comps, &Pubkey::new_unique()), None); + let mut key_not_in_list = Pubkey::new_unique(); + while comps.iter().any(|comp| comp.pub_ == key_not_in_list) { + key_not_in_list = Pubkey::new_unique(); + } + + assert_eq!(find_publisher_index(&comps, &key_not_in_list), None); } /// Test the find_publisher_index method works with a sorted list of components. @@ -374,6 +381,11 @@ mod test { assert_eq!(find_publisher_index(&comps, &comp.pub_), Some(idx)); }); - assert_eq!(find_publisher_index(&comps, &Pubkey::new_unique()), None); + let mut key_not_in_list = Pubkey::new_unique(); + while comps.iter().any(|comp| comp.pub_ == key_not_in_list) { + key_not_in_list = Pubkey::new_unique(); + } + + assert_eq!(find_publisher_index(&comps, &key_not_in_list), None); } } diff --git a/program/rust/src/tests/test_add_publisher.rs b/program/rust/src/tests/test_add_publisher.rs index 6511acdfb..45d6172c3 100644 --- a/program/rust/src/tests/test_add_publisher.rs +++ b/program/rust/src/tests/test_add_publisher.rs @@ -199,7 +199,7 @@ fn test_add_publisher() { // Make sure that publishers get sorted after adding the default publisher { let price_data = load_checked::(&price_account, PC_VERSION).unwrap(); - println!("{:?}", price_data.comp_.map(|x| x.pub_)); + assert!(price_data.num_ == PC_NUM_COMP); for i in 1..PC_NUM_COMP { assert!(price_data.comp_[i as usize].pub_ > price_data.comp_[(i - 1) as usize].pub_); }