1#[cfg(feature = "swap-metrics")]
19use std::{
20 cell::RefCell,
21 cmp::Reverse,
22 time::{Duration, Instant},
23};
24
25#[cfg(feature = "swap-metrics")]
26use metrics::counter;
27use num_bigint::BigUint;
28#[cfg(feature = "swap-metrics")]
29use rustc_hash::{FxHashMap, FxHashSet};
30#[cfg(feature = "swap-metrics")]
31use tracing::{debug, enabled, Level};
32use tycho_simulation::tycho_common::{
33 models::token::Token,
34 simulation::{
35 errors::SimulationError,
36 protocol_sim::{GetAmountOutResult, ProtocolSim},
37 },
38};
39
40use super::sim_guard::GuardedProtocolSim;
41use crate::{feed::market_data::MarketState, types::ComponentId};
42
43pub type StageLabel = &'static str;
48
49#[cfg(feature = "swap-metrics")]
52#[derive(Default, Clone, Copy)]
53struct ComponentSwaps {
54 calls: u64,
56 failed: u64,
59 cache_hits: u64,
61 interpolated: u64,
65 refused_without_calling: u64,
67 call_time: Duration,
69}
70
71#[cfg(feature = "swap-metrics")]
72impl ComponentSwaps {
73 fn add(&mut self, other: &ComponentSwaps) {
74 self.calls += other.calls;
75 self.failed += other.failed;
76 self.cache_hits += other.cache_hits;
77 self.interpolated += other.interpolated;
78 self.refused_without_calling += other.refused_without_calling;
79 self.call_time += other.call_time;
80 }
81}
82
83#[cfg(feature = "swap-metrics")]
84thread_local! {
85 static SOLVE_SWAPS: RefCell<FxHashMap<(ComponentId, StageLabel), ComponentSwaps>> =
95 RefCell::new(FxHashMap::default());
96}
97
98#[cfg(feature = "swap-metrics")]
99fn with_counts(
100 component_id: &ComponentId,
101 stage: StageLabel,
102 edit: impl FnOnce(&mut ComponentSwaps),
103) {
104 SOLVE_SWAPS.with_borrow_mut(|swaps| {
105 edit(
106 swaps
107 .entry((component_id.clone(), stage))
108 .or_default(),
109 );
110 });
111}
112
113#[cfg(feature = "swap-metrics")]
115pub fn start_solve() {
116 SOLVE_SWAPS.with_borrow_mut(FxHashMap::clear);
117}
118
119#[cfg(feature = "swap-metrics")]
121fn record_call(component_id: &ComponentId, stage: StageLabel, call_time: Duration, failed: bool) {
122 with_counts(component_id, stage, |counts| {
123 counts.calls += 1;
124 counts.call_time += call_time;
125 if failed {
126 counts.failed += 1;
127 }
128 });
129}
130
131#[cfg(feature = "swap-metrics")]
133pub fn record_cache_hit(component_id: &ComponentId, stage: StageLabel) {
134 with_counts(component_id, stage, |counts| counts.cache_hits += 1);
135}
136
137#[cfg(feature = "swap-metrics")]
139pub fn record_interpolation(component_id: &ComponentId, stage: StageLabel) {
140 with_counts(component_id, stage, |counts| counts.interpolated += 1);
141}
142
143#[cfg(feature = "swap-metrics")]
145pub fn record_refusal_without_calling(component_id: &ComponentId, stage: StageLabel) {
146 with_counts(component_id, stage, |counts| counts.refused_without_calling += 1);
147}
148
149#[cfg(feature = "swap-metrics")]
156pub fn report(algorithm: &str, market: &MarketState, solve_time_ms: impl FnOnce() -> u64) {
157 let solve_time_ms = solve_time_ms();
158 SOLVE_SWAPS.with_borrow(|swaps| report_swaps(algorithm, swaps, market, solve_time_ms));
159}
160
161#[cfg(feature = "swap-metrics")]
162fn report_swaps(
163 algorithm: &str,
164 swaps: &FxHashMap<(ComponentId, StageLabel), ComponentSwaps>,
165 market: &MarketState,
166 solve_time_ms: u64,
167) {
168 let mut by_protocol: FxHashMap<&str, ComponentSwaps> = FxHashMap::default();
169 for ((component_id, _), counts) in swaps {
170 let protocol = market
171 .get_component(component_id)
172 .map_or("unknown", |component| component.protocol_system.as_str());
173 by_protocol
174 .entry(protocol)
175 .or_default()
176 .add(counts);
177 }
178
179 let mut costliest_first: Vec<(&str, ComponentSwaps)> = by_protocol.into_iter().collect();
180 costliest_first
181 .sort_unstable_by_key(|(protocol, counts)| (Reverse(counts.call_time), *protocol));
182
183 for (protocol, counts) in &costliest_first {
184 counter!(format!("{algorithm}.get_amount_out_calls"), "protocol" => protocol.to_string())
185 .increment(counts.calls);
186 counter!(format!("{algorithm}.failed_calls"), "protocol" => protocol.to_string())
187 .increment(counts.failed);
188 counter!(format!("{algorithm}.cache_hits"), "protocol" => protocol.to_string())
189 .increment(counts.cache_hits);
190 counter!(format!("{algorithm}.interpolated_swaps"), "protocol" => protocol.to_string())
191 .increment(counts.interpolated);
192 counter!(format!("{algorithm}.refused_without_calling"), "protocol" => protocol.to_string())
193 .increment(counts.refused_without_calling);
194 }
195
196 if !enabled!(Level::DEBUG) {
197 return;
198 }
199
200 let mut by_stage: FxHashMap<StageLabel, ComponentSwaps> = FxHashMap::default();
201 let mut components: FxHashSet<&ComponentId> = FxHashSet::default();
202 let mut totals = ComponentSwaps::default();
203 for ((component_id, stage), counts) in swaps {
204 by_stage
205 .entry(stage)
206 .or_default()
207 .add(counts);
208 components.insert(component_id);
209 totals.add(counts);
210 }
211
212 let mut stages: Vec<(StageLabel, ComponentSwaps)> = by_stage.into_iter().collect();
213 stages.sort_unstable_by_key(|(_, counts)| Reverse(counts.call_time));
214 let per_stage = stages
215 .iter()
216 .map(|(stage, counts)| {
217 format!(
218 "{}: {} calls in {:.1}ms, {} answered without calling",
219 stage,
220 counts.calls,
221 counts.call_time.as_secs_f64() * 1000.0,
222 counts.cache_hits + counts.interpolated + counts.refused_without_calling,
223 )
224 })
225 .collect::<Vec<_>>()
226 .join(" | ");
227 debug!(solve_time_ms, "{algorithm} simulation by stage: {per_stage}");
228
229 let per_protocol = costliest_first
230 .iter()
231 .map(|(protocol, counts)| {
232 format!(
233 "{protocol}: {} calls ({} failed) in {:.1}ms, {} cache hits, {} interpolated, \
234 {} refused without calling",
235 counts.calls,
236 counts.failed,
237 counts.call_time.as_secs_f64() * 1000.0,
238 counts.cache_hits,
239 counts.interpolated,
240 counts.refused_without_calling,
241 )
242 })
243 .collect::<Vec<_>>()
244 .join(" | ");
245 debug!(
246 solve_time_ms,
247 components = components.len(),
248 get_amount_out_calls = totals.calls,
249 failed_calls = totals.failed,
250 cache_hits = totals.cache_hits,
251 interpolated = totals.interpolated,
252 refused_without_calling = totals.refused_without_calling,
253 call_time_ms = totals.call_time.as_secs_f64() * 1000.0,
254 "{algorithm} simulation cost: {per_protocol}",
255 );
256}
257
258#[cfg(not(feature = "swap-metrics"))]
261pub fn start_solve() {}
262
263#[cfg(not(feature = "swap-metrics"))]
265pub fn record_cache_hit(_component_id: &ComponentId, _stage: StageLabel) {}
266
267#[cfg(not(feature = "swap-metrics"))]
269pub fn record_interpolation(_component_id: &ComponentId, _stage: StageLabel) {}
270
271#[cfg(not(feature = "swap-metrics"))]
273pub fn record_refusal_without_calling(_component_id: &ComponentId, _stage: StageLabel) {}
274
275#[cfg(not(feature = "swap-metrics"))]
280pub fn report(_algorithm: &str, _market: &MarketState, _solve_time_ms: impl FnOnce() -> u64) {}
281
282pub trait MeteredProtocolSim {
287 fn get_amount_out_metered(
289 &self,
290 component_id: &ComponentId,
291 stage: StageLabel,
292 amount_in: BigUint,
293 token_in: &Token,
294 token_out: &Token,
295 ) -> Result<GetAmountOutResult, SimulationError>;
296}
297
298impl<T: ProtocolSim + ?Sized> MeteredProtocolSim for T {
299 fn get_amount_out_metered(
300 &self,
301 component_id: &ComponentId,
302 stage: StageLabel,
303 amount_in: BigUint,
304 token_in: &Token,
305 token_out: &Token,
306 ) -> Result<GetAmountOutResult, SimulationError> {
307 #[cfg(not(feature = "swap-metrics"))]
308 {
309 let _ = (component_id, stage);
310 self.get_amount_out_guarded(amount_in, token_in, token_out)
311 }
312 #[cfg(feature = "swap-metrics")]
313 {
314 let started = Instant::now();
315 let outcome = self.get_amount_out_guarded(amount_in, token_in, token_out);
316 record_call(component_id, stage, started.elapsed(), outcome.is_err());
317 outcome
318 }
319 }
320}
321
322#[cfg(all(test, feature = "swap-metrics"))]
323mod tests {
324 use std::time::Duration;
325
326 use super::*;
327
328 fn component(id: &str) -> ComponentId {
329 ComponentId::from(id)
330 }
331
332 fn counts_for(component_id: &ComponentId, stage: StageLabel) -> ComponentSwaps {
333 SOLVE_SWAPS.with_borrow(|swaps| {
334 swaps
335 .get(&(component_id.clone(), stage))
336 .copied()
337 .unwrap_or_default()
338 })
339 }
340
341 #[test]
344 fn test_records_each_stage_of_a_component_separately() {
345 start_solve();
346 let pool = component("pool-a");
347
348 record_call(&pool, "ranking", Duration::from_millis(3), false);
349 record_call(&pool, "ranking", Duration::from_millis(2), true);
350 record_cache_hit(&pool, "chunking");
351
352 let ranking = counts_for(&pool, "ranking");
353 assert_eq!(ranking.calls, 2);
354 assert_eq!(ranking.failed, 1);
355 assert_eq!(ranking.call_time, Duration::from_millis(5));
356 assert_eq!(ranking.cache_hits, 0);
357 assert_eq!(counts_for(&pool, "chunking").cache_hits, 1);
358 }
359
360 #[test]
363 fn test_counts_answers_that_never_reached_the_pool() {
364 start_solve();
365 let pool = component("pool-b");
366
367 record_cache_hit(&pool, "ranking");
368 record_interpolation(&pool, "ranking");
369 record_refusal_without_calling(&pool, "ranking");
370
371 let counts = counts_for(&pool, "ranking");
372 assert_eq!(counts.calls, 0);
373 assert_eq!(counts.cache_hits, 1);
374 assert_eq!(counts.interpolated, 1);
375 assert_eq!(counts.refused_without_calling, 1);
376 }
377
378 #[test]
380 fn test_start_solve_discards_the_previous_solve() {
381 start_solve();
382 let pool = component("pool-c");
383 record_call(&pool, "ranking", Duration::from_millis(1), false);
384
385 start_solve();
386
387 assert_eq!(counts_for(&pool, "ranking").calls, 0);
388 }
389
390 #[test]
393 fn test_unknown_component_is_reported_rather_than_dropped() {
394 start_solve();
395 record_call(&component("gone"), "ranking", Duration::from_millis(1), false);
396
397 let market = MarketState::default();
400 SOLVE_SWAPS.with_borrow(|swaps| {
401 let protocol = swaps.keys().map(|(component_id, _)| {
402 market
403 .get_component(component_id)
404 .map_or("unknown", |component| component.protocol_system.as_str())
405 });
406 assert!(protocol.eq(["unknown"]));
407 });
408 report("test", &market, || 1);
409 }
410}