Skip to main content

nautilus_data/defi/
engine.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! DeFi-specific data engine functionality.
17//!
18//! This module provides DeFi processing methods for the `DataEngine`.
19//! All code in this module requires the `defi` feature flag.
20
21use std::{rc::Rc, sync::Arc};
22
23use nautilus_common::{
24    defi,
25    messages::defi::{
26        DefiRequestCommand, DefiSubscribeCommand, DefiUnsubscribeCommand, PoolSnapshotResponse,
27        RequestPoolSnapshot,
28    },
29    msgbus::{self, TypedHandler},
30};
31use nautilus_core::UUID4;
32use nautilus_model::{
33    defi::{
34        Blockchain, DefiData, PoolProfiler,
35        data::{DexPoolData, block::BlockPosition},
36        pool_analysis::PoolSnapshot,
37    },
38    identifiers::{ClientId, InstrumentId},
39    instruments::{CurrencyPair, InstrumentAny},
40};
41
42use crate::engine::{
43    DataEngine,
44    pool::{
45        PoolCollectHandler, PoolFlashHandler, PoolLiquidityHandler, PoolSwapHandler, PoolUpdater,
46    },
47};
48
49/// Extracts the block position tuple from a `DexPoolData` event.
50fn get_event_block_position(event: &DexPoolData) -> (u64, u32, u32) {
51    match event {
52        DexPoolData::Swap(s) => (s.block, s.transaction_index, s.log_index),
53        DexPoolData::LiquidityUpdate(u) => (u.block, u.transaction_index, u.log_index),
54        DexPoolData::FeeCollect(c) => (c.block, c.transaction_index, c.log_index),
55        DexPoolData::FeeProtocolUpdate(u) => (u.block, u.transaction_index, u.log_index),
56        DexPoolData::FeeProtocolCollect(c) => (c.block, c.transaction_index, c.log_index),
57        DexPoolData::Flash(f) => (f.block, f.transaction_index, f.log_index),
58    }
59}
60
61/// Converts buffered `DefiData` events to `DexPoolData` and sorts by block position.
62fn convert_and_sort_buffered_events(buffered_events: Vec<DefiData>) -> Vec<DexPoolData> {
63    let mut events: Vec<DexPoolData> = buffered_events
64        .into_iter()
65        .filter_map(|event| match event {
66            DefiData::PoolSwap(swap) => Some(DexPoolData::Swap(swap)),
67            DefiData::PoolLiquidityUpdate(update) => Some(DexPoolData::LiquidityUpdate(update)),
68            DefiData::PoolFeeCollect(collect) => Some(DexPoolData::FeeCollect(collect)),
69            DefiData::PoolFeeProtocolUpdate(update) => Some(DexPoolData::FeeProtocolUpdate(update)),
70            DefiData::PoolFeeProtocolCollect(collect) => {
71                Some(DexPoolData::FeeProtocolCollect(collect))
72            }
73            DefiData::PoolFlash(flash) => Some(DexPoolData::Flash(flash)),
74            _ => None,
75        })
76        .collect();
77
78    events.sort_by(|a, b| {
79        let pos_a = get_event_block_position(a);
80        let pos_b = get_event_block_position(b);
81        pos_a.cmp(&pos_b)
82    });
83
84    events
85}
86
87impl DataEngine {
88    /// Returns all blockchains for which blocks subscriptions exist.
89    #[must_use]
90    pub fn subscribed_blocks(&self) -> Vec<Blockchain> {
91        self.collect_subscriptions(|client| &client.subscriptions_blocks)
92    }
93
94    /// Returns all instrument IDs for which pool subscriptions exist.
95    #[must_use]
96    pub fn subscribed_pools(&self) -> Vec<InstrumentId> {
97        self.collect_subscriptions(|client| &client.subscriptions_pools)
98    }
99
100    /// Returns all instrument IDs for which swap subscriptions exist.
101    #[must_use]
102    pub fn subscribed_pool_swaps(&self) -> Vec<InstrumentId> {
103        self.collect_subscriptions(|client| &client.subscriptions_pool_swaps)
104    }
105
106    /// Returns all instrument IDs for which liquidity update subscriptions exist.
107    #[must_use]
108    pub fn subscribed_pool_liquidity_updates(&self) -> Vec<InstrumentId> {
109        self.collect_subscriptions(|client| &client.subscriptions_pool_liquidity_updates)
110    }
111
112    /// Returns all instrument IDs for which fee collect subscriptions exist.
113    #[must_use]
114    pub fn subscribed_pool_fee_collects(&self) -> Vec<InstrumentId> {
115        self.collect_subscriptions(|client| &client.subscriptions_pool_fee_collects)
116    }
117
118    /// Returns all instrument IDs for which flash loan subscriptions exist.
119    #[must_use]
120    pub fn subscribed_pool_flash(&self) -> Vec<InstrumentId> {
121        self.collect_subscriptions(|client| &client.subscriptions_pool_flash)
122    }
123
124    /// Handles a subscribe command, updating internal state and forwarding to the client.
125    ///
126    /// # Errors
127    ///
128    /// Returns an error if the subscription is invalid (e.g., synthetic instrument for book data),
129    /// or if the underlying client operation fails.
130    pub fn execute_defi_subscribe(&mut self, cmd: DefiSubscribeCommand) -> anyhow::Result<()> {
131        if let Some(client_id) = cmd.client_id()
132            && self.is_external_client(*client_id)
133        {
134            if self.config().debug {
135                log::debug!("Skipping defi subscribe for external client {client_id}: {cmd:?}");
136            }
137            return Ok(());
138        }
139
140        let Some(client) = self.get_client(cmd.client_id(), cmd.venue()) else {
141            log::error!(
142                "Cannot handle command: no client found for client_id={:?}, venue={:?}",
143                cmd.client_id(),
144                cmd.venue(),
145            );
146            return Ok(());
147        };
148
149        #[rustfmt::skip]
150        let pool = match &cmd {
151            DefiSubscribeCommand::Pool(cmd) => Some((cmd.instrument_id, cmd.client_id)),
152            DefiSubscribeCommand::PoolSwaps(cmd) => Some((cmd.instrument_id, cmd.client_id)),
153            DefiSubscribeCommand::PoolLiquidityUpdates(cmd) => Some((cmd.instrument_id, cmd.client_id)),
154            DefiSubscribeCommand::PoolFeeCollects(cmd) => Some((cmd.instrument_id, cmd.client_id)),
155            DefiSubscribeCommand::PoolFlashEvents(cmd) => Some((cmd.instrument_id, cmd.client_id)),
156            DefiSubscribeCommand::Blocks(_) => None, // No pool setup needed for blocks
157        };
158
159        log::info!("Forwarding subscription to client {}", client.client_id);
160        client.execute_defi_subscribe_with_retained(cmd, true);
161
162        if let Some((instrument_id, client_id)) = pool {
163            self.setup_pool_updater(&instrument_id, client_id.as_ref());
164        }
165
166        Ok(())
167    }
168
169    /// Handles an unsubscribe command, updating internal state and forwarding to the client.
170    ///
171    /// # Errors
172    ///
173    /// Returns an error if the underlying client operation fails.
174    pub fn execute_defi_unsubscribe(&mut self, cmd: &DefiUnsubscribeCommand) -> anyhow::Result<()> {
175        if let Some(client_id) = cmd.client_id()
176            && self.is_external_client(*client_id)
177        {
178            if self.config().debug {
179                log::debug!("Skipping defi unsubscribe for external client {client_id}: {cmd:?}");
180            }
181            return Ok(());
182        }
183
184        let Some(client) = self.get_client(cmd.client_id(), cmd.venue()) else {
185            log::error!(
186                "Cannot handle command: no client found for client_id={:?}, venue={:?}",
187                cmd.client_id(),
188                cmd.venue(),
189            );
190            return Ok(());
191        };
192
193        client.execute_defi_unsubscribe(cmd);
194
195        #[rustfmt::skip]
196        let instrument_id = match cmd {
197            DefiUnsubscribeCommand::Pool(cmd) => cmd.instrument_id,
198            DefiUnsubscribeCommand::PoolSwaps(cmd) => cmd.instrument_id,
199            DefiUnsubscribeCommand::PoolLiquidityUpdates(cmd) => cmd.instrument_id,
200            DefiUnsubscribeCommand::PoolFeeCollects(cmd) => cmd.instrument_id,
201            DefiUnsubscribeCommand::PoolFlashEvents(cmd) => cmd.instrument_id,
202            DefiUnsubscribeCommand::Blocks(_) => return Ok(()),
203        };
204
205        // The updater serves every pool subscription shape on every client
206        if !self
207            .get_clients()
208            .iter()
209            .any(|client| client.has_defi_pool_demand(&instrument_id))
210        {
211            self.release_pool_updater(&instrument_id);
212        }
213
214        Ok(())
215    }
216
217    /// Sends a [`DefiRequestCommand`] to a suitable data client implementation.
218    ///
219    /// # Errors
220    ///
221    /// Returns an error if no client is found for the given client ID or venue,
222    /// or if the client fails to process the request.
223    pub fn execute_defi_request(&mut self, req: DefiRequestCommand) -> anyhow::Result<()> {
224        // Skip requests for external clients
225        if let Some(cid) = req.client_id()
226            && self.is_external_client(*cid)
227        {
228            if self.config().debug {
229                log::debug!("Skipping defi data request for external client {cid}: {req:?}");
230            }
231            return Ok(());
232        }
233
234        if let Some(client) = self.get_client(req.client_id(), req.venue()) {
235            client.execute_defi_request(req)
236        } else {
237            anyhow::bail!(
238                "Cannot handle request: no client found for {:?} {:?}",
239                req.client_id(),
240                req.venue()
241            );
242        }
243    }
244
245    /// Processes DeFi-specific data events.
246    pub fn process_defi_data(&mut self, data: DefiData) {
247        self.increment_data_count();
248
249        match data {
250            DefiData::Block(block) => {
251                let topic = defi::switchboard::get_defi_blocks_topic(block.chain());
252                msgbus::publish_defi_block(topic, &block);
253            }
254            DefiData::Pool(pool) => {
255                if let Err(e) = self.cache().borrow_mut().add_pool(pool.clone()) {
256                    log::error!("Failed to add Pool to cache: {e}");
257                }
258
259                match CurrencyPair::try_from(&pool) {
260                    Ok(instrument) => {
261                        self.handle_instrument(&InstrumentAny::CurrencyPair(instrument));
262                    }
263                    Err(e) => {
264                        log::error!(
265                            "Failed to create instrument for Pool {}: {e}",
266                            pool.instrument_id
267                        );
268                    }
269                }
270
271                let topic = defi::switchboard::get_defi_pool_topic(pool.instrument_id);
272                msgbus::publish_defi_pool(topic, &pool);
273            }
274            DefiData::PoolSnapshot(snapshot) => {
275                // Replayed snapshots carry no request identity, so any pending bootstrap accepts them
276                if !self
277                    .pool_snapshot_pending
278                    .contains_key(&snapshot.instrument_id)
279                {
280                    log::warn!(
281                        "Received unexpected pool snapshot for {} (not in pending set)",
282                        snapshot.instrument_id
283                    );
284                    return;
285                }
286
287                self.hydrate_pool_profiler(&snapshot);
288            }
289            DefiData::PoolSwap(swap) => {
290                let instrument_id = swap.instrument_id;
291                // Subscribers receive events immediately; keep a copy for profiler hydration
292                // while the snapshot is pending.
293                if self.pool_snapshot_pending.contains_key(&instrument_id) {
294                    log::debug!("Buffering swap event for {instrument_id} (waiting for snapshot)");
295                    self.pool_event_buffers
296                        .entry(instrument_id)
297                        .or_default()
298                        .push(DefiData::PoolSwap(swap.clone()));
299                }
300
301                let topic = defi::switchboard::get_defi_pool_swaps_topic(instrument_id);
302                msgbus::publish_defi_swap(topic, &swap);
303            }
304            DefiData::PoolLiquidityUpdate(update) => {
305                let instrument_id = update.instrument_id;
306                // Subscribers receive events immediately; keep a copy for profiler hydration
307                // while the snapshot is pending.
308                if self.pool_snapshot_pending.contains_key(&instrument_id) {
309                    log::debug!(
310                        "Buffering liquidity update event for {instrument_id} (waiting for snapshot)"
311                    );
312                    self.pool_event_buffers
313                        .entry(instrument_id)
314                        .or_default()
315                        .push(DefiData::PoolLiquidityUpdate(update.clone()));
316                }
317
318                let topic = defi::switchboard::get_defi_liquidity_topic(instrument_id);
319                msgbus::publish_defi_liquidity(topic, &update);
320            }
321            DefiData::PoolFeeCollect(collect) => {
322                let instrument_id = collect.instrument_id;
323                // Subscribers receive events immediately; keep a copy for profiler hydration
324                // while the snapshot is pending.
325                if self.pool_snapshot_pending.contains_key(&instrument_id) {
326                    log::debug!(
327                        "Buffering fee collect event for {instrument_id} (waiting for snapshot)"
328                    );
329                    self.pool_event_buffers
330                        .entry(instrument_id)
331                        .or_default()
332                        .push(DefiData::PoolFeeCollect(collect.clone()));
333                }
334
335                let topic = defi::switchboard::get_defi_collect_topic(instrument_id);
336                msgbus::publish_defi_collect(topic, &collect);
337            }
338            DefiData::PoolFeeProtocolUpdate(update) => {
339                let instrument_id = update.instrument_id;
340                // A protocol-fee change is profiler infrastructure, not a subscriber data stream:
341                // it only needs to keep the profiler's `fee_protocol` consistent for swap and flash
342                // fee splitting. Buffer until the snapshot lands, otherwise apply in place.
343                if self.pool_snapshot_pending.contains_key(&instrument_id) {
344                    log::debug!(
345                        "Buffering fee protocol update for {instrument_id} (waiting for snapshot)"
346                    );
347                    self.pool_event_buffers
348                        .entry(instrument_id)
349                        .or_default()
350                        .push(DefiData::PoolFeeProtocolUpdate(update));
351                } else if let Some(profiler) =
352                    self.cache().borrow_mut().pool_profiler_mut(&instrument_id)
353                    && let Err(e) = profiler.process_fee_protocol_update(&update)
354                {
355                    log::error!("Failed to process pool fee protocol update: {e}");
356                }
357            }
358            DefiData::PoolFeeProtocolCollect(collect) => {
359                let instrument_id = collect.instrument_id;
360                // A protocol-fee withdrawal is profiler infrastructure, not a subscriber data
361                // stream: it only decrements the profiler's accrued protocol-fee balances. Buffer
362                // until the snapshot lands, otherwise apply in place.
363                if self.pool_snapshot_pending.contains_key(&instrument_id) {
364                    log::debug!(
365                        "Buffering fee protocol collect event for {instrument_id} (waiting for snapshot)"
366                    );
367                    self.pool_event_buffers
368                        .entry(instrument_id)
369                        .or_default()
370                        .push(DefiData::PoolFeeProtocolCollect(collect));
371                } else if let Some(profiler) =
372                    self.cache().borrow_mut().pool_profiler_mut(&instrument_id)
373                    && let Err(e) = profiler.process_fee_protocol_collect(&collect)
374                {
375                    log::error!("Failed to process pool fee protocol collect event: {e}");
376                }
377            }
378            DefiData::PoolFlash(flash) => {
379                let instrument_id = flash.instrument_id;
380                // Subscribers receive events immediately; keep a copy for profiler hydration
381                // while the snapshot is pending.
382                if self.pool_snapshot_pending.contains_key(&instrument_id) {
383                    log::debug!("Buffering flash event for {instrument_id} (waiting for snapshot)");
384                    self.pool_event_buffers
385                        .entry(instrument_id)
386                        .or_default()
387                        .push(DefiData::PoolFlash(flash.clone()));
388                }
389
390                let topic = defi::switchboard::get_defi_flash_topic(instrument_id);
391                msgbus::publish_defi_flash(topic, &flash);
392            }
393        }
394    }
395
396    pub(crate) fn handle_pool_snapshot_response(&mut self, response: &PoolSnapshotResponse) {
397        let instrument_id = response.data.instrument_id;
398
399        if self.pool_snapshot_pending.get(&instrument_id) != Some(&response.correlation_id) {
400            log::warn!(
401                "Discarding pool snapshot for {instrument_id} from stale request {}",
402                response.correlation_id
403            );
404            return;
405        }
406
407        self.hydrate_pool_profiler(&response.data);
408    }
409
410    fn hydrate_pool_profiler(&mut self, snapshot: &PoolSnapshot) {
411        let instrument_id = snapshot.instrument_id;
412        log::info!(
413            "Received pool snapshot for {instrument_id} at block {} with {} positions and {} ticks",
414            snapshot.block_position.number,
415            snapshot.positions.len(),
416            snapshot.ticks.len()
417        );
418
419        // Get pool from cache
420        let pool = self
421            .cache()
422            .borrow()
423            .pool(&instrument_id)
424            .map(|pool| Arc::new(pool.clone()));
425
426        let Some(pool) = pool else {
427            log::error!("Pool {instrument_id} not found in cache when processing snapshot");
428            self.abandon_pool_snapshot(&instrument_id);
429            return;
430        };
431
432        // Defensive: refuse stub snapshots that slipped past the bootstrap-side
433        // guards. Installing one would leave Python actors observing an initialized
434        // profiler with zero liquidity; better to leave the pool without a profiler
435        // so the bad state is visible.
436        if snapshot.positions.is_empty()
437            && snapshot.ticks.is_empty()
438            && snapshot.block_position.number == pool.creation_block
439        {
440            log::warn!(
441                "Refusing empty stub snapshot for {instrument_id} at pool creation block {}; pool will remain without profiler",
442                snapshot.block_position.number,
443            );
444            self.abandon_pool_snapshot(&instrument_id);
445            return;
446        }
447
448        // Create profiler and restore from snapshot
449        let mut profiler = PoolProfiler::new(pool);
450        if let Err(e) = profiler.restore_from_snapshot(snapshot.clone()) {
451            log::error!("Failed to restore profiler from snapshot for {instrument_id}: {e}");
452            self.abandon_pool_snapshot(&instrument_id);
453            return;
454        }
455
456        log::debug!("Restored pool profiler for {instrument_id} from snapshot");
457
458        // Process buffered events
459        let buffered_events = self
460            .pool_event_buffers
461            .remove(&instrument_id)
462            .unwrap_or_default();
463
464        if !buffered_events.is_empty() {
465            log::info!(
466                "Processing {} buffered events for {instrument_id}",
467                buffered_events.len()
468            );
469
470            let events_to_apply = convert_and_sort_buffered_events(buffered_events);
471            let applied_count = Self::apply_buffered_events_to_profiler(
472                &mut profiler,
473                events_to_apply,
474                &snapshot.block_position,
475                instrument_id,
476            );
477
478            log::info!("Applied {applied_count} buffered events to profiler for {instrument_id}");
479        }
480
481        // Add profiler to cache
482        let result = self.cache().borrow_mut().add_pool_profiler(profiler);
483        if let Err(e) = result {
484            log::error!("Failed to add pool profiler to cache for {instrument_id}: {e}");
485            self.abandon_pool_snapshot(&instrument_id);
486            return;
487        }
488
489        // Create updater and subscribe to topics
490        self.pool_snapshot_pending.remove(&instrument_id);
491        let updater = Rc::new(PoolUpdater::new(&instrument_id, self.cache().clone()));
492
493        self.subscribe_pool_updater_topics(instrument_id, updater.clone());
494        self.pool_updaters.insert(instrument_id, updater);
495
496        log::info!("Pool profiler setup completed for {instrument_id}, now processing live events");
497    }
498
499    fn abandon_pool_snapshot(&mut self, instrument_id: &InstrumentId) {
500        self.pool_snapshot_pending.remove(instrument_id);
501        self.pool_event_buffers.remove(instrument_id);
502    }
503
504    /// Subscribes a pool updater to all relevant pool data topics using typed handlers.
505    fn subscribe_pool_updater_topics(&self, instrument_id: InstrumentId, updater: Rc<PoolUpdater>) {
506        let priority = Some(self.msgbus_priority());
507
508        // Subscribe swap handler
509        let swap_topic = defi::switchboard::get_defi_pool_swaps_topic(instrument_id);
510        let swap_handler = TypedHandler(Rc::new(PoolSwapHandler::new(updater.clone())));
511        msgbus::subscribe_defi_swaps(swap_topic.into(), swap_handler, priority);
512
513        // Subscribe liquidity handler
514        let liq_topic = defi::switchboard::get_defi_liquidity_topic(instrument_id);
515        let liq_handler = TypedHandler(Rc::new(PoolLiquidityHandler::new(updater.clone())));
516        msgbus::subscribe_defi_liquidity(liq_topic.into(), liq_handler, priority);
517
518        // Subscribe collect handler
519        let collect_topic = defi::switchboard::get_defi_collect_topic(instrument_id);
520        let collect_handler = TypedHandler(Rc::new(PoolCollectHandler::new(updater.clone())));
521        msgbus::subscribe_defi_collects(collect_topic.into(), collect_handler, priority);
522
523        // Subscribe flash handler
524        let flash_topic = defi::switchboard::get_defi_flash_topic(instrument_id);
525        let flash_handler = TypedHandler(Rc::new(PoolFlashHandler::new(updater)));
526        msgbus::subscribe_defi_flash(flash_topic.into(), flash_handler, priority);
527    }
528
529    /// Applies buffered events to a pool profiler, filtering to events after the snapshot.
530    ///
531    /// Returns the count of successfully applied events.
532    fn apply_buffered_events_to_profiler(
533        profiler: &mut PoolProfiler,
534        events: Vec<DexPoolData>,
535        snapshot_block: &BlockPosition,
536        instrument_id: InstrumentId,
537    ) -> usize {
538        let mut applied_count = 0;
539
540        for event in events {
541            let event_block = get_event_block_position(&event);
542
543            // Only apply events that occurred after the snapshot
544            let is_after_snapshot = event_block.0 > snapshot_block.number
545                || (event_block.0 == snapshot_block.number
546                    && event_block.1 > snapshot_block.transaction_index)
547                || (event_block.0 == snapshot_block.number
548                    && event_block.1 == snapshot_block.transaction_index
549                    && event_block.2 > snapshot_block.log_index);
550
551            if is_after_snapshot {
552                if let Err(e) = profiler.process(&event) {
553                    log::error!(
554                        "Failed to apply buffered event to profiler for {instrument_id}: {e}"
555                    );
556                } else {
557                    applied_count += 1;
558                }
559            }
560        }
561
562        applied_count
563    }
564
565    fn setup_pool_updater(&mut self, instrument_id: &InstrumentId, client_id: Option<&ClientId>) {
566        // Early return if updater already exists or we are in the middle of setting it up.
567        if self.pool_updaters.contains_key(instrument_id)
568            || self.pool_snapshot_pending.contains_key(instrument_id)
569        {
570            log::debug!("Pool updater for {instrument_id} already exists");
571            return;
572        }
573
574        log::info!("Setting up pool updater for {instrument_id}");
575
576        // Check cache state and ensure profiler exists
577        {
578            let mut cache = self.cache().borrow_mut();
579
580            if cache.pool_profiler(instrument_id).is_some() {
581                // Profiler already exists, proceed to create updater
582                log::debug!("Pool profiler already exists for {instrument_id}");
583            } else if let Some(pool) = cache.pool(instrument_id) {
584                // Pool exists but no profiler, create profiler from pool
585                let pool = Arc::new(pool.clone());
586                let mut pool_profiler = PoolProfiler::new(pool.clone());
587
588                if let Some(initial_sqrt_price_x96) = pool.initial_sqrt_price_x96 {
589                    if let Err(e) = pool_profiler.initialize(initial_sqrt_price_x96) {
590                        log::error!("Failed to initialize pool profiler for {instrument_id}: {e}");
591                        drop(cache);
592                        return;
593                    }
594                    log::debug!(
595                        "Initialized pool profiler for {instrument_id} with sqrt_price {initial_sqrt_price_x96}"
596                    );
597                } else {
598                    log::debug!("Created pool profiler for {instrument_id}");
599                }
600
601                if let Err(e) = cache.add_pool_profiler(pool_profiler) {
602                    log::error!("Failed to add pool profiler for {instrument_id}: {e}");
603                    drop(cache);
604                    return;
605                }
606                drop(cache);
607            } else {
608                // Neither profiler nor pool exists, request snapshot
609                drop(cache);
610
611                let request_id = UUID4::new();
612                let ts_init = self.clock().borrow().timestamp_ns();
613
614                let request = RequestPoolSnapshot::new(
615                    *instrument_id,
616                    client_id.copied(),
617                    request_id,
618                    ts_init,
619                    None,
620                );
621
622                if let Err(e) = self.execute_defi_request(DefiRequestCommand::PoolSnapshot(request))
623                {
624                    log::warn!("Failed to request pool snapshot for {instrument_id}: {e}");
625                } else {
626                    log::debug!("Requested pool snapshot for {instrument_id}");
627                    self.pool_snapshot_pending
628                        .insert(*instrument_id, request_id);
629                    self.pool_event_buffers.entry(*instrument_id).or_default();
630                }
631                return;
632            }
633        }
634
635        // Profiler exists, create updater and subscribe to topics
636        let updater = Rc::new(PoolUpdater::new(instrument_id, self.cache().clone()));
637
638        self.subscribe_pool_updater_topics(*instrument_id, updater.clone());
639        self.pool_updaters.insert(*instrument_id, updater);
640
641        log::debug!("Created PoolUpdater for instrument ID {instrument_id}");
642    }
643
644    fn release_pool_updater(&mut self, instrument_id: &InstrumentId) {
645        if let Some(updater) = self.pool_updaters.remove(instrument_id) {
646            Self::unsubscribe_pool_updater_topics(*instrument_id, updater);
647            log::debug!("Removed PoolUpdater for instrument ID {instrument_id}");
648        }
649
650        self.abandon_pool_snapshot(instrument_id);
651    }
652
653    pub(crate) fn clear_pool_updaters(&mut self) {
654        let updaters: Vec<_> = self.pool_updaters.drain().collect();
655        for (instrument_id, updater) in updaters {
656            Self::unsubscribe_pool_updater_topics(instrument_id, updater);
657        }
658
659        self.abandon_pool_snapshots();
660    }
661
662    pub(crate) fn abandon_pool_snapshots(&mut self) {
663        self.pool_snapshot_pending.clear();
664        self.pool_event_buffers.clear();
665    }
666
667    fn unsubscribe_pool_updater_topics(instrument_id: InstrumentId, updater: Rc<PoolUpdater>) {
668        let swap_topic = defi::switchboard::get_defi_pool_swaps_topic(instrument_id);
669        let swap_handler = TypedHandler(Rc::new(PoolSwapHandler::new(updater.clone())));
670        msgbus::unsubscribe_defi_swaps(swap_topic.into(), &swap_handler);
671
672        let liq_topic = defi::switchboard::get_defi_liquidity_topic(instrument_id);
673        let liq_handler = TypedHandler(Rc::new(PoolLiquidityHandler::new(updater.clone())));
674        msgbus::unsubscribe_defi_liquidity(liq_topic.into(), &liq_handler);
675
676        let collect_topic = defi::switchboard::get_defi_collect_topic(instrument_id);
677        let collect_handler = TypedHandler(Rc::new(PoolCollectHandler::new(updater.clone())));
678        msgbus::unsubscribe_defi_collects(collect_topic.into(), &collect_handler);
679
680        let flash_topic = defi::switchboard::get_defi_flash_topic(instrument_id);
681        let flash_handler = TypedHandler(Rc::new(PoolFlashHandler::new(updater)));
682        msgbus::unsubscribe_defi_flash(flash_topic.into(), &flash_handler);
683    }
684}
685
686#[cfg(test)]
687mod tests {
688    use std::sync::Arc;
689
690    use alloy_primitives::{Address, I256, U160, U256};
691    use nautilus_core::UnixNanos;
692    use nautilus_model::{
693        defi::{
694            Chain, DefiData, PoolFeeCollect, PoolFeeProtocolUpdate, PoolFlash, PoolIdentifier,
695            PoolLiquidityUpdate, PoolLiquidityUpdateType, PoolSwap,
696            chain::chains,
697            data::DexPoolData,
698            dex::{AmmType, Dex, DexType},
699        },
700        identifiers::{InstrumentId, Symbol, Venue},
701    };
702    use rstest::*;
703
704    use super::*;
705
706    #[fixture]
707    fn test_instrument_id() -> InstrumentId {
708        InstrumentId::new(Symbol::from("ETH/USDC"), Venue::from("UNISWAPV3"))
709    }
710
711    #[fixture]
712    fn test_chain() -> Arc<Chain> {
713        Arc::new(chains::ETHEREUM.clone())
714    }
715
716    #[fixture]
717    fn test_dex(test_chain: Arc<Chain>) -> Arc<Dex> {
718        Arc::new(Dex::new(
719            (*test_chain).clone(),
720            DexType::UniswapV3,
721            "0x1F98431c8aD98523631AE4a59f267346ea31F984",
722            12369621,
723            AmmType::CLAMM,
724            "PoolCreated(address,address,uint24,int24,address)",
725            "Swap(address,address,int256,int256,uint160,uint128,int24)",
726            "Mint(address,address,int24,int24,uint128,uint256,uint256)",
727            "Burn(address,int24,int24,uint128,uint256,uint256)",
728            "Collect(address,address,int24,int24,uint128,uint128)",
729        ))
730    }
731
732    fn create_test_swap(
733        test_instrument_id: InstrumentId,
734        test_chain: Arc<Chain>,
735        test_dex: Arc<Dex>,
736        block: u64,
737        tx_index: u32,
738        log_index: u32,
739    ) -> PoolSwap {
740        PoolSwap::new(
741            test_chain,
742            test_dex,
743            test_instrument_id,
744            PoolIdentifier::from_address(Address::ZERO),
745            block,
746            format!("0x{block:064x}"),
747            tx_index,
748            log_index,
749            UnixNanos::default(),
750            UnixNanos::default(),
751            Address::ZERO,
752            Address::ZERO,
753            I256::ZERO,
754            I256::ZERO,
755            U160::ZERO,
756            0,
757            0,
758        )
759    }
760
761    fn create_test_liquidity_update(
762        test_instrument_id: InstrumentId,
763        test_chain: Arc<Chain>,
764        test_dex: Arc<Dex>,
765        block: u64,
766        tx_index: u32,
767        log_index: u32,
768    ) -> PoolLiquidityUpdate {
769        PoolLiquidityUpdate::new(
770            test_chain,
771            test_dex,
772            test_instrument_id,
773            PoolIdentifier::from_address(Address::ZERO),
774            PoolLiquidityUpdateType::Mint,
775            block,
776            format!("0x{block:064x}"),
777            tx_index,
778            log_index,
779            None,
780            Address::ZERO,
781            0,
782            U256::ZERO,
783            U256::ZERO,
784            0,
785            0,
786            UnixNanos::default(),
787            UnixNanos::default(),
788        )
789    }
790
791    fn create_test_fee_collect(
792        test_instrument_id: InstrumentId,
793        test_chain: Arc<Chain>,
794        test_dex: Arc<Dex>,
795        block: u64,
796        tx_index: u32,
797        log_index: u32,
798    ) -> PoolFeeCollect {
799        PoolFeeCollect::new(
800            test_chain,
801            test_dex,
802            test_instrument_id,
803            PoolIdentifier::from_address(Address::ZERO),
804            block,
805            format!("0x{block:064x}"),
806            tx_index,
807            log_index,
808            Address::ZERO,
809            0,
810            0,
811            0,
812            0,
813            UnixNanos::default(),
814            UnixNanos::default(),
815        )
816    }
817
818    fn create_test_flash(
819        test_instrument_id: InstrumentId,
820        test_chain: Arc<Chain>,
821        test_dex: Arc<Dex>,
822        block: u64,
823        tx_index: u32,
824        log_index: u32,
825    ) -> PoolFlash {
826        PoolFlash::new(
827            test_chain,
828            test_dex,
829            test_instrument_id,
830            PoolIdentifier::from_address(Address::ZERO),
831            block,
832            format!("0x{block:064x}"),
833            tx_index,
834            log_index,
835            UnixNanos::default(),
836            UnixNanos::default(),
837            Address::ZERO,
838            Address::ZERO,
839            U256::ZERO,
840            U256::ZERO,
841            U256::ZERO,
842            U256::ZERO,
843        )
844    }
845
846    fn create_test_fee_protocol_update(
847        test_instrument_id: InstrumentId,
848        test_chain: Arc<Chain>,
849        test_dex: Arc<Dex>,
850        block: u64,
851        tx_index: u32,
852        log_index: u32,
853    ) -> PoolFeeProtocolUpdate {
854        PoolFeeProtocolUpdate::new(
855            test_chain,
856            test_dex,
857            test_instrument_id,
858            PoolIdentifier::from_address(Address::ZERO),
859            block,
860            format!("0x{block:064x}"),
861            tx_index,
862            log_index,
863            4,
864            4,
865            UnixNanos::default(),
866            UnixNanos::default(),
867        )
868    }
869
870    #[rstest]
871    fn test_get_event_block_position_swap(
872        test_instrument_id: InstrumentId,
873        test_chain: Arc<Chain>,
874        test_dex: Arc<Dex>,
875    ) {
876        let swap = create_test_swap(test_instrument_id, test_chain, test_dex, 100, 5, 3);
877        let pos = get_event_block_position(&DexPoolData::Swap(swap));
878        assert_eq!(pos, (100, 5, 3));
879    }
880
881    #[rstest]
882    fn test_get_event_block_position_liquidity_update(
883        test_instrument_id: InstrumentId,
884        test_chain: Arc<Chain>,
885        test_dex: Arc<Dex>,
886    ) {
887        let update =
888            create_test_liquidity_update(test_instrument_id, test_chain, test_dex, 200, 10, 7);
889        let pos = get_event_block_position(&DexPoolData::LiquidityUpdate(update));
890        assert_eq!(pos, (200, 10, 7));
891    }
892
893    #[rstest]
894    fn test_get_event_block_position_fee_collect(
895        test_instrument_id: InstrumentId,
896        test_chain: Arc<Chain>,
897        test_dex: Arc<Dex>,
898    ) {
899        let collect = create_test_fee_collect(test_instrument_id, test_chain, test_dex, 300, 15, 2);
900        let pos = get_event_block_position(&DexPoolData::FeeCollect(collect));
901        assert_eq!(pos, (300, 15, 2));
902    }
903
904    #[rstest]
905    fn test_get_event_block_position_flash(
906        test_instrument_id: InstrumentId,
907        test_chain: Arc<Chain>,
908        test_dex: Arc<Dex>,
909    ) {
910        let flash = create_test_flash(test_instrument_id, test_chain, test_dex, 400, 20, 8);
911        let pos = get_event_block_position(&DexPoolData::Flash(flash));
912        assert_eq!(pos, (400, 20, 8));
913    }
914
915    #[rstest]
916    fn test_get_event_block_position_fee_protocol_update(
917        test_instrument_id: InstrumentId,
918        test_chain: Arc<Chain>,
919        test_dex: Arc<Dex>,
920    ) {
921        let update =
922            create_test_fee_protocol_update(test_instrument_id, test_chain, test_dex, 500, 25, 4);
923        let pos = get_event_block_position(&DexPoolData::FeeProtocolUpdate(update));
924        assert_eq!(pos, (500, 25, 4));
925    }
926
927    #[rstest]
928    fn test_convert_and_sort_empty_events() {
929        let events = convert_and_sort_buffered_events(vec![]);
930        assert!(events.is_empty());
931    }
932
933    #[rstest]
934    fn test_convert_and_sort_filters_non_pool_events(
935        test_instrument_id: InstrumentId,
936        test_chain: Arc<Chain>,
937        test_dex: Arc<Dex>,
938    ) {
939        let events = vec![
940            DefiData::PoolSwap(create_test_swap(
941                test_instrument_id,
942                test_chain,
943                test_dex,
944                100,
945                0,
946                0,
947            )),
948            // Block events would be filtered out
949        ];
950        let sorted = convert_and_sort_buffered_events(events);
951        assert_eq!(sorted.len(), 1);
952    }
953
954    #[rstest]
955    fn test_convert_and_sort_single_event(
956        test_instrument_id: InstrumentId,
957        test_chain: Arc<Chain>,
958        test_dex: Arc<Dex>,
959    ) {
960        let swap = create_test_swap(test_instrument_id, test_chain, test_dex, 100, 5, 3);
961        let events = vec![DefiData::PoolSwap(swap)];
962        let sorted = convert_and_sort_buffered_events(events);
963        assert_eq!(sorted.len(), 1);
964        assert_eq!(get_event_block_position(&sorted[0]), (100, 5, 3));
965    }
966
967    #[rstest]
968    fn test_convert_and_sort_already_sorted(
969        test_instrument_id: InstrumentId,
970        test_chain: Arc<Chain>,
971        test_dex: Arc<Dex>,
972    ) {
973        let events = vec![
974            DefiData::PoolSwap(create_test_swap(
975                test_instrument_id,
976                test_chain.clone(),
977                test_dex.clone(),
978                100,
979                0,
980                0,
981            )),
982            DefiData::PoolSwap(create_test_swap(
983                test_instrument_id,
984                test_chain.clone(),
985                test_dex.clone(),
986                100,
987                0,
988                1,
989            )),
990            DefiData::PoolSwap(create_test_swap(
991                test_instrument_id,
992                test_chain,
993                test_dex,
994                100,
995                1,
996                0,
997            )),
998        ];
999        let sorted = convert_and_sort_buffered_events(events);
1000        assert_eq!(sorted.len(), 3);
1001        assert_eq!(get_event_block_position(&sorted[0]), (100, 0, 0));
1002        assert_eq!(get_event_block_position(&sorted[1]), (100, 0, 1));
1003        assert_eq!(get_event_block_position(&sorted[2]), (100, 1, 0));
1004    }
1005
1006    #[rstest]
1007    fn test_convert_and_sort_reverse_order(
1008        test_instrument_id: InstrumentId,
1009        test_chain: Arc<Chain>,
1010        test_dex: Arc<Dex>,
1011    ) {
1012        let events = vec![
1013            DefiData::PoolSwap(create_test_swap(
1014                test_instrument_id,
1015                test_chain.clone(),
1016                test_dex.clone(),
1017                100,
1018                2,
1019                5,
1020            )),
1021            DefiData::PoolSwap(create_test_swap(
1022                test_instrument_id,
1023                test_chain.clone(),
1024                test_dex.clone(),
1025                100,
1026                1,
1027                3,
1028            )),
1029            DefiData::PoolSwap(create_test_swap(
1030                test_instrument_id,
1031                test_chain,
1032                test_dex,
1033                100,
1034                0,
1035                1,
1036            )),
1037        ];
1038        let sorted = convert_and_sort_buffered_events(events);
1039        assert_eq!(sorted.len(), 3);
1040        assert_eq!(get_event_block_position(&sorted[0]), (100, 0, 1));
1041        assert_eq!(get_event_block_position(&sorted[1]), (100, 1, 3));
1042        assert_eq!(get_event_block_position(&sorted[2]), (100, 2, 5));
1043    }
1044
1045    #[rstest]
1046    fn test_convert_and_sort_mixed_blocks(
1047        test_instrument_id: InstrumentId,
1048        test_chain: Arc<Chain>,
1049        test_dex: Arc<Dex>,
1050    ) {
1051        let events = vec![
1052            DefiData::PoolSwap(create_test_swap(
1053                test_instrument_id,
1054                test_chain.clone(),
1055                test_dex.clone(),
1056                102,
1057                0,
1058                0,
1059            )),
1060            DefiData::PoolSwap(create_test_swap(
1061                test_instrument_id,
1062                test_chain.clone(),
1063                test_dex.clone(),
1064                100,
1065                5,
1066                2,
1067            )),
1068            DefiData::PoolSwap(create_test_swap(
1069                test_instrument_id,
1070                test_chain,
1071                test_dex,
1072                101,
1073                3,
1074                1,
1075            )),
1076        ];
1077        let sorted = convert_and_sort_buffered_events(events);
1078        assert_eq!(sorted.len(), 3);
1079        assert_eq!(get_event_block_position(&sorted[0]), (100, 5, 2));
1080        assert_eq!(get_event_block_position(&sorted[1]), (101, 3, 1));
1081        assert_eq!(get_event_block_position(&sorted[2]), (102, 0, 0));
1082    }
1083
1084    #[rstest]
1085    fn test_convert_and_sort_mixed_event_types(
1086        test_instrument_id: InstrumentId,
1087        test_chain: Arc<Chain>,
1088        test_dex: Arc<Dex>,
1089    ) {
1090        let events = vec![
1091            DefiData::PoolSwap(create_test_swap(
1092                test_instrument_id,
1093                test_chain.clone(),
1094                test_dex.clone(),
1095                100,
1096                2,
1097                0,
1098            )),
1099            DefiData::PoolLiquidityUpdate(create_test_liquidity_update(
1100                test_instrument_id,
1101                test_chain.clone(),
1102                test_dex.clone(),
1103                100,
1104                0,
1105                0,
1106            )),
1107            DefiData::PoolFeeCollect(create_test_fee_collect(
1108                test_instrument_id,
1109                test_chain.clone(),
1110                test_dex.clone(),
1111                100,
1112                1,
1113                0,
1114            )),
1115            DefiData::PoolFlash(create_test_flash(
1116                test_instrument_id,
1117                test_chain.clone(),
1118                test_dex.clone(),
1119                100,
1120                3,
1121                0,
1122            )),
1123            DefiData::PoolFeeProtocolUpdate(create_test_fee_protocol_update(
1124                test_instrument_id,
1125                test_chain,
1126                test_dex,
1127                100,
1128                4,
1129                0,
1130            )),
1131        ];
1132        let sorted = convert_and_sort_buffered_events(events);
1133        assert_eq!(sorted.len(), 5);
1134        assert_eq!(get_event_block_position(&sorted[0]), (100, 0, 0));
1135        assert_eq!(get_event_block_position(&sorted[1]), (100, 1, 0));
1136        assert_eq!(get_event_block_position(&sorted[2]), (100, 2, 0));
1137        assert_eq!(get_event_block_position(&sorted[3]), (100, 3, 0));
1138        assert_eq!(get_event_block_position(&sorted[4]), (100, 4, 0));
1139        assert!(matches!(sorted[4], DexPoolData::FeeProtocolUpdate(_)));
1140    }
1141
1142    #[rstest]
1143    fn test_convert_and_sort_same_block_and_tx_different_log_index(
1144        test_instrument_id: InstrumentId,
1145        test_chain: Arc<Chain>,
1146        test_dex: Arc<Dex>,
1147    ) {
1148        let events = vec![
1149            DefiData::PoolSwap(create_test_swap(
1150                test_instrument_id,
1151                test_chain.clone(),
1152                test_dex.clone(),
1153                100,
1154                5,
1155                10,
1156            )),
1157            DefiData::PoolSwap(create_test_swap(
1158                test_instrument_id,
1159                test_chain.clone(),
1160                test_dex.clone(),
1161                100,
1162                5,
1163                5,
1164            )),
1165            DefiData::PoolSwap(create_test_swap(
1166                test_instrument_id,
1167                test_chain,
1168                test_dex,
1169                100,
1170                5,
1171                1,
1172            )),
1173        ];
1174        let sorted = convert_and_sort_buffered_events(events);
1175        assert_eq!(sorted.len(), 3);
1176        assert_eq!(get_event_block_position(&sorted[0]), (100, 5, 1));
1177        assert_eq!(get_event_block_position(&sorted[1]), (100, 5, 5));
1178        assert_eq!(get_event_block_position(&sorted[2]), (100, 5, 10));
1179    }
1180}