diff --git a/state/src/caching/order_metadata_index.rs b/state/src/caching/order_metadata_index.rs index 64a3b228..ed5001c3 100644 --- a/state/src/caching/order_metadata_index.rs +++ b/state/src/caching/order_metadata_index.rs @@ -22,13 +22,19 @@ type OrderSideIndex = HashMap>; pub struct OrderMetadataIndex { /// The mapping from pair -> side -> (matchable_amount, order_id) index: RwLockHashMap, + /// A reverse mapping from order_id to pair and side + /// + /// This is used to efficiently query the index by order_id for updates and + /// deletion + reverse_index: RwLockHashMap, } impl OrderMetadataIndex { /// Construct a new order metadata index pub fn new() -> Self { let index = RwLock::new(HashMap::new()); - Self { index } + let reverse_index = RwLock::new(HashMap::new()); + Self { index, reverse_index } } // --- Getters --- // @@ -43,6 +49,11 @@ impl OrderMetadataIndex { } } + /// Get the pair and side for a given order_id + pub async fn get_pair_and_side(&self, order_id: &OrderIdentifier) -> Option<(Pair, OrderSide)> { + self.reverse_index.read().await.get(order_id).cloned() + } + // --- Setters --- // /// Add an order to the index @@ -54,8 +65,56 @@ impl OrderMetadataIndex { ) { let (pair, side) = order.pair_and_side(); let mut index = self.index.write().await; - let entry = index.entry(pair).or_insert_with(OrderSideIndex::new); + let entry = index.entry(pair.clone()).or_insert_with(OrderSideIndex::new); entry.entry(side).or_insert_with(SortedVec::new).insert((matchable_amount, order_id)); + + // Update the reverse index + let mut reverse_index = self.reverse_index.write().await; + reverse_index.insert(order_id, (pair, side)); + } + + /// Update the matchable amount for an order + pub async fn update_matchable_amount( + &self, + order_id: OrderIdentifier, + matchable_amount: Amount, + ) { + let (pair, side) = self.get_pair_and_side(&order_id).await.unwrap(); + let mut index = self.index.write().await; + let entry = index.entry(pair).or_insert_with(OrderSideIndex::new); + let sorted_vec = entry.entry(side).or_insert_with(SortedVec::new); + + // Remove the old entry + if let Some(idx) = sorted_vec.find_index(|(_, oid)| *oid == order_id) { + sorted_vec.remove(idx); + } + + // Insert the new entry (this will maintain the sort order) + sorted_vec.insert((matchable_amount, order_id)); + } + + /// Remove an order from the index + /// + /// Note that we do not clean up sub-index entries when their + /// lists become empty. + pub async fn remove_order(&self, order_id: &OrderIdentifier) -> Option<()> { + // Remove from the reverse index + let mut reverse_index = self.reverse_index.write().await; + reverse_index.remove(order_id); + + // Get the pair and side from the reverse index + let (pair, side) = self.get_pair_and_side(order_id).await?; + + // Remove from the main index + let mut index = self.index.write().await; + let side_index = index.get_mut(&pair)?; + let sorted_vec = side_index.get_mut(&side)?; + + if let Some(idx) = sorted_vec.find_index(|(_, oid)| oid == order_id) { + sorted_vec.remove(idx); + } + + Some(()) } } @@ -73,8 +132,19 @@ impl SortedVec { // --- Getters --- // + /// Get the element at index i + #[cfg(test)] + pub fn get(&self, i: usize) -> Option<&T> { + self.vec.get(i) + } + + /// Find the index of an element in the vector using the given filter method + pub fn find_index(&self, filter: impl Fn(&T) -> bool) -> Option { + self.vec.iter().position(filter) + } + /// Get the vector - #[allow(unused)] + #[cfg(test)] pub fn vec(&self) -> &Vec { &self.vec } @@ -98,6 +168,11 @@ impl SortedVec { }; self.vec.insert(index, element); } + + /// Remove an element from the vector + pub fn remove(&mut self, index: usize) -> T { + self.vec.remove(index) + } } #[cfg(test)] @@ -259,4 +334,123 @@ mod order_index_tests { assert_eq!(orders2_correct_side.len(), 2); assert!(orders2_incorrect_side.is_empty()); } + + #[tokio::test] + async fn test_update_matchable_amount() { + let index = OrderMetadataIndex::new(); + let order_id = OrderIdentifier::new_v4(); + let order = mock_order(); + let pair = order.pair(); + + // Add an order with an initial matchable amount + let initial_amount = 100; + index.add_order(order_id, &order, initial_amount).await; + + // Update the matchable amount + let updated_amount = 200; + index.update_matchable_amount(order_id, updated_amount).await; + + // Get the orders and verify the order is in the correct position + let orders = index.get_orders(&pair, order.side).await; + assert_eq!(orders.len(), 1); + assert_eq!(orders[0], order_id); + + // Verify the updated amount by checking the internal state + let index_map = index.index.read().await; + let side_index = index_map.get(&pair).unwrap(); + let sorted_vec = side_index.get(&OrderSide::Buy).unwrap(); + assert_eq!(sorted_vec.get(0).unwrap().0, updated_amount); + } + + #[tokio::test] + async fn test_update_matchable_amount_sort_order() { + let index = OrderMetadataIndex::new(); + let order = mock_order(); + let pair = order.pair(); + + // Add two orders with initial amounts + let order_id1 = OrderIdentifier::new_v4(); + let order_id2 = OrderIdentifier::new_v4(); + + index.add_order(order_id1, &order, 200).await; + index.add_order(order_id2, &order, 100).await; + + // Verify initial sort order (descending by amount) + let orders = index.get_orders(&pair, OrderSide::Buy).await; + assert_eq!(orders.len(), 2); + assert_eq!(orders[0], order_id1); // 200 + assert_eq!(orders[1], order_id2); // 100 + + // Update order2's amount to be larger than order1 + index.update_matchable_amount(order_id2, 300).await; + + // Verify the sort order has changed + let orders = index.get_orders(&pair, OrderSide::Buy).await; + assert_eq!(orders.len(), 2); + assert_eq!(orders[0], order_id2); // 300 + assert_eq!(orders[1], order_id1); // 200 + } + + #[tokio::test] + async fn test_remove_order() { + let index = OrderMetadataIndex::new(); + let order_id = OrderIdentifier::new_v4(); + let order = mock_order(); + let pair = order.pair(); + + // Add an order + index.add_order(order_id, &order, 100).await; + + // Verify it was added + let orders = index.get_orders(&pair, order.side).await; + assert_eq!(orders.len(), 1); + assert_eq!(orders[0], order_id); + + // Remove the order + let result = index.remove_order(&order_id).await; + assert!(result.is_some()); + + // Verify it was removed + let orders = index.get_orders(&pair, order.side).await; + assert!(orders.is_empty()); + + // Verify it was removed from reverse index + let pair_and_side = index.get_pair_and_side(&order_id).await; + assert!(pair_and_side.is_none()); + } + + #[tokio::test] + async fn test_remove_nonexistent_order() { + let index = OrderMetadataIndex::new(); + let order_id = OrderIdentifier::new_v4(); + + // Try to remove a nonexistent order + let result = index.remove_order(&order_id).await; + assert!(result.is_none()); + } + + #[tokio::test] + async fn test_remove_order_maintains_sort() { + let index = OrderMetadataIndex::new(); + let order = mock_order(); + let pair = order.pair(); + + // Add three orders + let order_id1 = OrderIdentifier::new_v4(); + let order_id2 = OrderIdentifier::new_v4(); + let order_id3 = OrderIdentifier::new_v4(); + + index.add_order(order_id1, &order, 300).await; + index.add_order(order_id2, &order, 200).await; + index.add_order(order_id3, &order, 100).await; + + // Remove the middle order + index.remove_order(&order_id2).await; + + // Verify remaining orders are still sorted + let orders = index.get_orders(&pair, OrderSide::Buy).await; + assert_eq!(orders.len(), 2); + assert_eq!(orders[0], order_id1); // 300 + assert_eq!(orders[1], order_id3); // 100 + } }