#[cfg(feature = "swap-metrics")]
use std::{
cell::RefCell,
cmp::Reverse,
time::{Duration, Instant},
};
#[cfg(feature = "swap-metrics")]
use metrics::counter;
use num_bigint::BigUint;
#[cfg(feature = "swap-metrics")]
use rustc_hash::{FxHashMap, FxHashSet};
#[cfg(feature = "swap-metrics")]
use tracing::{debug, enabled, Level};
use tycho_simulation::tycho_common::{
models::token::Token,
simulation::{
errors::SimulationError,
protocol_sim::{GetAmountOutResult, ProtocolSim},
},
};
use super::sim_guard::GuardedProtocolSim;
use crate::{feed::market_data::MarketState, types::ComponentId};
pub(crate) type StageLabel = &'static str;
#[cfg(feature = "swap-metrics")]
#[derive(Default, Clone, Copy)]
struct ComponentSwaps {
calls: u64,
failed: u64,
cache_hits: u64,
interpolated: u64,
refused_without_calling: u64,
call_time: Duration,
}
#[cfg(feature = "swap-metrics")]
impl ComponentSwaps {
fn add(&mut self, other: &ComponentSwaps) {
self.calls += other.calls;
self.failed += other.failed;
self.cache_hits += other.cache_hits;
self.interpolated += other.interpolated;
self.refused_without_calling += other.refused_without_calling;
self.call_time += other.call_time;
}
}
#[cfg(feature = "swap-metrics")]
thread_local! {
static SOLVE_SWAPS: RefCell<FxHashMap<(ComponentId, StageLabel), ComponentSwaps>> =
RefCell::new(FxHashMap::default());
}
#[cfg(feature = "swap-metrics")]
fn with_counts(
component_id: &ComponentId,
stage: StageLabel,
edit: impl FnOnce(&mut ComponentSwaps),
) {
SOLVE_SWAPS.with_borrow_mut(|swaps| {
edit(
swaps
.entry((component_id.clone(), stage))
.or_default(),
);
});
}
#[cfg(feature = "swap-metrics")]
pub(crate) fn start_solve() {
SOLVE_SWAPS.with_borrow_mut(FxHashMap::clear);
}
#[cfg(feature = "swap-metrics")]
fn record_call(component_id: &ComponentId, stage: StageLabel, call_time: Duration, failed: bool) {
with_counts(component_id, stage, |counts| {
counts.calls += 1;
counts.call_time += call_time;
if failed {
counts.failed += 1;
}
});
}
#[cfg(feature = "swap-metrics")]
pub(crate) fn record_cache_hit(component_id: &ComponentId, stage: StageLabel) {
with_counts(component_id, stage, |counts| counts.cache_hits += 1);
}
#[cfg(feature = "swap-metrics")]
pub(crate) fn record_interpolation(component_id: &ComponentId, stage: StageLabel) {
with_counts(component_id, stage, |counts| counts.interpolated += 1);
}
#[cfg(feature = "swap-metrics")]
pub(crate) fn record_refusal_without_calling(component_id: &ComponentId, stage: StageLabel) {
with_counts(component_id, stage, |counts| counts.refused_without_calling += 1);
}
#[cfg(feature = "swap-metrics")]
pub(crate) fn report(market: &MarketState, solve_time_ms: impl FnOnce() -> u64) {
let solve_time_ms = solve_time_ms();
SOLVE_SWAPS.with_borrow(|swaps| report_swaps(swaps, market, solve_time_ms));
}
#[cfg(feature = "swap-metrics")]
fn report_swaps(
swaps: &FxHashMap<(ComponentId, StageLabel), ComponentSwaps>,
market: &MarketState,
solve_time_ms: u64,
) {
let mut by_protocol: FxHashMap<&str, ComponentSwaps> = FxHashMap::default();
for ((component_id, _), counts) in swaps {
let protocol = market
.get_component(component_id)
.map_or("unknown", |component| component.protocol_system.as_str());
by_protocol
.entry(protocol)
.or_default()
.add(counts);
}
let mut costliest_first: Vec<(&str, ComponentSwaps)> = by_protocol.into_iter().collect();
costliest_first
.sort_unstable_by_key(|(protocol, counts)| (Reverse(counts.call_time), *protocol));
for (protocol, counts) in &costliest_first {
counter!("water_fill.get_amount_out_calls", "protocol" => protocol.to_string())
.increment(counts.calls);
counter!("water_fill.failed_calls", "protocol" => protocol.to_string())
.increment(counts.failed);
counter!("water_fill.cache_hits", "protocol" => protocol.to_string())
.increment(counts.cache_hits);
counter!("water_fill.interpolated_swaps", "protocol" => protocol.to_string())
.increment(counts.interpolated);
counter!("water_fill.refused_without_calling", "protocol" => protocol.to_string())
.increment(counts.refused_without_calling);
}
if !enabled!(Level::DEBUG) {
return;
}
let mut by_stage: FxHashMap<StageLabel, ComponentSwaps> = FxHashMap::default();
let mut components: FxHashSet<&ComponentId> = FxHashSet::default();
let mut totals = ComponentSwaps::default();
for ((component_id, stage), counts) in swaps {
by_stage
.entry(stage)
.or_default()
.add(counts);
components.insert(component_id);
totals.add(counts);
}
let mut stages: Vec<(StageLabel, ComponentSwaps)> = by_stage.into_iter().collect();
stages.sort_unstable_by_key(|(_, counts)| Reverse(counts.call_time));
let per_stage = stages
.iter()
.map(|(stage, counts)| {
format!(
"{}: {} calls in {:.1}ms, {} answered without calling",
stage,
counts.calls,
counts.call_time.as_secs_f64() * 1000.0,
counts.cache_hits + counts.interpolated + counts.refused_without_calling,
)
})
.collect::<Vec<_>>()
.join(" | ");
debug!(solve_time_ms, "water-fill simulation by stage: {per_stage}");
let per_protocol = costliest_first
.iter()
.map(|(protocol, counts)| {
format!(
"{protocol}: {} calls ({} failed) in {:.1}ms, {} cache hits, {} interpolated, \
{} refused without calling",
counts.calls,
counts.failed,
counts.call_time.as_secs_f64() * 1000.0,
counts.cache_hits,
counts.interpolated,
counts.refused_without_calling,
)
})
.collect::<Vec<_>>()
.join(" | ");
debug!(
solve_time_ms,
components = components.len(),
get_amount_out_calls = totals.calls,
failed_calls = totals.failed,
cache_hits = totals.cache_hits,
interpolated = totals.interpolated,
refused_without_calling = totals.refused_without_calling,
call_time_ms = totals.call_time.as_secs_f64() * 1000.0,
"water-fill simulation cost: {per_protocol}",
);
}
#[cfg(not(feature = "swap-metrics"))]
pub(crate) fn start_solve() {}
#[cfg(not(feature = "swap-metrics"))]
pub(crate) fn record_cache_hit(_component_id: &ComponentId, _stage: StageLabel) {}
#[cfg(not(feature = "swap-metrics"))]
pub(crate) fn record_interpolation(_component_id: &ComponentId, _stage: StageLabel) {}
#[cfg(not(feature = "swap-metrics"))]
pub(crate) fn record_refusal_without_calling(_component_id: &ComponentId, _stage: StageLabel) {}
#[cfg(not(feature = "swap-metrics"))]
pub(crate) fn report(_market: &MarketState, _solve_time_ms: impl FnOnce() -> u64) {}
pub(crate) trait MeteredProtocolSim {
fn get_amount_out_metered(
&self,
component_id: &ComponentId,
stage: StageLabel,
amount_in: BigUint,
token_in: &Token,
token_out: &Token,
) -> Result<GetAmountOutResult, SimulationError>;
}
impl<T: ProtocolSim + ?Sized> MeteredProtocolSim for T {
fn get_amount_out_metered(
&self,
component_id: &ComponentId,
stage: StageLabel,
amount_in: BigUint,
token_in: &Token,
token_out: &Token,
) -> Result<GetAmountOutResult, SimulationError> {
#[cfg(not(feature = "swap-metrics"))]
{
let _ = (component_id, stage);
self.get_amount_out_guarded(amount_in, token_in, token_out)
}
#[cfg(feature = "swap-metrics")]
{
let started = Instant::now();
let outcome = self.get_amount_out_guarded(amount_in, token_in, token_out);
record_call(component_id, stage, started.elapsed(), outcome.is_err());
outcome
}
}
}
#[cfg(all(test, feature = "swap-metrics"))]
mod tests {
use std::time::Duration;
use super::*;
fn component(id: &str) -> ComponentId {
ComponentId::from(id)
}
fn counts_for(component_id: &ComponentId, stage: StageLabel) -> ComponentSwaps {
SOLVE_SWAPS.with_borrow(|swaps| {
swaps
.get(&(component_id.clone(), stage))
.copied()
.unwrap_or_default()
})
}
#[test]
fn test_records_each_stage_of_a_component_separately() {
start_solve();
let pool = component("pool-a");
record_call(&pool, "ranking", Duration::from_millis(3), false);
record_call(&pool, "ranking", Duration::from_millis(2), true);
record_cache_hit(&pool, "chunking");
let ranking = counts_for(&pool, "ranking");
assert_eq!(ranking.calls, 2);
assert_eq!(ranking.failed, 1);
assert_eq!(ranking.call_time, Duration::from_millis(5));
assert_eq!(ranking.cache_hits, 0);
assert_eq!(counts_for(&pool, "chunking").cache_hits, 1);
}
#[test]
fn test_counts_answers_that_never_reached_the_pool() {
start_solve();
let pool = component("pool-b");
record_cache_hit(&pool, "ranking");
record_interpolation(&pool, "ranking");
record_refusal_without_calling(&pool, "ranking");
let counts = counts_for(&pool, "ranking");
assert_eq!(counts.calls, 0);
assert_eq!(counts.cache_hits, 1);
assert_eq!(counts.interpolated, 1);
assert_eq!(counts.refused_without_calling, 1);
}
#[test]
fn test_start_solve_discards_the_previous_solve() {
start_solve();
let pool = component("pool-c");
record_call(&pool, "ranking", Duration::from_millis(1), false);
start_solve();
assert_eq!(counts_for(&pool, "ranking").calls, 0);
}
#[test]
fn test_unknown_component_is_reported_rather_than_dropped() {
start_solve();
record_call(&component("gone"), "ranking", Duration::from_millis(1), false);
let market = MarketState::default();
SOLVE_SWAPS.with_borrow(|swaps| {
let protocol = swaps.keys().map(|(component_id, _)| {
market
.get_component(component_id)
.map_or("unknown", |component| component.protocol_system.as_str())
});
assert!(protocol.eq(["unknown"]));
});
report(&market, || 1);
}
}