use std::{any::Any, str::FromStr, sync::Arc, time::Instant};
use rustc_hash::FxHashMap;
use tokio::sync::RwLock;
use tracing::warn;
use tycho_simulation::tycho_common::models::Address;
use super::{
computation::{ComputationId, DerivedComputation, FailedItem, FailedItemError},
computations::{ComponentDepthComputation, SpotPriceComputation, TokenGasPriceComputation},
types::{
ComponentDepthKey, ComponentDepths, SpotPriceKey, SpotPrices, TokenGasPriceKey,
TokenGasPrices, TokenPricesWithDeps,
},
};
use crate::derived::SharedDerivedDataRef;
#[derive(Debug)]
struct ComputedValue<T> {
data: T,
block: u64,
}
#[derive(Debug, Clone, Copy)]
pub struct ComputationStatus {
block: u64,
updated_at: Instant,
}
impl ComputationStatus {
pub fn block(&self) -> u64 {
self.block
}
pub fn age_ms_at(&self, now: Instant) -> u64 {
let elapsed = now.saturating_duration_since(self.updated_at);
u64::try_from(elapsed.as_millis()).unwrap_or_else(|_| {
warn!("Failed to compute age_ms at {:?}", elapsed);
u64::MAX
})
}
}
struct ComputedSlot {
data: Box<dyn Any + Send + Sync>,
status: ComputationStatus,
}
impl std::fmt::Debug for ComputedSlot {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ComputedSlot")
.field("status", &self.status)
.finish_non_exhaustive()
}
}
#[derive(Debug, Default)]
pub struct DerivedData {
slots: FxHashMap<ComputationId, ComputedSlot>,
token_prices_failed: FxHashMap<TokenGasPriceKey, (u64, FailedItemError)>,
token_prices_deps: Option<ComputedValue<TokenPricesWithDeps>>,
component_depths_failed: FxHashMap<ComponentDepthKey, (u64, FailedItemError)>,
spot_prices_failed: FxHashMap<SpotPriceKey, (u64, FailedItemError)>,
}
fn parse_pair_key(s: &str) -> Option<(String, Address, Address)> {
let mut parts = s.rsplitn(3, '/');
let token_out_str = parts.next()?;
let token_in_str = parts.next()?;
let component_id = parts.next()?;
let token_in = Address::from_str(token_in_str).ok()?;
let token_out = Address::from_str(token_out_str).ok()?;
Some((component_id.to_string(), token_in, token_out))
}
impl DerivedData {
pub fn new() -> Self {
Self::default()
}
pub fn new_shared() -> SharedDerivedDataRef {
Arc::new(RwLock::new(Self::new()))
}
pub(crate) fn set_output<T: Any + Send + Sync>(
&mut self,
id: ComputationId,
data: T,
block: u64,
) {
self.set_output_at(id, data, block, Instant::now());
}
fn set_output_at<T: Any + Send + Sync>(
&mut self,
id: ComputationId,
data: T,
block: u64,
updated_at: Instant,
) {
self.slots.insert(
id,
ComputedSlot { data: Box::new(data), status: ComputationStatus { block, updated_at } },
);
}
pub(crate) fn output<T: Any>(&self, id: ComputationId) -> Option<&T> {
self.output_with_status(id)
.map(|(data, _)| data)
}
fn output_with_status<T: Any>(&self, id: ComputationId) -> Option<(&T, ComputationStatus)> {
let slot = self.slots.get(&id)?;
debug_assert!(slot.data.is::<T>(), "derived output {id} read as the wrong type");
slot.data
.downcast_ref::<T>()
.map(|data| (data, slot.status))
}
pub(crate) fn output_block(&self, id: ComputationId) -> Option<u64> {
self.output_status(id)
.map(|status| status.block())
}
fn output_status(&self, id: ComputationId) -> Option<ComputationStatus> {
self.slots
.get(id)
.map(|slot| slot.status)
}
fn clear_output(&mut self, id: ComputationId) {
self.slots.remove(id);
}
pub fn derived_data_ready(&self) -> bool {
self.token_prices_block().is_some() &&
self.token_prices_deps_block().is_some() &&
self.component_depths_block().is_some() &&
self.spot_prices_block().is_some()
}
pub fn token_prices(&self) -> Option<&TokenGasPrices> {
self.token_prices_slot()
.map(Arc::as_ref)
}
pub fn token_prices_shared(&self) -> Option<Arc<TokenGasPrices>> {
self.token_prices_slot().cloned()
}
fn token_prices_slot(&self) -> Option<&Arc<TokenGasPrices>> {
self.output(TokenGasPriceComputation::ID)
}
pub fn token_prices_block(&self) -> Option<u64> {
self.output_block(TokenGasPriceComputation::ID)
}
pub fn token_prices_with_status(&self) -> Option<(&Arc<TokenGasPrices>, ComputationStatus)> {
self.output_with_status(TokenGasPriceComputation::ID)
}
pub fn set_token_prices(
&mut self,
prices: TokenGasPrices,
failed_items: Vec<FailedItem>,
block: u64,
is_full_recompute: bool,
) {
let new_failures: FxHashMap<TokenGasPriceKey, (u64, FailedItemError)> = failed_items
.into_iter()
.filter_map(|f| {
Address::from_str(&f.key)
.ok()
.map(|k| (k, (block, f.error)))
})
.collect();
if is_full_recompute {
self.token_prices_failed = new_failures;
} else {
self.token_prices_failed
.retain(|k, _| !prices.contains_key(k));
self.token_prices_failed
.extend(new_failures);
}
self.set_output(TokenGasPriceComputation::ID, Arc::new(prices), block);
}
pub fn token_price_failure(&self, key: &TokenGasPriceKey) -> Option<(u64, &FailedItemError)> {
self.token_prices_failed
.get(key)
.map(|(block, error)| (*block, error))
}
pub fn clear_token_prices(&mut self) {
self.clear_output(TokenGasPriceComputation::ID);
self.token_prices_failed.clear();
}
pub fn token_prices_deps(&self) -> Option<&TokenPricesWithDeps> {
self.token_prices_deps
.as_ref()
.map(|v| &v.data)
}
pub fn token_prices_deps_block(&self) -> Option<u64> {
self.token_prices_deps
.as_ref()
.map(|v| v.block)
}
pub fn set_token_prices_deps(&mut self, prices: TokenPricesWithDeps, block: u64) {
self.token_prices_deps = Some(ComputedValue { data: prices, block });
}
pub fn clear_token_prices_deps(&mut self) {
self.token_prices_deps = None;
}
pub fn component_depths(&self) -> Option<&ComponentDepths> {
self.output(ComponentDepthComputation::ID)
}
pub fn component_depths_block(&self) -> Option<u64> {
self.output_block(ComponentDepthComputation::ID)
}
pub fn component_depths_with_status(&self) -> Option<(&ComponentDepths, ComputationStatus)> {
self.output_with_status(ComponentDepthComputation::ID)
}
pub fn set_component_depths(
&mut self,
depths: ComponentDepths,
failed_items: Vec<FailedItem>,
block: u64,
is_full_recompute: bool,
) {
let new_failures: FxHashMap<ComponentDepthKey, (u64, FailedItemError)> = failed_items
.into_iter()
.filter_map(|f| parse_pair_key(&f.key).map(|k| (k, (block, f.error))))
.collect();
if is_full_recompute {
self.component_depths_failed = new_failures;
} else {
self.component_depths_failed
.retain(|k, _| !depths.contains_key(k));
self.component_depths_failed
.extend(new_failures);
}
self.set_output(ComponentDepthComputation::ID, depths, block);
}
pub fn component_depth_failure(
&self,
key: &ComponentDepthKey,
) -> Option<(u64, &FailedItemError)> {
self.component_depths_failed
.get(key)
.map(|(block, error)| (*block, error))
}
pub fn clear_component_depths(&mut self) {
self.clear_output(ComponentDepthComputation::ID);
self.component_depths_failed.clear();
}
pub fn spot_prices(&self) -> Option<&SpotPrices> {
self.output(SpotPriceComputation::ID)
}
pub fn spot_prices_block(&self) -> Option<u64> {
self.output_block(SpotPriceComputation::ID)
}
pub fn spot_prices_with_status(&self) -> Option<(&SpotPrices, ComputationStatus)> {
self.output_with_status(SpotPriceComputation::ID)
}
pub fn set_spot_prices(
&mut self,
prices: SpotPrices,
failed_items: Vec<FailedItem>,
block: u64,
is_full_recompute: bool,
) {
let new_failures: FxHashMap<SpotPriceKey, (u64, FailedItemError)> = failed_items
.into_iter()
.filter_map(|f| parse_pair_key(&f.key).map(|k| (k, (block, f.error))))
.collect();
if is_full_recompute {
self.spot_prices_failed = new_failures;
} else {
self.spot_prices_failed
.retain(|k, _| !prices.contains_key(k));
self.spot_prices_failed
.extend(new_failures);
}
self.set_output(SpotPriceComputation::ID, prices, block);
}
pub fn spot_price_failure(&self, key: &SpotPriceKey) -> Option<(u64, &FailedItemError)> {
self.spot_prices_failed
.get(key)
.map(|(block, error)| (*block, error))
}
pub fn clear_spot_prices(&mut self) {
self.clear_output(SpotPriceComputation::ID);
self.spot_prices_failed.clear();
}
pub fn clear_all(&mut self) {
self.slots.clear();
self.token_prices_failed.clear();
self.token_prices_deps = None;
self.component_depths_failed.clear();
self.spot_prices_failed.clear();
}
}
#[cfg(test)]
mod tests {
use std::time::{Duration, Instant};
use super::*;
use crate::{algorithm::test_utils::addr, derived::types::SpotPrices};
fn failed(key: &str, error: FailedItemError) -> FailedItem {
FailedItem { key: key.to_string(), error }
}
fn pair_key(comp: &str, b_in: u8, b_out: u8) -> SpotPriceKey {
(comp.to_string(), addr(b_in), addr(b_out))
}
#[test]
fn test_token_prices_status_age() {
let updated_at = Instant::now();
let mut store = DerivedData::new();
store.set_output_at(
TokenGasPriceComputation::ID,
Arc::new(TokenGasPrices::default()),
42,
updated_at,
);
let (prices, status) = store
.token_prices_with_status()
.unwrap();
assert!(prices.is_empty());
assert_eq!(status.block(), 42);
assert_eq!(status.age_ms_at(updated_at + Duration::from_millis(1_234)), 1_234);
}
#[test]
fn test_computation_statuses_update_independently() {
let started_at = Instant::now();
let mut store = DerivedData::new();
store.set_output_at(
TokenGasPriceComputation::ID,
Arc::new(TokenGasPrices::default()),
10,
started_at,
);
store.set_output_at(
SpotPriceComputation::ID,
SpotPrices::default(),
11,
started_at + Duration::from_millis(10),
);
store.set_output_at(
ComponentDepthComputation::ID,
ComponentDepths::default(),
12,
started_at + Duration::from_millis(20),
);
store.set_output_at(
TokenGasPriceComputation::ID,
Arc::new(TokenGasPrices::default()),
20,
started_at + Duration::from_millis(30),
);
let compared_at = started_at + Duration::from_millis(50);
let (_, token_prices) = store
.token_prices_with_status()
.unwrap();
assert_eq!(token_prices.block(), 20);
assert_eq!(token_prices.age_ms_at(compared_at), 20);
let (_, spot_prices) = store.spot_prices_with_status().unwrap();
assert_eq!(spot_prices.block(), 11);
assert_eq!(spot_prices.age_ms_at(compared_at), 40);
let (_, component_depths) = store
.component_depths_with_status()
.unwrap();
assert_eq!(component_depths.block(), 12);
assert_eq!(component_depths.age_ms_at(compared_at), 30);
}
#[test]
fn test_computation_status_earlier_comparison() {
let updated_at = Instant::now();
let mut store = DerivedData::new();
store.set_output_at(
TokenGasPriceComputation::ID,
Arc::new(TokenGasPrices::default()),
10,
updated_at,
);
let compared_at = updated_at
.checked_sub(Duration::from_millis(1))
.unwrap();
assert_eq!(
store
.token_prices_with_status()
.unwrap()
.1
.age_ms_at(compared_at),
0
);
}
#[test]
fn test_token_prices_block_tracks_independently() {
let mut store = DerivedData::new();
assert_eq!(store.token_prices_block(), None);
store.set_token_prices(Default::default(), vec![], 42, true);
assert_eq!(store.token_prices_block(), Some(42));
assert_eq!(store.spot_prices_block(), None);
assert_eq!(store.component_depths_block(), None);
}
#[test]
fn test_spot_prices_block_tracks_independently() {
let mut store = DerivedData::new();
store.set_spot_prices(Default::default(), vec![], 10, true);
assert_eq!(store.spot_prices_block(), Some(10));
assert_eq!(store.token_prices_block(), None);
}
#[test]
fn test_component_depths_block_tracks_independently() {
let mut store = DerivedData::new();
store.set_component_depths(Default::default(), vec![], 7, true);
assert_eq!(store.component_depths_block(), Some(7));
assert_eq!(store.token_prices_block(), None);
}
#[test]
fn test_derived_data_ready() {
let mut store = DerivedData::new();
assert!(!store.derived_data_ready());
store.set_spot_prices(Default::default(), vec![], 5, true);
assert!(!store.derived_data_ready());
store.set_token_prices(Default::default(), vec![], 10, true);
assert!(!store.derived_data_ready());
store.set_token_prices_deps(Default::default(), 10);
assert!(!store.derived_data_ready());
store.set_component_depths(Default::default(), vec![], 9, true);
assert!(store.derived_data_ready());
}
#[test]
fn test_clear_all_resets_all_fields() {
let mut store = DerivedData::new();
store.set_token_prices(Default::default(), vec![], 1, true);
store.set_spot_prices(Default::default(), vec![], 1, true);
store.set_component_depths(Default::default(), vec![], 1, true);
store.clear_all();
assert!(store
.token_prices_with_status()
.is_none());
assert!(store
.spot_prices_with_status()
.is_none());
assert!(store
.component_depths_with_status()
.is_none());
assert!(!store.derived_data_ready());
}
#[test]
fn test_token_price_failure_stored_with_block() {
let token_addr = addr(0xab);
let key_str = format!("{token_addr}");
let mut store = DerivedData::new();
store.set_token_prices(
Default::default(),
vec![failed(&key_str, FailedItemError::SimulationFailed("sim error".into()))],
42,
true,
);
assert_eq!(
store.token_price_failure(&token_addr),
Some((42, &FailedItemError::SimulationFailed("sim error".into())))
);
assert_eq!(store.token_price_failure(&addr(0xcd)), None);
}
#[test]
fn test_spot_price_failure_stored_with_block() {
let key = pair_key("component1", 0x01, 0x02);
let key_str = format!("component1/{}/{}", addr(0x01), addr(0x02));
let mut store = DerivedData::new();
store.set_spot_prices(
Default::default(),
vec![failed(&key_str, FailedItemError::SimulationFailed("sim error".into()))],
10,
true,
);
assert_eq!(
store.spot_price_failure(&key),
Some((10, &FailedItemError::SimulationFailed("sim error".into())))
);
assert_eq!(store.spot_price_failure(&pair_key("component1", 0x01, 0x03)), None);
}
#[test]
fn test_component_depth_failure_stored_with_block() {
let key: ComponentDepthKey = pair_key("component1", 0x01, 0x02);
let key_str = format!("component1/{}/{}", addr(0x01), addr(0x02));
let mut store = DerivedData::new();
store.set_component_depths(
Default::default(),
vec![failed(&key_str, FailedItemError::SimulationFailed("depth error".into()))],
7,
true,
);
assert_eq!(
store.component_depth_failure(&key),
Some((7, &FailedItemError::SimulationFailed("depth error".into())))
);
assert_eq!(store.component_depth_failure(&pair_key("component2", 0x01, 0x02)), None);
}
#[test]
fn test_rerunning_with_empty_failures_clears_old_reasons() {
let key = pair_key("component1", 0x01, 0x02);
let key_str = format!("component1/{}/{}", addr(0x01), addr(0x02));
let mut store = DerivedData::new();
store.set_spot_prices(
Default::default(),
vec![failed(&key_str, FailedItemError::MissingSimulationState)],
1,
true,
);
assert!(store.spot_price_failure(&key).is_some());
store.set_spot_prices(Default::default(), vec![], 2, true);
assert_eq!(store.spot_price_failure(&key), None);
}
#[test]
fn test_clear_token_prices_clears_failure_map() {
let token_addr = addr(0xab);
let key_str = format!("{token_addr}");
let mut store = DerivedData::new();
store.set_token_prices(
Default::default(),
vec![failed(&key_str, FailedItemError::MissingSellRoute("no route".to_string()))],
1,
true,
);
store.clear_token_prices();
assert_eq!(store.token_price_failure(&token_addr), None);
assert!(store
.token_prices_with_status()
.is_none());
}
#[test]
fn test_clear_spot_prices_clears_failure_map() {
let key = pair_key("component1", 0x01, 0x02);
let key_str = format!("component1/{}/{}", addr(0x01), addr(0x02));
let mut store = DerivedData::new();
store.set_spot_prices(
Default::default(),
vec![failed(&key_str, FailedItemError::MissingSimulationState)],
1,
true,
);
store.clear_spot_prices();
assert_eq!(store.spot_price_failure(&key), None);
}
#[test]
fn test_clear_component_depths_clears_failure_map() {
let key: ComponentDepthKey = pair_key("component1", 0x01, 0x02);
let key_str = format!("component1/{}/{}", addr(0x01), addr(0x02));
let mut store = DerivedData::new();
store.set_component_depths(
Default::default(),
vec![failed(&key_str, FailedItemError::MissingSpotPrice)],
1,
true,
);
store.clear_component_depths();
assert_eq!(store.component_depth_failure(&key), None);
}
#[test]
fn test_incremental_run_preserves_failures_for_unattempted_items() {
let key_a = pair_key("component_a", 0x01, 0x02);
let key_a_str = format!("component_a/{}/{}", addr(0x01), addr(0x02));
let key_b = pair_key("component_b", 0x03, 0x04);
let key_b_str = format!("component_b/{}/{}", addr(0x03), addr(0x04));
let mut store = DerivedData::new();
store.set_spot_prices(
Default::default(),
vec![
failed(&key_a_str, FailedItemError::MissingSimulationState),
failed(&key_b_str, FailedItemError::MissingTokenMetadata),
],
10,
true,
);
assert_eq!(
store.spot_price_failure(&key_a),
Some((10, &FailedItemError::MissingSimulationState))
);
assert_eq!(
store.spot_price_failure(&key_b),
Some((10, &FailedItemError::MissingTokenMetadata))
);
let mut prices = SpotPrices::default();
prices.insert(key_b.clone(), 1.0);
store.set_spot_prices(prices, vec![], 11, false);
assert_eq!(
store.spot_price_failure(&key_a),
Some((10, &FailedItemError::MissingSimulationState))
);
assert_eq!(store.spot_price_failure(&key_b), None);
}
#[test]
fn test_incremental_run_updates_block_on_repeated_failure() {
let key = pair_key("component_a", 0x01, 0x02);
let key_str = format!("component_a/{}/{}", addr(0x01), addr(0x02));
let mut store = DerivedData::new();
store.set_spot_prices(
Default::default(),
vec![failed(&key_str, FailedItemError::MissingSimulationState)],
10,
true,
);
assert_eq!(
store.spot_price_failure(&key),
Some((10, &FailedItemError::MissingSimulationState))
);
store.set_spot_prices(
Default::default(),
vec![failed(&key_str, FailedItemError::MissingTokenMetadata)],
11,
false,
);
assert_eq!(
store.spot_price_failure(&key),
Some((11, &FailedItemError::MissingTokenMetadata))
);
}
#[test]
fn test_clear_all_clears_all_failure_maps() {
let token_addr = addr(0xab);
let token_str = format!("{token_addr}");
let pair = pair_key("component1", 0x01, 0x02);
let pair_str = format!("component1/{}/{}", addr(0x01), addr(0x02));
let mut store = DerivedData::new();
store.set_token_prices(
Default::default(),
vec![failed(&token_str, FailedItemError::MissingSellRoute("no route".to_string()))],
1,
true,
);
store.set_spot_prices(
Default::default(),
vec![failed(&pair_str, FailedItemError::MissingSimulationState)],
1,
true,
);
store.set_component_depths(
Default::default(),
vec![failed(&pair_str, FailedItemError::MissingSpotPrice)],
1,
true,
);
store.clear_all();
assert_eq!(store.token_price_failure(&token_addr), None);
assert_eq!(store.spot_price_failure(&pair), None);
assert_eq!(store.component_depth_failure(&pair), None);
}
}