tycho_simulation/price_level_stream/
stream.rs

1use std::{
2    collections::{HashMap, HashSet},
3    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
4};
5
6use async_stream::stream;
7use num_bigint::BigUint;
8use tokio_stream::{Stream, StreamExt};
9use tycho_common::{models::token::Token, Bytes};
10
11use super::{
12    config::{
13        default_denied_pamms, default_served_pamms, PriceLevelStreamConfig,
14        DEFAULT_AUTO_DETECTED_GAS_COST,
15    },
16    titan::{self, ConnectionSettings, TITAN_PRICE_LEVEL_URL, TITAN_PRICE_LEVEL_URL_ENV},
17    tracker::{FreshnessTracker, Now, TrackerSettings, DEFAULT_STALE_AFTER},
18};
19use crate::protocol::models::Update;
20
21/// Static attribute under which each emitted component carries its pAMM venue address.
22pub const PAMM_ADDRESS_ATTRIBUTE: &str = "pamm_address";
23
24/// The longest [`stale_after`](PriceLevelStreamBuilder::stale_after) the builder accepts;
25/// longer values are capped to it. Quotes target the block being built, so serving a ladder for
26/// longer than this is never intended, and deadlines this far ahead stay representable on the
27/// monotonic clock.
28pub const MAX_STALE_AFTER: Duration = Duration::from_secs(3600);
29
30/// Builds a stream of [`Update`]s from the Titan pAMM price level WebSocket.
31///
32/// A new builder serves no pAMMs: register the known venues via
33/// [`with_known_pamms`](Self::with_known_pamms), individual ones via
34/// [`add_pamm`](Self::add_pamm), or opt into serving unknown streamed venues via
35/// [`auto_detect`](Self::auto_detect); [`with_tokens`](Self::with_tokens) provides the token
36/// metadata pairs are interpreted with.
37///
38/// One component is emitted per (pAMM, token pair), identified by the concatenation
39/// `pamm ++ token0 ++ token1` (tokens sorted ascending), under the protocol system
40/// `fallback:{pamm}` — or `pricelevelstream:{pamm}` after
41/// [`without_fallback_router`](Self::without_fallback_router). The venue address is exposed
42/// through the [`PAMM_ADDRESS_ATTRIBUTE`] static attribute for downstream encoding.
43pub struct PriceLevelStreamBuilder {
44    registry: HashMap<Bytes, PriceLevelStreamConfig>,
45    denied: HashSet<Bytes>,
46    tokens: HashMap<Bytes, Token>,
47    url: Option<String>,
48    auto_detect: bool,
49    auto_detected_gas_cost: Option<BigUint>,
50    connection: ConnectionSettings,
51    /// Whether components are emitted under the `fallback:` family, executed through
52    /// `TychoFallbackRouter`, instead of the direct `pricelevelstream:` family.
53    fallback_router: bool,
54    /// See [`stale_after`](Self::stale_after).
55    stale_after: Duration,
56    /// See [`without_quote_guard`](Self::without_quote_guard).
57    quote_guard: bool,
58}
59
60impl Default for PriceLevelStreamBuilder {
61    fn default() -> Self {
62        Self {
63            registry: HashMap::new(),
64            denied: HashSet::new(),
65            tokens: HashMap::new(),
66            url: None,
67            auto_detect: false,
68            auto_detected_gas_cost: None,
69            connection: ConnectionSettings::default(),
70            fallback_router: true,
71            stale_after: DEFAULT_STALE_AFTER,
72            quote_guard: true,
73        }
74    }
75}
76
77impl PriceLevelStreamBuilder {
78    pub fn new() -> Self {
79        Self::default()
80    }
81
82    /// Enables serving pAMMs that are not registered via
83    /// [`with_known_pamms`](Self::with_known_pamms) or [`add_pamm`](Self::add_pamm)
84    /// (disabled by default).
85    ///
86    /// When enabled, any unknown streamed venue — except denied ones (see
87    /// [`deny_pamm`](Self::deny_pamm)) — is served under its full lowercase hex address
88    /// as the name, with the default gas cost. A venue's protocol system therefore changes from
89    /// the address form (`pricelevelstream:{0xaddress}`) to a name (`pricelevelstream:{name}`)
90    /// once it gets registered — via [`add_pamm`](Self::add_pamm) or a release's
91    /// [`default_served_pamms`] recognizing it; the name-independent identifiers — the component id
92    /// and the [`PAMM_ADDRESS_ATTRIBUTE`] — stay stable across such renames.
93    pub fn auto_detect(mut self, enabled: bool) -> Self {
94        self.auto_detect = enabled;
95        self
96    }
97
98    /// Overrides the per-swap gas cost that auto-detected pAMMs (see
99    /// [`auto_detect`](Self::auto_detect)) are served with. Defaults to the maximum over the
100    /// known venue profiles, as the conservative choice. Registered venues are unaffected —
101    /// their gas cost comes from their [`PriceLevelStreamConfig`].
102    pub fn auto_detected_gas_cost(mut self, gas_cost: BigUint) -> Self {
103        self.auto_detected_gas_cost = Some(gas_cost);
104        self
105    }
106
107    /// Overrides the stream endpoint, e.g. to connect to a closer Titan region than the default
108    /// (see <https://docs.titanbuilder.xyz/propamms/takers>). Without it, the
109    /// `TITAN_PAMM_PRICE_LEVEL_URL` environment variable is used when set, else the built-in
110    /// default.
111    pub fn endpoint(mut self, url: impl Into<String>) -> Self {
112        self.url = Some(url.into());
113        self
114    }
115
116    /// Overrides how long a single connection attempt may take before it is aborted and retried
117    /// (default: 10s), so a hung TCP/TLS handshake cannot block the stream forever.
118    pub fn connect_timeout(mut self, timeout: Duration) -> Self {
119        self.connection.connect_timeout = timeout;
120        self
121    }
122
123    /// Overrides the longest gap between parsed Titan frames tolerated before the connection is
124    /// treated as dead and re-established (default: 10s). Titan pushes one frame per second and
125    /// sends no keepalives, so a multi-second silence means a stalled or half-open connection.
126    /// Control frames and unparsable text do not reset this timeout, and neither does the time
127    /// the consumer spends between polls: the gap is measured while the stream waits on the
128    /// socket.
129    pub fn read_idle_timeout(mut self, timeout: Duration) -> Self {
130        self.connection.read_idle_timeout = timeout;
131        self
132    }
133
134    /// Overrides the cap on the exponential reconnect backoff of `2^attempt` seconds
135    /// (default: 32s).
136    pub fn max_backoff(mut self, max_backoff: Duration) -> Self {
137        self.connection.max_backoff = max_backoff;
138        self
139    }
140
141    /// Registers a pAMM to be served under the given configuration, overriding any default,
142    /// denied, or auto-detected one for the same address.
143    ///
144    /// Between [`add_pamm`](Self::add_pamm) and [`deny_pamm`](Self::deny_pamm) for the same
145    /// address, the later call wins; the defaults applied by
146    /// [`with_known_pamms`](Self::with_known_pamms) never override either, in any call order.
147    pub fn add_pamm(mut self, config: PriceLevelStreamConfig) -> Self {
148        self.denied.remove(&config.address);
149        self.registry
150            .insert(config.address.clone(), config);
151        self
152    }
153
154    /// Excludes a venue from being served: drops its current registration (default or explicit)
155    /// and blocks auto-detecting it.
156    ///
157    /// Between [`add_pamm`](Self::add_pamm) and [`deny_pamm`](Self::deny_pamm) for the same
158    /// address, the later call wins; the defaults applied by
159    /// [`with_known_pamms`](Self::with_known_pamms) never override either, in any call order —
160    /// so denying a venue from the default set works whether the denial comes before or after
161    /// [`with_known_pamms`](Self::with_known_pamms).
162    pub fn deny_pamm(mut self, address: Bytes) -> Self {
163        self.registry.remove(&address);
164        self.denied.insert(address);
165        self
166    }
167
168    /// Applies what is known about the streamed venues: registers the known-good ones
169    /// ([`default_served_pamms`]) to be served and denies the known-bad ones
170    /// ([`default_denied_pamms`]) — venues that stream quotes but whose swaps are not executable.
171    ///
172    /// These defaults never override an explicit [`add_pamm`](Self::add_pamm) or
173    /// [`deny_pamm`](Self::deny_pamm) for the same address, regardless of call order.
174    pub fn with_known_pamms(mut self) -> Self {
175        for config in default_served_pamms() {
176            if self.denied.contains(&config.address) {
177                continue;
178            }
179            self.registry
180                .entry(config.address.clone())
181                .or_insert(config);
182        }
183        for address in default_denied_pamms() {
184            if self.registry.contains_key(&address) {
185                continue;
186            }
187            self.denied.insert(address);
188        }
189        self
190    }
191
192    /// Provides the token metadata used to build components and interpret amounts. Pairs whose
193    /// tokens are missing here are skipped.
194    pub fn with_tokens(mut self, tokens: HashMap<Bytes, Token>) -> Self {
195        self.tokens = tokens;
196        self
197    }
198
199    /// Keeps every venue on the direct `pricelevelstream:{name}` path, so swaps execute on the
200    /// venues themselves and a stale maker quote reverts the route.
201    ///
202    /// By default components are emitted under `fallback:{name}`, so tycho-execution routes
203    /// their swaps through `TychoFallbackRouter`. Opt out when the direct call is what you want
204    /// to measure or execute.
205    pub fn without_fallback_router(mut self) -> Self {
206        self.fallback_router = false;
207        self
208    }
209
210    /// Overrides how long a component stays served after the last accepted frame that carried
211    /// it (default: 24s, two slots). A component no accepted frame has carried for this long
212    /// turns stale and is emitted in `removed_pairs`; the next accepted frame carrying it
213    /// re-adds it in `new_pairs`. Frames whose `timestamp` is this old or older are rejected.
214    /// Independent of this setting, a state refuses to quote once its frame is one slot old
215    /// (see [`without_quote_guard`](Self::without_quote_guard)).
216    ///
217    /// Values above [`MAX_STALE_AFTER`] are capped to it. A value shorter than the age frames
218    /// arrive with rejects every frame as `too_old`, which the
219    /// `price_level_stream_frames_rejected_total` counter and a WARN log show.
220    pub fn stale_after(mut self, duration: Duration) -> Self {
221        self.stale_after = duration.min(MAX_STALE_AFTER);
222        self
223    }
224
225    /// Emits states that never refuse to quote.
226    ///
227    /// By default every emitted state refuses `spot_price`, `get_amount_out` and `get_limits`
228    /// once its frame is one slot ([`QUOTE_TTL`](super::state::QUOTE_TTL)) old: Titan quotes
229    /// the block being built, and a venue rejects a fill against an older ladder as stale, so
230    /// such a quote is not executable. Opt out only for a consumer that quotes a state more
231    /// than one slot after it arrived by design — a batch simulator, a validation harness — and
232    /// that accepts a quote the venue may no longer fill. The component still turns stale and
233    /// is removed after [`stale_after`](Self::stale_after), which then becomes the only bound
234    /// on how old a quoted ladder can be. Never disable it on a live router.
235    pub fn without_quote_guard(mut self) -> Self {
236        self.quote_guard = false;
237        self
238    }
239
240    /// Consumes the builder and opens the stream.
241    ///
242    /// Components are emitted under `fallback:{name}`, so tycho-execution routes their swaps
243    /// through `TychoFallbackRouter`, which retries a reverted pAMM swap — a stale maker quote
244    /// reverts in any simulation against a mined block — on the fallback pool the solver names.
245    /// [`without_fallback_router`](Self::without_fallback_router) keeps them on the direct
246    /// `pricelevelstream:` path.
247    ///
248    /// The connection is established lazily on first poll and maintained (with reconnects) for as
249    /// long as the stream is polled; it never terminates on its own, and dropping the stream
250    /// closes the connection.
251    ///
252    /// Every accepted frame yields an update with the states of the served pairs it carries,
253    /// with `new_pairs` for pairs not currently served. The update does not mention pairs the
254    /// frame does not carry, so consumers keep their previous state. A component no accepted
255    /// frame has carried for [`stale_after`](Self::stale_after) turns stale and is emitted in
256    /// `removed_pairs`, together with every other component turning stale at that instant, and
257    /// re-added by the next accepted frame carrying it. Frames that are too old, from the
258    /// future, out of order, or whose block regresses or jumps more than one block per elapsed
259    /// slot plus 2 are rejected without effect. Frames that contain no served pAMM produce no
260    /// update. Pairs whose tokens are missing from the provided token metadata are skipped.
261    ///
262    /// Every update is stamped with the block its frame targets, a removal with the newest
263    /// accepted block. Block numbers never decrease while something is served. Once every
264    /// component has turned stale, the next accepted frame is judged as a first frame and may
265    /// carry a lower block than the removal did: that is how the stream recovers from a frame
266    /// with an implausible block, so consumers must not rely on the block number to order
267    /// updates across such a gap. See the [module documentation](super) for the full contract.
268    pub fn build(self) -> impl Stream<Item = Update> + Send {
269        let Self {
270            registry,
271            denied,
272            tokens,
273            url,
274            auto_detect,
275            auto_detected_gas_cost,
276            connection,
277            fallback_router,
278            stale_after,
279            quote_guard,
280        } = self;
281        if registry.is_empty() && !auto_detect {
282            tracing::warn!(
283                "No pAMMs registered and auto-detection is off; the stream will never produce \
284                 an update"
285            );
286        }
287        if tokens.is_empty() {
288            tracing::warn!(
289                "No token metadata provided; every streamed pair will be skipped and the stream \
290                 will never produce an update"
291            );
292        }
293        let url = url.unwrap_or_else(|| {
294            std::env::var(TITAN_PRICE_LEVEL_URL_ENV)
295                .unwrap_or_else(|_| TITAN_PRICE_LEVEL_URL.to_string())
296        });
297        let auto_detected_gas_cost =
298            auto_detected_gas_cost.unwrap_or_else(|| BigUint::from(DEFAULT_AUTO_DETECTED_GAS_COST));
299        let mut tracker = FreshnessTracker::new(TrackerSettings {
300            registry,
301            denied,
302            tokens,
303            auto_detect,
304            auto_detected_gas_cost,
305            stale_after,
306            via_fallback_router: fallback_router,
307            quote_guard,
308        });
309
310        stream! {
311            let frames = titan::messages(url, connection);
312            tokio::pin!(frames);
313            loop {
314                let deadline = tracker.stale_deadline();
315                let sleep_until_deadline = tokio::time::sleep_until(
316                    deadline.map_or_else(tokio::time::Instant::now, tokio::time::Instant::from_std),
317                );
318                let update = tokio::select! {
319                    Some(frame) = frames.next() => tracker.on_frame(frame, now()),
320                    () = sleep_until_deadline, if deadline.is_some() => {
321                        tracker.on_stale_deadline(timer_now())
322                    }
323                };
324                if let Some(update) = update {
325                    yield update;
326                }
327            }
328        }
329    }
330}
331
332/// The clocks a frame is judged against: the wall clock for Titan's `timestamp`, and the
333/// timer's clock for every deadline. This is the only place the stream reads a clock; the
334/// tracker only ever sees the values it is handed.
335fn now() -> Now {
336    Now { wall_nanos: wall_clock_nanos(), monotonic: timer_now() }
337}
338
339/// Nanoseconds since the Unix epoch. A system clock unrepresentable as unix nanoseconds yields
340/// 0, which rejects every frame as `in_future`.
341fn wall_clock_nanos() -> u64 {
342    let since_epoch = SystemTime::now()
343        .duration_since(UNIX_EPOCH)
344        .ok()
345        .and_then(|since_epoch| u64::try_from(since_epoch.as_nanos()).ok());
346    match since_epoch {
347        Some(wall_nanos) => wall_nanos,
348        None => {
349            tracing::error!(
350                "System clock unrepresentable as unix nanoseconds; every price level frame will \
351                 be rejected as in_future"
352            );
353            0
354        }
355    }
356}
357
358/// The clock the deadline timer runs on. Deadlines are set and swept with it so that a sweep
359/// fired by the timer finds the component due, also under `tokio::time::pause()`, where the
360/// runtime's virtual time and `std::time::Instant` diverge.
361fn timer_now() -> Instant {
362    tokio::time::Instant::now().into_std()
363}
364
365#[cfg(test)]
366mod tests {
367    use std::{
368        pin::Pin,
369        str::FromStr,
370        sync::{
371            atomic::{AtomicBool, Ordering},
372            Arc,
373        },
374        time::Duration,
375    };
376
377    use futures::{future::BoxFuture, SinkExt};
378    use num_bigint::BigUint;
379    use rstest::rstest;
380    use tokio_tungstenite::tungstenite::Message;
381
382    use super::{
383        super::{
384            config::{default_denied_pamms, PriceLevelStreamConfig},
385            state::PriceLevelStreamState,
386            telemetry::{
387                recorded::{counter_value, record_async},
388                RECONNECTS,
389            },
390            test_support::{
391                fermiswap, frame_text, frame_then_repeat, tokens, wall_nanos_now, FakeConnection,
392                FakeTitan, PAMM,
393            },
394        },
395        *,
396    };
397
398    #[test]
399    fn explicit_add_and_deny_are_last_wins() {
400        let address = Bytes::from_str(PAMM).unwrap();
401        let custom =
402            || PriceLevelStreamConfig::new("custom", Bytes::from_str(PAMM).unwrap(), 1u64.into());
403
404        let builder = PriceLevelStreamBuilder::new()
405            .add_pamm(custom())
406            .deny_pamm(address.clone());
407        assert!(!builder.registry.contains_key(&address));
408        assert!(builder.denied.contains(&address));
409
410        let builder = PriceLevelStreamBuilder::new()
411            .deny_pamm(address.clone())
412            .add_pamm(custom());
413        assert_eq!(builder.registry[&address].protocol, "custom");
414        assert!(builder.denied.is_empty());
415    }
416
417    #[test]
418    fn quote_guard_is_on_unless_opted_out() {
419        assert!(PriceLevelStreamBuilder::new().quote_guard);
420        assert!(
421            !PriceLevelStreamBuilder::new()
422                .without_quote_guard()
423                .quote_guard
424        );
425    }
426
427    /// A builder with short timings: `stale_after` 2 s, `read_idle_timeout` 100 ms,
428    /// `max_backoff` 20 ms.
429    fn fast_builder(fake: &FakeTitan) -> PriceLevelStreamBuilder {
430        PriceLevelStreamBuilder::new()
431            .endpoint(fake.url())
432            .add_pamm(fermiswap())
433            .with_tokens(tokens())
434            .stale_after(STALE_AFTER)
435            .connect_timeout(Duration::from_secs(1))
436            .read_idle_timeout(Duration::from_millis(100))
437            .max_backoff(Duration::from_millis(20))
438    }
439
440    /// The `stale_after` of [`fast_builder`]: long enough that a loaded runner does not reject
441    /// the first frame as too old, short enough that a removal arrives within [`WAIT`].
442    const STALE_AFTER: Duration = Duration::from_secs(2);
443
444    /// How long a test waits for one update.
445    const WAIT: Duration = Duration::from_secs(5);
446
447    async fn next_within(
448        stream: &mut Pin<&mut impl Stream<Item = Update>>,
449        limit: Duration,
450    ) -> Option<Update> {
451        tokio::time::timeout(limit, stream.next())
452            .await
453            .ok()
454            .flatten()
455    }
456
457    /// Waits for the first update and checks that it adds the one served component.
458    async fn expect_first_update(stream: &mut Pin<&mut impl Stream<Item = Update>>) -> Update {
459        let first = next_within(stream, WAIT)
460            .await
461            .expect("first update");
462        assert_eq!(first.new_pairs.len(), 1);
463        assert!(first.removed_pairs.is_empty());
464        first
465    }
466
467    /// Waits up to [`WAIT`] for an update that removes components, skipping the updates that
468    /// only refresh served ones; an update that re-adds a component before the removal fails.
469    async fn expect_removal(stream: &mut Pin<&mut impl Stream<Item = Update>>) -> Update {
470        tokio::time::timeout(WAIT, async {
471            loop {
472                let update = stream
473                    .next()
474                    .await
475                    .expect("stream ended");
476                if !update.removed_pairs.is_empty() {
477                    return update;
478                }
479                assert!(update.new_pairs.is_empty(), "component re-added before the removal");
480            }
481        })
482        .await
483        .expect("removal")
484    }
485
486    fn assert_removal_only(update: &Update, expected_removed: usize) {
487        assert!(update.states.is_empty());
488        assert!(update.new_pairs.is_empty());
489        assert!(update.sync_states.is_empty());
490        assert!(update.is_partial);
491        assert_eq!(update.removed_pairs.len(), expected_removed);
492    }
493
494    fn fresh_frame() -> Message {
495        Message::Text(frame_text(100, wall_nanos_now()).into())
496    }
497
498    /// Sends one fresh frame on the first connection only; later connections stay silent.
499    async fn first_connection_sends_one_frame(index: usize, mut socket: FakeConnection) {
500        if index == 0 {
501            let _ = socket.send(fresh_frame()).await;
502        }
503        std::future::pending::<()>().await;
504    }
505
506    #[tokio::test]
507    async fn silence_past_stale_after_removes_every_served_component() {
508        let fake = FakeTitan::spawn(first_connection_sends_one_frame).await;
509        let stream = fast_builder(&fake).build();
510        tokio::pin!(stream);
511
512        expect_first_update(&mut stream).await;
513        // The stream reconnects on idle timeout, but the later connections send no frame, so no
514        // component deadline moves.
515        let removal = next_within(&mut stream, WAIT)
516            .await
517            .expect("removal");
518
519        assert_removal_only(&removal, 1);
520        assert_eq!(removal.block_number_or_timestamp, 100);
521        assert!(fake.connections.load(Ordering::SeqCst) >= 2, "no reconnect on idle timeout");
522    }
523
524    #[tokio::test]
525    async fn repeated_immediate_closes_remove_within_stale_after() {
526        let fake = FakeTitan::spawn(|index, mut socket| async move {
527            if index == 0 {
528                let _ = socket.send(fresh_frame()).await;
529                tokio::time::sleep(Duration::from_millis(20)).await;
530            }
531            let _ = socket.close(None).await;
532        })
533        .await;
534        let stream = fast_builder(&fake).build();
535        tokio::pin!(stream);
536
537        expect_first_update(&mut stream).await;
538        let removal = next_within(&mut stream, WAIT)
539            .await
540            .expect("removal");
541
542        assert_removal_only(&removal, 1);
543    }
544
545    #[test]
546    fn refused_reconnects_remove_within_stale_after() {
547        let (removal, snapshot) = record_async(async {
548            let mut fake = FakeTitan::spawn(|_, mut socket| async move {
549                let _ = socket.send(fresh_frame()).await;
550                tokio::time::sleep(Duration::from_millis(20)).await;
551                let _ = socket.close(None).await;
552            })
553            .await;
554            let stream = fast_builder(&fake).build();
555            tokio::pin!(stream);
556
557            expect_first_update(&mut stream).await;
558            // From here every connect is refused at TCP level. The fake closes 20 ms in and the
559            // backoff is capped at 20 ms, so one reconnect can be accepted before `shutdown`
560            // drops the listener; its fresh frame refreshes the component without re-adding it.
561            fake.shutdown();
562            expect_removal(&mut stream).await
563        });
564
565        assert_removal_only(&removal, 1);
566        assert!(
567            counter_value(&snapshot, RECONNECTS, &[("reason", "connect_failed")]) >= 1,
568            "no connect was refused"
569        );
570    }
571
572    #[tokio::test]
573    async fn replayed_frames_remove_within_stale_after_and_never_re_add() {
574        // One frame, stamped once, replayed every 50 ms forever.
575        let replay = fresh_frame();
576        let fake =
577            FakeTitan::spawn(frame_then_repeat(replay.clone(), replay, Duration::from_millis(50)))
578                .await;
579        let stream = fast_builder(&fake)
580            .read_idle_timeout(Duration::from_secs(5))
581            .build();
582        tokio::pin!(stream);
583
584        expect_first_update(&mut stream).await;
585        // Replays are accepted while fresh but cannot extend the deadline; the deadline fires
586        // even though a frame is ready on every poll.
587        let removal = expect_removal(&mut stream).await;
588
589        assert_removal_only(&removal, 1);
590        // Every later replay is too old to be accepted: nothing comes back.
591        assert!(next_within(&mut stream, Duration::from_millis(500))
592            .await
593            .is_none());
594    }
595
596    #[tokio::test]
597    async fn fresh_frame_after_removal_re_adds_the_component() {
598        let fake = FakeTitan::spawn(|_, mut socket| async move {
599            let _ = socket.send(fresh_frame()).await;
600            tokio::time::sleep(STALE_AFTER + Duration::from_millis(500)).await;
601            let _ = socket
602                .send(Message::Text(frame_text(101, wall_nanos_now()).into()))
603                .await;
604            std::future::pending::<()>().await;
605        })
606        .await;
607        let stream = fast_builder(&fake)
608            .read_idle_timeout(Duration::from_secs(5))
609            .build();
610        tokio::pin!(stream);
611
612        expect_first_update(&mut stream).await;
613        let removal = next_within(&mut stream, WAIT)
614            .await
615            .expect("removal");
616        let re_added = next_within(&mut stream, WAIT)
617            .await
618            .expect("re-add");
619
620        assert_removal_only(&removal, 1);
621        assert_eq!(re_added.new_pairs.len(), 1);
622        assert!(re_added.removed_pairs.is_empty());
623        assert_eq!(re_added.block_number_or_timestamp, 101);
624    }
625
626    #[tokio::test]
627    async fn frames_are_forwarded_without_waiting_on_timers() {
628        let fake = FakeTitan::spawn(fresh_frame_every(Duration::from_millis(50))).await;
629        let stream = fast_builder(&fake)
630            .stale_after(Duration::from_secs(24))
631            .build();
632        tokio::pin!(stream);
633
634        expect_first_update(&mut stream).await;
635        for _ in 0..5 {
636            let update = next_within(&mut stream, Duration::from_secs(1))
637                .await
638                .expect("steady-state frame");
639            assert!(update.removed_pairs.is_empty());
640        }
641    }
642
643    #[tokio::test]
644    async fn no_connection_before_first_poll() {
645        let fake = FakeTitan::spawn(first_connection_sends_one_frame).await;
646        let stream = fast_builder(&fake).build();
647        tokio::pin!(stream);
648
649        tokio::time::sleep(Duration::from_millis(150)).await;
650
651        assert_eq!(fake.connections.load(Ordering::SeqCst), 0, "connected before first poll");
652    }
653
654    #[tokio::test]
655    async fn drop_closes_the_socket() {
656        let server_saw_close = Arc::new(AtomicBool::new(false));
657        let fake = {
658            let server_saw_close = server_saw_close.clone();
659            FakeTitan::spawn(move |_, mut socket| {
660                let server_saw_close = server_saw_close.clone();
661                async move {
662                    let _ = socket.send(fresh_frame()).await;
663                    // Read until the client goes away.
664                    while let Some(Ok(message)) = socket.next().await {
665                        if matches!(message, Message::Close(_)) {
666                            break;
667                        }
668                    }
669                    server_saw_close.store(true, Ordering::SeqCst);
670                }
671            })
672            .await
673        };
674        // `Box::pin` rather than `tokio::pin!`, so that `drop` below drops the stream itself
675        // rather than a `Pin<&mut _>` pointing at a value that outlives the assertion.
676        let mut stream = Box::pin(
677            fast_builder(&fake)
678                .read_idle_timeout(Duration::from_secs(5))
679                .build(),
680        );
681        expect_first_update(&mut stream.as_mut()).await;
682
683        drop(stream);
684        tokio::time::sleep(Duration::from_millis(200)).await;
685
686        assert!(server_saw_close.load(Ordering::SeqCst), "socket not closed on drop");
687        assert_eq!(fake.connections.load(Ordering::SeqCst), 1, "reconnected after drop");
688    }
689
690    /// `Instant + Duration` panics on overflow, so an absurd `stale_after` must be capped before
691    /// the first accepted frame computes a deadline from it.
692    #[tokio::test]
693    async fn stale_after_above_the_cap_still_serves() {
694        assert_eq!(
695            PriceLevelStreamBuilder::new()
696                .stale_after(Duration::MAX)
697                .stale_after,
698            MAX_STALE_AFTER
699        );
700        let fake = FakeTitan::spawn(first_connection_sends_one_frame).await;
701        let stream = fast_builder(&fake)
702            .stale_after(Duration::MAX)
703            .build();
704        tokio::pin!(stream);
705
706        expect_first_update(&mut stream).await;
707    }
708
709    #[tokio::test]
710    async fn without_quote_guard_emits_states_that_never_expire() {
711        let fake = FakeTitan::spawn(first_connection_sends_one_frame).await;
712        let stream = fast_builder(&fake)
713            .without_quote_guard()
714            .build();
715        tokio::pin!(stream);
716
717        let first = expect_first_update(&mut stream).await;
718
719        let state = first
720            .states
721            .values()
722            .next()
723            .expect("one state")
724            .as_any()
725            .downcast_ref::<PriceLevelStreamState>()
726            .expect("price level state");
727        assert!(state.quotable_until().is_none());
728    }
729
730    /// A [`FakeTitan`] handler that sends a freshly stamped frame every `interval` until the
731    /// socket closes.
732    fn fresh_frame_every(
733        interval: Duration,
734    ) -> impl Fn(usize, FakeConnection) -> BoxFuture<'static, ()> + Send + Sync + 'static {
735        move |_, mut socket| {
736            Box::pin(async move {
737                while socket.send(fresh_frame()).await.is_ok() {
738                    tokio::time::sleep(interval).await;
739                }
740            })
741        }
742    }
743
744    /// The fallback router path is the default; `without_fallback_router` is the way off it.
745    #[test]
746    fn fallback_router_is_on_unless_opted_out() {
747        assert!(PriceLevelStreamBuilder::new().fallback_router);
748        assert!(
749            !PriceLevelStreamBuilder::new()
750                .without_fallback_router()
751                .fallback_router
752        );
753    }
754
755    #[test]
756    fn defaults_never_override_explicit_calls() {
757        // Denying a venue from the default set works in either call order.
758        let fermiswap_router = Bytes::from_str(PAMM).unwrap();
759        for builder in [
760            PriceLevelStreamBuilder::new()
761                .deny_pamm(fermiswap_router.clone())
762                .with_known_pamms(),
763            PriceLevelStreamBuilder::new()
764                .with_known_pamms()
765                .deny_pamm(fermiswap_router.clone()),
766        ] {
767            assert!(!builder
768                .registry
769                .contains_key(&fermiswap_router));
770            assert!(builder
771                .denied
772                .contains(&fermiswap_router));
773            // The other defaults are unaffected.
774            assert!(!builder.registry.is_empty());
775        }
776
777        // Registering a venue from the default deny set works in either call order. Any one of
778        // them exercises that; the set is empty while every streamed venue is executable.
779        let Some(denied_venue) = default_denied_pamms().pop() else { return };
780        let custom = || PriceLevelStreamConfig::new("custom", denied_venue.clone(), 1u64.into());
781        for builder in [
782            PriceLevelStreamBuilder::new()
783                .add_pamm(custom())
784                .with_known_pamms(),
785            PriceLevelStreamBuilder::new()
786                .with_known_pamms()
787                .add_pamm(custom()),
788        ] {
789            assert_eq!(builder.registry[&denied_venue].protocol, "custom");
790            assert!(!builder.denied.contains(&denied_venue));
791        }
792    }
793
794    #[test]
795    fn with_known_pamms_registers_known_venues() {
796        // PAMM is the FermiSwap router, one of the default venues.
797        let fermiswap_router = Bytes::from_str(PAMM).unwrap();
798
799        let builder = PriceLevelStreamBuilder::new();
800        assert!(builder.registry.is_empty());
801        assert!(builder.denied.is_empty());
802
803        let builder = builder.with_known_pamms();
804        assert_eq!(builder.registry[&fermiswap_router].protocol, "fermiswap");
805        // The known-bad venues get denied alongside, and never overlap the served defaults.
806        assert_eq!(
807            builder.denied,
808            default_denied_pamms()
809                .into_iter()
810                .collect()
811        );
812        assert!(builder.denied.is_disjoint(
813            &builder
814                .registry
815                .keys()
816                .cloned()
817                .collect()
818        ));
819
820        // An `add_pamm` entry wins over the default for the same address, in either call order.
821        let custom =
822            || PriceLevelStreamConfig::new("custom", fermiswap_router.clone(), BigUint::from(1u64));
823        for builder in [
824            PriceLevelStreamBuilder::new()
825                .add_pamm(custom())
826                .with_known_pamms(),
827            PriceLevelStreamBuilder::new()
828                .with_known_pamms()
829                .add_pamm(custom()),
830        ] {
831            assert_eq!(builder.registry[&fermiswap_router].protocol, "custom");
832            assert_eq!(builder.registry[&fermiswap_router].gas_cost, BigUint::from(1u64));
833        }
834    }
835
836    /// The families this stream emits are the ones tycho-execution resolves an encoder for. A
837    /// drift between the two makes every route through a pAMM fail to encode.
838    #[test]
839    fn families_match_the_execution_side_prefixes() {
840        use tycho_execution::encoding::evm::{FALLBACK_PREFIX, PRICE_LEVEL_STREAM_PREFIX};
841
842        use super::super::config::{FALLBACK_FAMILY, PRICE_LEVEL_STREAM_FAMILY};
843
844        assert_eq!(format!("{PRICE_LEVEL_STREAM_FAMILY}:"), PRICE_LEVEL_STREAM_PREFIX);
845        assert_eq!(format!("{FALLBACK_FAMILY}:"), FALLBACK_PREFIX);
846    }
847
848    /// Traffic that is not a parsed frame keeps neither the socket nor the component alive: the
849    /// idle timeout reconnects, every reconnect resends the same frame, which refreshes the
850    /// component but cannot move its deadline, and the removal follows.
851    #[rstest]
852    #[case::ping_only(Message::Ping(Vec::new().into()))]
853    #[case::malformed_text(Message::Text("nonsense".into()))]
854    #[tokio::test]
855    async fn non_frame_traffic_removes_within_stale_after(#[case] filler: Message) {
856        let fake =
857            FakeTitan::spawn(frame_then_repeat(fresh_frame(), filler, Duration::from_millis(10)))
858                .await;
859        let stream = fast_builder(&fake).build();
860        tokio::pin!(stream);
861
862        expect_first_update(&mut stream).await;
863        let removal = expect_removal(&mut stream).await;
864
865        assert_removal_only(&removal, 1);
866        assert!(fake.connections.load(Ordering::SeqCst) >= 2, "no reconnect on idle timeout");
867    }
868}
tycho_simulation/price_level_stream/stream.rs

tycho_simulation/price_level_stream/
stream.rs
