nautilus_blockchain/data/
subscription.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
16use ahash::{AHashMap, AHashSet};
17use alloy::primitives::{Address, keccak256};
18use nautilus_core::hex;
19use nautilus_model::defi::DexType;
20
21/// Manages subscriptions to DeFi protocol events (swaps, mints, burns, collects) across different DEXs.
22///
23/// This manager tracks which pool addresses are subscribed for each event type
24/// and maintains the event signature encodings for efficient filtering.
25#[derive(Debug, Default)]
26pub struct DefiDataSubscriptionManager {
27    block_demand_explicit: bool,
28    block_demand_pool_events: bool,
29    block_feed_backend: Option<BlockFeedBackend>,
30    pool_swap_event_encoded: AHashMap<DexType, String>,
31    pool_mint_event_encoded: AHashMap<DexType, String>,
32    pool_burn_event_encoded: AHashMap<DexType, String>,
33    pool_collect_event_encoded: AHashMap<DexType, String>,
34    pool_flash_event_encoded: AHashMap<DexType, String>,
35    pool_fee_protocol_update_event_encoded: AHashMap<DexType, String>,
36    pool_fee_protocol_collect_event_encoded: AHashMap<DexType, String>,
37    subscribed_pool_swaps: AHashMap<DexType, AHashSet<Address>>,
38    subscribed_pool_mints: AHashMap<DexType, AHashSet<Address>>,
39    subscribed_pool_burns: AHashMap<DexType, AHashSet<Address>>,
40    subscribed_pool_collects: AHashMap<DexType, AHashSet<Address>>,
41    subscribed_pool_flashes: AHashMap<DexType, AHashSet<Address>>,
42    subscribed_pool_fee_protocol_updates: AHashMap<DexType, AHashSet<Address>>,
43    subscribed_pool_fee_protocol_collects: AHashMap<DexType, AHashSet<Address>>,
44}
45
46impl DefiDataSubscriptionManager {
47    /// Creates a new [`DefiDataSubscriptionManager`] instance.
48    #[must_use]
49    pub fn new() -> Self {
50        Self::default()
51    }
52
53    pub(crate) fn add_block_demand(
54        &mut self,
55        owner: BlockFeedOwner,
56        preferred_backend: BlockFeedBackend,
57    ) -> Option<BlockFeedBackend> {
58        match owner {
59            BlockFeedOwner::Explicit => self.block_demand_explicit = true,
60            BlockFeedOwner::PoolEvents => self.block_demand_pool_events = true,
61        }
62
63        self.block_feed_backend
64            .is_none()
65            .then_some(preferred_backend)
66    }
67
68    pub(crate) fn block_feed_started(&mut self, backend: BlockFeedBackend) {
69        self.block_feed_backend = Some(backend);
70    }
71
72    pub(crate) fn remove_block_demand(
73        &mut self,
74        owner: BlockFeedOwner,
75    ) -> Option<BlockFeedBackend> {
76        match owner {
77            BlockFeedOwner::Explicit => self.block_demand_explicit = false,
78            BlockFeedOwner::PoolEvents => self.block_demand_pool_events = false,
79        }
80
81        if self.has_block_demand() {
82            None
83        } else {
84            self.block_feed_backend
85        }
86    }
87
88    pub(crate) fn block_feed_stopped(&mut self, backend: BlockFeedBackend) {
89        if self.block_feed_backend == Some(backend) {
90            self.block_feed_backend = None;
91        }
92    }
93
94    pub(crate) fn clear_block_demand(&mut self) {
95        self.block_demand_explicit = false;
96        self.block_demand_pool_events = false;
97        self.block_feed_backend = None;
98    }
99
100    pub(crate) fn has_pool_event_subscriptions(&self) -> bool {
101        self.subscribed_pool_swaps
102            .values()
103            .chain(self.subscribed_pool_mints.values())
104            .chain(self.subscribed_pool_burns.values())
105            .chain(self.subscribed_pool_collects.values())
106            .chain(self.subscribed_pool_flashes.values())
107            .chain(self.subscribed_pool_fee_protocol_updates.values())
108            .chain(self.subscribed_pool_fee_protocol_collects.values())
109            .any(|addresses| !addresses.is_empty())
110    }
111
112    fn has_block_demand(&self) -> bool {
113        self.block_demand_explicit || self.block_demand_pool_events
114    }
115
116    /// Gets all unique contract addresses subscribed for any event type for a given DEX.
117    #[must_use]
118    pub fn get_subscribed_dex_contract_addresses(&self, dex: &DexType) -> Vec<Address> {
119        let mut unique_addresses = AHashSet::new();
120
121        if let Some(addresses) = self.subscribed_pool_swaps.get(dex) {
122            unique_addresses.extend(addresses.iter().copied());
123        }
124
125        if let Some(addresses) = self.subscribed_pool_mints.get(dex) {
126            unique_addresses.extend(addresses.iter().copied());
127        }
128
129        if let Some(addresses) = self.subscribed_pool_burns.get(dex) {
130            unique_addresses.extend(addresses.iter().copied());
131        }
132
133        if let Some(addresses) = self.subscribed_pool_collects.get(dex) {
134            unique_addresses.extend(addresses.iter().copied());
135        }
136
137        if let Some(addresses) = self.subscribed_pool_flashes.get(dex) {
138            unique_addresses.extend(addresses.iter().copied());
139        }
140
141        if let Some(addresses) = self.subscribed_pool_fee_protocol_updates.get(dex) {
142            unique_addresses.extend(addresses.iter().copied());
143        }
144
145        if let Some(addresses) = self.subscribed_pool_fee_protocol_collects.get(dex) {
146            unique_addresses.extend(addresses.iter().copied());
147        }
148
149        unique_addresses.into_iter().collect()
150    }
151
152    /// Gets all event signatures (keccak256 hashes) registered for a given DEX.
153    #[must_use]
154    pub fn get_subscribed_dex_event_signatures(&self, dex: &DexType) -> Vec<String> {
155        let mut result = Vec::new();
156
157        if let Some(swap_event_signature) = self.pool_swap_event_encoded.get(dex) {
158            result.push(swap_event_signature.clone());
159        }
160
161        if let Some(mint_event_signature) = self.pool_mint_event_encoded.get(dex) {
162            result.push(mint_event_signature.clone());
163        }
164
165        if let Some(burn_event_signature) = self.pool_burn_event_encoded.get(dex) {
166            result.push(burn_event_signature.clone());
167        }
168
169        if let Some(collect_event_signature) = self.pool_collect_event_encoded.get(dex) {
170            result.push(collect_event_signature.clone());
171        }
172
173        if let Some(flash_event_signature) = self.pool_flash_event_encoded.get(dex) {
174            result.push(flash_event_signature.clone());
175        }
176
177        if let Some(fee_protocol_update_event_signature) =
178            self.pool_fee_protocol_update_event_encoded.get(dex)
179        {
180            result.push(fee_protocol_update_event_signature.clone());
181        }
182
183        if let Some(fee_protocol_collect_event_signature) =
184            self.pool_fee_protocol_collect_event_encoded.get(dex)
185        {
186            result.push(fee_protocol_collect_event_signature.clone());
187        }
188
189        result
190    }
191
192    /// Gets event signatures with at least one subscribed pool address for a given DEX.
193    #[must_use]
194    pub fn get_active_subscribed_dex_event_signatures(&self, dex: &DexType) -> Vec<String> {
195        let mut result = Vec::new();
196
197        if self
198            .subscribed_pool_swaps
199            .get(dex)
200            .is_some_and(|addresses| !addresses.is_empty())
201            && let Some(signature) = self.pool_swap_event_encoded.get(dex)
202        {
203            result.push(signature.clone());
204        }
205
206        if self
207            .subscribed_pool_mints
208            .get(dex)
209            .is_some_and(|addresses| !addresses.is_empty())
210            && let Some(signature) = self.pool_mint_event_encoded.get(dex)
211        {
212            result.push(signature.clone());
213        }
214
215        if self
216            .subscribed_pool_burns
217            .get(dex)
218            .is_some_and(|addresses| !addresses.is_empty())
219            && let Some(signature) = self.pool_burn_event_encoded.get(dex)
220        {
221            result.push(signature.clone());
222        }
223
224        if self
225            .subscribed_pool_collects
226            .get(dex)
227            .is_some_and(|addresses| !addresses.is_empty())
228            && let Some(signature) = self.pool_collect_event_encoded.get(dex)
229        {
230            result.push(signature.clone());
231        }
232
233        if self
234            .subscribed_pool_flashes
235            .get(dex)
236            .is_some_and(|addresses| !addresses.is_empty())
237            && let Some(signature) = self.pool_flash_event_encoded.get(dex)
238        {
239            result.push(signature.clone());
240        }
241
242        if self
243            .subscribed_pool_fee_protocol_updates
244            .get(dex)
245            .is_some_and(|addresses| !addresses.is_empty())
246            && let Some(signature) = self.pool_fee_protocol_update_event_encoded.get(dex)
247        {
248            result.push(signature.clone());
249        }
250
251        if self
252            .subscribed_pool_fee_protocol_collects
253            .get(dex)
254            .is_some_and(|addresses| !addresses.is_empty())
255            && let Some(signature) = self.pool_fee_protocol_collect_event_encoded.get(dex)
256        {
257            result.push(signature.clone());
258        }
259
260        result
261    }
262
263    /// Gets pool addresses subscribed to swap events for a given DEX.
264    #[must_use]
265    pub fn get_subscribed_pool_swap_addresses(&self, dex: &DexType) -> Vec<Address> {
266        self.subscribed_pool_swaps
267            .get(dex)
268            .map(|addresses| addresses.iter().copied().collect())
269            .unwrap_or_default()
270    }
271
272    /// Gets pool addresses subscribed to mint events for a given DEX.
273    #[must_use]
274    pub fn get_subscribed_pool_mint_addresses(&self, dex: &DexType) -> Vec<Address> {
275        self.subscribed_pool_mints
276            .get(dex)
277            .map(|addresses| addresses.iter().copied().collect())
278            .unwrap_or_default()
279    }
280
281    /// Gets pool addresses subscribed to burn events for a given DEX.
282    #[must_use]
283    pub fn get_subscribed_pool_burn_addresses(&self, dex: &DexType) -> Vec<Address> {
284        self.subscribed_pool_burns
285            .get(dex)
286            .map(|addresses| addresses.iter().copied().collect())
287            .unwrap_or_default()
288    }
289
290    /// Gets pool addresses subscribed to collect events for a given DEX.
291    #[must_use]
292    pub fn get_subscribed_pool_collect_addresses(&self, dex: &DexType) -> Vec<Address> {
293        self.subscribed_pool_collects
294            .get(dex)
295            .map(|addresses| addresses.iter().copied().collect())
296            .unwrap_or_default()
297    }
298
299    /// Gets pool addresses subscribed to flash events for a given DEX.
300    #[must_use]
301    pub fn get_subscribed_pool_flash_addresses(&self, dex: &DexType) -> Vec<Address> {
302        self.subscribed_pool_flashes
303            .get(dex)
304            .map(|addresses| addresses.iter().copied().collect())
305            .unwrap_or_default()
306    }
307
308    /// Gets pool addresses subscribed to fee-protocol update events for a given DEX.
309    #[must_use]
310    pub fn get_subscribed_pool_fee_protocol_update_addresses(&self, dex: &DexType) -> Vec<Address> {
311        self.subscribed_pool_fee_protocol_updates
312            .get(dex)
313            .map(|addresses| addresses.iter().copied().collect())
314            .unwrap_or_default()
315    }
316
317    /// Gets pool addresses subscribed to fee-protocol collect events for a given DEX.
318    #[must_use]
319    pub fn get_subscribed_pool_fee_protocol_collect_addresses(
320        &self,
321        dex: &DexType,
322    ) -> Vec<Address> {
323        self.subscribed_pool_fee_protocol_collects
324            .get(dex)
325            .map(|addresses| addresses.iter().copied().collect())
326            .unwrap_or_default()
327    }
328
329    /// Gets the swap event signature for a specific DEX.
330    #[must_use]
331    pub fn get_dex_pool_swap_event_signature(&self, dex: &DexType) -> Option<String> {
332        self.pool_swap_event_encoded.get(dex).cloned()
333    }
334
335    /// Gets the mint event signature for a specific DEX.
336    #[must_use]
337    pub fn get_dex_pool_mint_event_signature(&self, dex: &DexType) -> Option<String> {
338        self.pool_mint_event_encoded.get(dex).cloned()
339    }
340    /// Gets the burn event signature for a specific DEX.
341    #[must_use]
342    pub fn get_dex_pool_burn_event_signature(&self, dex: &DexType) -> Option<String> {
343        self.pool_burn_event_encoded.get(dex).cloned()
344    }
345
346    /// Gets the collect event signature for a specific DEX.
347    #[must_use]
348    pub fn get_dex_pool_collect_event_signature(&self, dex: &DexType) -> Option<String> {
349        self.pool_collect_event_encoded.get(dex).cloned()
350    }
351
352    /// Gets the flash event signature for a specific DEX.
353    #[must_use]
354    pub fn get_dex_pool_flash_event_signature(&self, dex: &DexType) -> Option<String> {
355        self.pool_flash_event_encoded.get(dex).cloned()
356    }
357
358    /// Gets the fee-protocol update event signature for a specific DEX.
359    #[must_use]
360    pub fn get_dex_pool_fee_protocol_update_event_signature(
361        &self,
362        dex: &DexType,
363    ) -> Option<String> {
364        self.pool_fee_protocol_update_event_encoded
365            .get(dex)
366            .cloned()
367    }
368
369    /// Gets the fee-protocol collect event signature for a specific DEX.
370    #[must_use]
371    pub fn get_dex_pool_fee_protocol_collect_event_signature(
372        &self,
373        dex: &DexType,
374    ) -> Option<String> {
375        self.pool_fee_protocol_collect_event_encoded
376            .get(dex)
377            .cloned()
378    }
379
380    /// Normalizes an event signature to a consistent format.
381    ///
382    /// Accepts:
383    /// - A raw event signature like "Swap(address,address,int256,int256,uint160,uint128,int24)".
384    /// - A pre-encoded topic like "0xc42079f94a6350d7e6235f29174924f928cc2ac818eb64fed8004e115fbcca67".
385    /// - A hex string without 0x prefix.
386    ///
387    /// Returns a normalized "0x..." format string.
388    fn normalize_topic(sig: &str) -> String {
389        let s = sig.trim();
390
391        // Check if it's already a properly formatted hex string with 0x prefix
392        if let Some(rest) = s.strip_prefix("0x")
393            && rest.len() == 64
394            && rest.chars().all(|c| c.is_ascii_hexdigit())
395        {
396            return format!("0x{}", rest.to_ascii_lowercase());
397        }
398
399        // Check if it's a hex string without 0x prefix
400        if s.len() == 64 && s.chars().all(|c| c.is_ascii_hexdigit()) {
401            return format!("0x{}", s.to_ascii_lowercase());
402        }
403
404        // Otherwise, it's a raw signature that needs hashing
405        hex::encode_prefixed(keccak256(s.as_bytes()))
406    }
407
408    /// Registers a DEX with its event signatures for subscription management.
409    ///
410    /// This must be called before subscribing to any events for a DEX.
411    /// Event signatures can be either raw signatures or pre-encoded keccak256 hashes.
412    pub fn register_dex_for_subscriptions(
413        &mut self,
414        dex: DexType,
415        swap_event_signature: &str,
416        mint_event_signature: &str,
417        burn_event_signature: &str,
418        collect_event_signature: &str,
419        flash_event_signature: Option<&str>,
420    ) {
421        self.subscribed_pool_swaps.insert(dex, AHashSet::new());
422        self.pool_swap_event_encoded
423            .insert(dex, Self::normalize_topic(swap_event_signature));
424
425        self.subscribed_pool_mints.insert(dex, AHashSet::new());
426        self.pool_mint_event_encoded
427            .insert(dex, Self::normalize_topic(mint_event_signature));
428
429        self.subscribed_pool_burns.insert(dex, AHashSet::new());
430        self.pool_burn_event_encoded
431            .insert(dex, Self::normalize_topic(burn_event_signature));
432
433        self.subscribed_pool_collects.insert(dex, AHashSet::new());
434        self.pool_collect_event_encoded
435            .insert(dex, Self::normalize_topic(collect_event_signature));
436
437        if let Some(flash_event_signature) = flash_event_signature {
438            self.subscribed_pool_flashes.insert(dex, AHashSet::new());
439            self.pool_flash_event_encoded
440                .insert(dex, Self::normalize_topic(flash_event_signature));
441        }
442
443        log::debug!("Registered DEX for subscriptions: {dex:?}");
444    }
445
446    /// Registers optional fee-protocol event signatures for subscription management.
447    pub fn register_dex_fee_protocol_events(
448        &mut self,
449        dex: DexType,
450        fee_protocol_update_event_signature: Option<&str>,
451        fee_protocol_collect_event_signature: Option<&str>,
452    ) {
453        if let Some(fee_protocol_update_event_signature) = fee_protocol_update_event_signature {
454            self.subscribed_pool_fee_protocol_updates
455                .insert(dex, AHashSet::new());
456            self.pool_fee_protocol_update_event_encoded.insert(
457                dex,
458                Self::normalize_topic(fee_protocol_update_event_signature),
459            );
460        }
461
462        if let Some(fee_protocol_collect_event_signature) = fee_protocol_collect_event_signature {
463            self.subscribed_pool_fee_protocol_collects
464                .insert(dex, AHashSet::new());
465            self.pool_fee_protocol_collect_event_encoded.insert(
466                dex,
467                Self::normalize_topic(fee_protocol_collect_event_signature),
468            );
469        }
470    }
471
472    /// Subscribes to swap events for a specific pool address on a DEX.
473    pub fn subscribe_swaps(&mut self, dex: DexType, address: Address) {
474        if let Some(pool_set) = self.subscribed_pool_swaps.get_mut(&dex) {
475            pool_set.insert(address);
476        } else {
477            log::error!("DEX not registered for swap subscriptions: {dex:?}");
478        }
479    }
480
481    /// Subscribes to mint events for a specific pool address on a DEX.
482    pub fn subscribe_mints(&mut self, dex: DexType, address: Address) {
483        if let Some(pool_set) = self.subscribed_pool_mints.get_mut(&dex) {
484            pool_set.insert(address);
485        } else {
486            log::error!("DEX not registered for mint subscriptions: {dex:?}");
487        }
488    }
489
490    /// Subscribes to burn events for a specific pool address on a DEX.
491    pub fn subscribe_burns(&mut self, dex: DexType, address: Address) {
492        if let Some(pool_set) = self.subscribed_pool_burns.get_mut(&dex) {
493            pool_set.insert(address);
494        } else {
495            log::warn!("DEX not registered for burn subscriptions: {dex:?}");
496        }
497    }
498
499    /// Unsubscribes from swap events for a specific pool address on a DEX.
500    pub fn unsubscribe_swaps(&mut self, dex: DexType, address: Address) {
501        if let Some(pool_set) = self.subscribed_pool_swaps.get_mut(&dex) {
502            pool_set.remove(&address);
503        } else {
504            log::error!("DEX not registered for swap subscriptions: {dex:?}");
505        }
506    }
507
508    /// Unsubscribes from mint events for a specific pool address on a DEX.
509    pub fn unsubscribe_mints(&mut self, dex: DexType, address: Address) {
510        if let Some(pool_set) = self.subscribed_pool_mints.get_mut(&dex) {
511            pool_set.remove(&address);
512        } else {
513            log::error!("DEX not registered for mint subscriptions: {dex:?}");
514        }
515    }
516
517    /// Unsubscribes from burn events for a specific pool address on a DEX.
518    pub fn unsubscribe_burns(&mut self, dex: DexType, address: Address) {
519        if let Some(pool_set) = self.subscribed_pool_burns.get_mut(&dex) {
520            pool_set.remove(&address);
521        } else {
522            log::error!("DEX not registered for burn subscriptions: {dex:?}");
523        }
524    }
525
526    /// Subscribes to collect events for a specific pool address on a DEX.
527    pub fn subscribe_collects(&mut self, dex: DexType, address: Address) {
528        if let Some(pool_set) = self.subscribed_pool_collects.get_mut(&dex) {
529            pool_set.insert(address);
530        } else {
531            log::error!("DEX not registered for collect subscriptions: {dex:?}");
532        }
533    }
534
535    /// Unsubscribes from collect events for a specific pool address on a DEX.
536    pub fn unsubscribe_collects(&mut self, dex: DexType, address: Address) {
537        if let Some(pool_set) = self.subscribed_pool_collects.get_mut(&dex) {
538            pool_set.remove(&address);
539        } else {
540            log::error!("DEX not registered for collect subscriptions: {dex:?}");
541        }
542    }
543
544    /// Subscribes to flash events for a specific pool address on a DEX.
545    pub fn subscribe_flashes(&mut self, dex: DexType, address: Address) {
546        if let Some(pool_set) = self.subscribed_pool_flashes.get_mut(&dex) {
547            pool_set.insert(address);
548        } else {
549            log::error!("DEX not registered for flash subscriptions: {dex:?}");
550        }
551    }
552
553    /// Unsubscribes from flash events for a specific pool address on a DEX.
554    pub fn unsubscribe_flashes(&mut self, dex: DexType, address: Address) {
555        if let Some(pool_set) = self.subscribed_pool_flashes.get_mut(&dex) {
556            pool_set.remove(&address);
557        } else {
558            log::error!("DEX not registered for flash subscriptions: {dex:?}");
559        }
560    }
561
562    /// Subscribes to fee-protocol update events for a specific pool address on a DEX.
563    pub fn subscribe_fee_protocol_updates(&mut self, dex: DexType, address: Address) {
564        if let Some(pool_set) = self.subscribed_pool_fee_protocol_updates.get_mut(&dex) {
565            pool_set.insert(address);
566        }
567    }
568
569    /// Unsubscribes from fee-protocol update events for a specific pool address on a DEX.
570    pub fn unsubscribe_fee_protocol_updates(&mut self, dex: DexType, address: Address) {
571        if let Some(pool_set) = self.subscribed_pool_fee_protocol_updates.get_mut(&dex) {
572            pool_set.remove(&address);
573        }
574    }
575
576    /// Subscribes to fee-protocol collect events for a specific pool address on a DEX.
577    pub fn subscribe_fee_protocol_collects(&mut self, dex: DexType, address: Address) {
578        if let Some(pool_set) = self.subscribed_pool_fee_protocol_collects.get_mut(&dex) {
579            pool_set.insert(address);
580        }
581    }
582
583    /// Unsubscribes from fee-protocol collect events for a specific pool address on a DEX.
584    pub fn unsubscribe_fee_protocol_collects(&mut self, dex: DexType, address: Address) {
585        if let Some(pool_set) = self.subscribed_pool_fee_protocol_collects.get_mut(&dex) {
586            pool_set.remove(&address);
587        }
588    }
589}
590
591#[derive(Debug, Clone, Copy, PartialEq, Eq)]
592pub(crate) enum BlockFeedBackend {
593    Rpc,
594    HyperSync,
595}
596
597#[derive(Debug, Clone, Copy, PartialEq, Eq)]
598pub(crate) enum BlockFeedOwner {
599    Explicit,
600    PoolEvents,
601}
602
603#[cfg(test)]
604mod tests {
605    use alloy::primitives::address;
606    use rstest::{fixture, rstest};
607
608    use super::*;
609
610    #[fixture]
611    fn manager() -> DefiDataSubscriptionManager {
612        DefiDataSubscriptionManager::new()
613    }
614
615    #[fixture]
616    fn registered_manager() -> DefiDataSubscriptionManager {
617        let mut manager = DefiDataSubscriptionManager::new();
618        manager.register_dex_for_subscriptions(
619            DexType::UniswapV3,
620            "Swap(address,address,int256,int256,uint160,uint128,int24)",
621            "Mint(address,address,int24,int24,uint128,uint256,uint256)",
622            "Burn(address,int24,int24,uint128,uint256,uint256)",
623            "Collect(address,address,int24,int24,uint128,uint128)",
624            Some("Flash(address,address,uint256,uint256,uint256,uint256)"),
625        );
626        manager
627    }
628
629    #[rstest]
630    #[case(BlockFeedBackend::Rpc)]
631    #[case(BlockFeedBackend::HyperSync)]
632    fn block_feed_pool_only_stops_with_final_owner(#[case] backend: BlockFeedBackend) {
633        let mut manager = DefiDataSubscriptionManager::new();
634
635        assert_eq!(
636            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
637            Some(backend)
638        );
639        manager.block_feed_started(backend);
640
641        assert_eq!(
642            manager.remove_block_demand(BlockFeedOwner::PoolEvents),
643            Some(backend)
644        );
645        manager.block_feed_stopped(backend);
646
647        assert_eq!(
648            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
649            Some(backend)
650        );
651    }
652
653    #[rstest]
654    #[case(BlockFeedBackend::Rpc)]
655    #[case(BlockFeedBackend::HyperSync)]
656    fn block_feed_repeated_commands_are_idempotent(#[case] backend: BlockFeedBackend) {
657        let mut manager = DefiDataSubscriptionManager::new();
658
659        assert_eq!(
660            manager.add_block_demand(BlockFeedOwner::Explicit, backend),
661            Some(backend)
662        );
663        manager.block_feed_started(backend);
664        assert_eq!(
665            manager.add_block_demand(BlockFeedOwner::Explicit, backend),
666            None
667        );
668
669        assert_eq!(
670            manager.remove_block_demand(BlockFeedOwner::Explicit),
671            Some(backend)
672        );
673        manager.block_feed_stopped(backend);
674        assert_eq!(manager.remove_block_demand(BlockFeedOwner::Explicit), None);
675    }
676
677    #[rstest]
678    #[case(BlockFeedBackend::Rpc)]
679    #[case(BlockFeedBackend::HyperSync)]
680    fn block_feed_explicit_demand_survives_pool_unsubscribe(#[case] backend: BlockFeedBackend) {
681        let mut manager = DefiDataSubscriptionManager::new();
682
683        assert_eq!(
684            manager.add_block_demand(BlockFeedOwner::Explicit, backend),
685            Some(backend)
686        );
687        manager.block_feed_started(backend);
688        assert_eq!(
689            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
690            None
691        );
692
693        assert_eq!(
694            manager.remove_block_demand(BlockFeedOwner::PoolEvents),
695            None
696        );
697        assert_eq!(
698            manager.remove_block_demand(BlockFeedOwner::Explicit),
699            Some(backend)
700        );
701    }
702
703    #[rstest]
704    #[case(BlockFeedBackend::Rpc)]
705    #[case(BlockFeedBackend::HyperSync)]
706    fn block_feed_pool_demand_survives_explicit_unsubscribe(#[case] backend: BlockFeedBackend) {
707        let mut manager = DefiDataSubscriptionManager::new();
708
709        assert_eq!(
710            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
711            Some(backend)
712        );
713        manager.block_feed_started(backend);
714        assert_eq!(
715            manager.add_block_demand(BlockFeedOwner::Explicit, backend),
716            None
717        );
718
719        assert_eq!(manager.remove_block_demand(BlockFeedOwner::Explicit), None);
720        assert_eq!(
721            manager.remove_block_demand(BlockFeedOwner::PoolEvents),
722            Some(backend)
723        );
724    }
725
726    #[rstest]
727    #[case(BlockFeedBackend::Rpc)]
728    #[case(BlockFeedBackend::HyperSync)]
729    fn block_feed_pool_demand_spans_dexes_addresses_and_events(#[case] backend: BlockFeedBackend) {
730        let mut manager = DefiDataSubscriptionManager::new();
731        let pool1 = address!("1111111111111111111111111111111111111111");
732        let pool2 = address!("2222222222222222222222222222222222222222");
733        let pool3 = address!("3333333333333333333333333333333333333333");
734
735        for dex in [DexType::UniswapV3, DexType::PancakeSwapV3] {
736            manager.register_dex_for_subscriptions(
737                dex,
738                "Swap(address,address,int256,int256,uint160,uint128,int24)",
739                "Mint(address,address,int24,int24,uint128,uint256,uint256)",
740                "Burn(address,int24,int24,uint128,uint256,uint256)",
741                "Collect(address,address,int24,int24,uint128,uint128)",
742                Some("Flash(address,address,uint256,uint256,uint256,uint256)"),
743            );
744        }
745
746        manager.subscribe_swaps(DexType::UniswapV3, pool1);
747        manager.subscribe_mints(DexType::UniswapV3, pool2);
748        manager.subscribe_flashes(DexType::PancakeSwapV3, pool3);
749
750        assert!(manager.has_pool_event_subscriptions());
751        assert_eq!(
752            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
753            Some(backend)
754        );
755        manager.block_feed_started(backend);
756
757        manager.unsubscribe_swaps(DexType::UniswapV3, pool1);
758        assert!(manager.has_pool_event_subscriptions());
759        assert_eq!(
760            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
761            None
762        );
763
764        manager.unsubscribe_mints(DexType::UniswapV3, pool2);
765        assert!(manager.has_pool_event_subscriptions());
766        assert_eq!(
767            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
768            None
769        );
770
771        manager.unsubscribe_flashes(DexType::PancakeSwapV3, pool3);
772        assert!(!manager.has_pool_event_subscriptions());
773        assert_eq!(
774            manager.remove_block_demand(BlockFeedOwner::PoolEvents),
775            Some(backend)
776        );
777    }
778
779    #[rstest]
780    #[case(BlockFeedBackend::Rpc)]
781    #[case(BlockFeedBackend::HyperSync)]
782    fn block_feed_disconnect_clears_owners_and_backend(#[case] backend: BlockFeedBackend) {
783        let mut manager = DefiDataSubscriptionManager::new();
784
785        assert_eq!(
786            manager.add_block_demand(BlockFeedOwner::Explicit, backend),
787            Some(backend)
788        );
789        manager.block_feed_started(backend);
790        assert_eq!(
791            manager.add_block_demand(BlockFeedOwner::PoolEvents, backend),
792            None
793        );
794
795        manager.clear_block_demand();
796
797        assert!(!manager.block_demand_explicit);
798        assert!(!manager.block_demand_pool_events);
799        assert_eq!(manager.block_feed_backend, None);
800        assert_eq!(
801            manager.add_block_demand(BlockFeedOwner::Explicit, backend),
802            Some(backend)
803        );
804    }
805
806    #[rstest]
807    fn block_feed_rpc_fallback_stops_hypersync_backend() {
808        let mut manager = DefiDataSubscriptionManager::new();
809
810        assert_eq!(
811            manager.add_block_demand(BlockFeedOwner::Explicit, BlockFeedBackend::Rpc),
812            Some(BlockFeedBackend::Rpc)
813        );
814        manager.block_feed_started(BlockFeedBackend::HyperSync);
815
816        assert_eq!(
817            manager.remove_block_demand(BlockFeedOwner::Explicit),
818            Some(BlockFeedBackend::HyperSync)
819        );
820    }
821
822    #[rstest]
823    fn test_new_creates_empty_manager(manager: DefiDataSubscriptionManager) {
824        assert_eq!(
825            manager
826                .get_subscribed_dex_contract_addresses(&DexType::UniswapV3)
827                .len(),
828            0
829        );
830        assert_eq!(
831            manager
832                .get_subscribed_dex_event_signatures(&DexType::UniswapV3)
833                .len(),
834            0
835        );
836        assert!(
837            manager
838                .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
839                .is_none()
840        );
841        assert!(
842            manager
843                .get_dex_pool_mint_event_signature(&DexType::UniswapV3)
844                .is_none()
845        );
846        assert!(
847            manager
848                .get_dex_pool_burn_event_signature(&DexType::UniswapV3)
849                .is_none()
850        );
851    }
852
853    #[rstest]
854    fn test_register_dex_for_subscriptions(registered_manager: DefiDataSubscriptionManager) {
855        // Should have all four event signatures
856        let signatures =
857            registered_manager.get_subscribed_dex_event_signatures(&DexType::UniswapV3);
858        assert_eq!(signatures.len(), 5);
859
860        // Each signature should be properly encoded
861        assert!(
862            registered_manager
863                .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
864                .is_some()
865        );
866        assert!(
867            registered_manager
868                .get_dex_pool_mint_event_signature(&DexType::UniswapV3)
869                .is_some()
870        );
871        assert!(
872            registered_manager
873                .get_dex_pool_burn_event_signature(&DexType::UniswapV3)
874                .is_some()
875        );
876    }
877
878    #[rstest]
879    fn test_subscribe_and_get_addresses(mut registered_manager: DefiDataSubscriptionManager) {
880        let pool_address = address!("1234567890123456789012345678901234567890");
881
882        // Subscribe to swap events
883        registered_manager.subscribe_swaps(DexType::UniswapV3, pool_address);
884
885        let addresses =
886            registered_manager.get_subscribed_dex_contract_addresses(&DexType::UniswapV3);
887        assert_eq!(addresses.len(), 1);
888        assert_eq!(addresses[0], pool_address);
889    }
890
891    #[rstest]
892    fn test_subscribe_to_unregistered_dex(mut manager: DefiDataSubscriptionManager) {
893        let pool_address = address!("1234567890123456789012345678901234567890");
894
895        // Try to subscribe without registering - should log warning but not panic
896        manager.subscribe_swaps(DexType::UniswapV3, pool_address);
897        manager.subscribe_mints(DexType::UniswapV3, pool_address);
898        manager.subscribe_burns(DexType::UniswapV3, pool_address);
899
900        // Should return empty results
901        let addresses = manager.get_subscribed_dex_contract_addresses(&DexType::UniswapV3);
902        assert_eq!(addresses.len(), 0);
903    }
904
905    #[rstest]
906    fn test_unsubscribe_removes_address(mut registered_manager: DefiDataSubscriptionManager) {
907        let pool_address = address!("1234567890123456789012345678901234567890");
908
909        // Subscribe
910        registered_manager.subscribe_swaps(DexType::UniswapV3, pool_address);
911
912        // Verify subscription
913        assert_eq!(
914            registered_manager
915                .get_subscribed_dex_contract_addresses(&DexType::UniswapV3)
916                .len(),
917            1
918        );
919
920        // Unsubscribe
921        registered_manager.unsubscribe_swaps(DexType::UniswapV3, pool_address);
922
923        // Verify removal
924        assert_eq!(
925            registered_manager
926                .get_subscribed_dex_contract_addresses(&DexType::UniswapV3)
927                .len(),
928            0
929        );
930    }
931
932    #[rstest]
933    fn test_get_event_signatures(registered_manager: DefiDataSubscriptionManager) {
934        let swap_sig = registered_manager.get_dex_pool_swap_event_signature(&DexType::UniswapV3);
935        let mint_sig = registered_manager.get_dex_pool_mint_event_signature(&DexType::UniswapV3);
936        let burn_sig = registered_manager.get_dex_pool_burn_event_signature(&DexType::UniswapV3);
937
938        // All should be Some and start with 0x
939        assert!(swap_sig.is_some() && swap_sig.unwrap().starts_with("0x"));
940        assert!(mint_sig.is_some() && mint_sig.unwrap().starts_with("0x"));
941        assert!(burn_sig.is_some() && burn_sig.unwrap().starts_with("0x"));
942    }
943
944    #[rstest]
945    fn test_active_event_signatures_only_include_subscribed_event_types(
946        mut registered_manager: DefiDataSubscriptionManager,
947    ) {
948        let pool_address = address!("1234567890123456789012345678901234567890");
949        let swap_sig = registered_manager
950            .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
951            .unwrap();
952        let collect_sig = registered_manager
953            .get_dex_pool_collect_event_signature(&DexType::UniswapV3)
954            .unwrap();
955
956        assert!(
957            registered_manager
958                .get_active_subscribed_dex_event_signatures(&DexType::UniswapV3)
959                .is_empty()
960        );
961
962        registered_manager.subscribe_swaps(DexType::UniswapV3, pool_address);
963        assert_eq!(
964            registered_manager.get_active_subscribed_dex_event_signatures(&DexType::UniswapV3),
965            vec![swap_sig]
966        );
967
968        registered_manager.subscribe_collects(DexType::UniswapV3, pool_address);
969        let active_signatures =
970            registered_manager.get_active_subscribed_dex_event_signatures(&DexType::UniswapV3);
971
972        assert_eq!(active_signatures.len(), 2);
973        assert!(active_signatures.contains(&collect_sig));
974    }
975
976    #[rstest]
977    fn test_fee_protocol_events_are_active_only_when_subscribed(
978        mut registered_manager: DefiDataSubscriptionManager,
979    ) {
980        let pool_address = address!("1234567890123456789012345678901234567890");
981        registered_manager.register_dex_fee_protocol_events(
982            DexType::UniswapV3,
983            Some("SetFeeProtocol(uint8,uint8,uint8,uint8)"),
984            Some("CollectProtocol(address,address,uint128,uint128)"),
985        );
986
987        let update_sig = registered_manager
988            .get_dex_pool_fee_protocol_update_event_signature(&DexType::UniswapV3)
989            .unwrap();
990        let collect_sig = registered_manager
991            .get_dex_pool_fee_protocol_collect_event_signature(&DexType::UniswapV3)
992            .unwrap();
993
994        registered_manager.subscribe_fee_protocol_updates(DexType::UniswapV3, pool_address);
995        registered_manager.subscribe_fee_protocol_collects(DexType::UniswapV3, pool_address);
996
997        let active_signatures =
998            registered_manager.get_active_subscribed_dex_event_signatures(&DexType::UniswapV3);
999        let addresses =
1000            registered_manager.get_subscribed_dex_contract_addresses(&DexType::UniswapV3);
1001
1002        assert_eq!(addresses.len(), 1);
1003        assert!(addresses.contains(&pool_address));
1004        assert!(active_signatures.contains(&update_sig));
1005        assert!(active_signatures.contains(&collect_sig));
1006    }
1007
1008    #[rstest]
1009    fn test_multiple_subscriptions_same_pool(mut registered_manager: DefiDataSubscriptionManager) {
1010        let pool_address = address!("1234567890123456789012345678901234567890");
1011
1012        // Subscribe same address multiple times to same event type
1013        registered_manager.subscribe_swaps(DexType::UniswapV3, pool_address);
1014        registered_manager.subscribe_swaps(DexType::UniswapV3, pool_address);
1015
1016        // Should only appear once (HashSet behavior)
1017        let addresses =
1018            registered_manager.get_subscribed_dex_contract_addresses(&DexType::UniswapV3);
1019        assert_eq!(addresses.len(), 1);
1020    }
1021
1022    #[rstest]
1023    fn test_get_combined_addresses_from_all_events(
1024        mut registered_manager: DefiDataSubscriptionManager,
1025    ) {
1026        let pool1 = address!("1111111111111111111111111111111111111111");
1027        let pool2 = address!("2222222222222222222222222222222222222222");
1028        let pool3 = address!("3333333333333333333333333333333333333333");
1029
1030        // Subscribe different pools to different events
1031        registered_manager.subscribe_swaps(DexType::UniswapV3, pool1);
1032        registered_manager.subscribe_mints(DexType::UniswapV3, pool2);
1033        registered_manager.subscribe_burns(DexType::UniswapV3, pool3);
1034
1035        // Should get all unique addresses
1036        let addresses =
1037            registered_manager.get_subscribed_dex_contract_addresses(&DexType::UniswapV3);
1038        assert_eq!(addresses.len(), 3);
1039        assert!(addresses.contains(&pool1));
1040        assert!(addresses.contains(&pool2));
1041        assert!(addresses.contains(&pool3));
1042    }
1043
1044    #[rstest]
1045    fn test_event_signature_encoding(registered_manager: DefiDataSubscriptionManager) {
1046        // Known event signature and its expected keccak256 hash
1047        // Swap(address,address,int256,int256,uint160,uint128,int24) for UniswapV3
1048        let swap_sig = registered_manager
1049            .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
1050            .unwrap();
1051
1052        // Should be properly formatted hex string
1053        assert!(swap_sig.starts_with("0x"));
1054        assert_eq!(swap_sig.len(), 66); // 0x + 64 hex chars (32 bytes)
1055
1056        // Verify it's valid hex
1057        let hex_part = &swap_sig[2..];
1058        assert!(hex_part.chars().all(|c| c.is_ascii_hexdigit()));
1059    }
1060
1061    #[rstest]
1062    #[case(DexType::UniswapV3)]
1063    #[case(DexType::UniswapV2)]
1064    fn test_complete_subscription_workflow(#[case] dex_type: DexType) {
1065        let mut manager = DefiDataSubscriptionManager::new();
1066        let pool1 = address!("1111111111111111111111111111111111111111");
1067        let pool2 = address!("2222222222222222222222222222222222222222");
1068
1069        // Step 1: Register DEX
1070        manager.register_dex_for_subscriptions(
1071            dex_type,
1072            "Swap(address,uint256,uint256)",
1073            "Mint(address,uint256)",
1074            "Burn(address,uint256)",
1075            "Collect(address,uint256,uint256)",
1076            Some("Flash(address,address,uint256,uint256,uint256,uint256)"),
1077        );
1078
1079        // Step 2: Subscribe to events
1080        manager.subscribe_swaps(dex_type, pool1);
1081        manager.subscribe_swaps(dex_type, pool2);
1082        manager.subscribe_mints(dex_type, pool1);
1083        manager.subscribe_burns(dex_type, pool2);
1084
1085        // Step 3: Verify subscriptions
1086        let addresses = manager.get_subscribed_dex_contract_addresses(&dex_type);
1087        assert_eq!(addresses.len(), 2);
1088        assert!(addresses.contains(&pool1));
1089        assert!(addresses.contains(&pool2));
1090
1091        // Step 4: Get event signatures
1092        let signatures = manager.get_subscribed_dex_event_signatures(&dex_type);
1093        assert_eq!(signatures.len(), 5);
1094
1095        // Step 5: Unsubscribe from some events
1096        manager.unsubscribe_swaps(dex_type, pool1);
1097        manager.unsubscribe_burns(dex_type, pool2);
1098
1099        // Step 6: Verify remaining subscriptions (only pool1 mint remains)
1100        let remaining = manager.get_subscribed_dex_contract_addresses(&dex_type);
1101        assert!(remaining.contains(&pool1)); // Still has mint subscription
1102        assert!(remaining.contains(&pool2)); // Still has swap subscription
1103    }
1104
1105    #[rstest]
1106    fn test_register_with_raw_signatures() {
1107        let mut manager = DefiDataSubscriptionManager::new();
1108
1109        // Register with raw event signatures
1110        manager.register_dex_for_subscriptions(
1111            DexType::UniswapV3,
1112            "Swap(address,address,int256,int256,uint160,uint128,int24)",
1113            "Mint(address,address,int24,int24,uint128,uint256,uint256)",
1114            "Burn(address,int24,int24,uint128,uint256,uint256)",
1115            "Collect(address,address,int24,int24,uint128,uint128)",
1116            Some("Flash(address,address,uint256,uint256,uint256,uint256)"),
1117        );
1118
1119        // Known keccak256 hashes for UniswapV3 events
1120        let swap_sig = manager
1121            .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
1122            .unwrap();
1123        let mint_sig = manager
1124            .get_dex_pool_mint_event_signature(&DexType::UniswapV3)
1125            .unwrap();
1126        let burn_sig = manager
1127            .get_dex_pool_burn_event_signature(&DexType::UniswapV3)
1128            .unwrap();
1129
1130        // Verify the exact hash values
1131        assert_eq!(
1132            swap_sig,
1133            "0xc42079f94a6350d7e6235f29174924f928cc2ac818eb64fed8004e115fbcca67"
1134        );
1135        assert_eq!(
1136            mint_sig,
1137            "0x7a53080ba414158be7ec69b987b5fb7d07dee101fe85488f0853ae16239d0bde"
1138        );
1139        assert_eq!(
1140            burn_sig,
1141            "0x0c396cd989a39f4459b5fa1aed6a9a8dcdbc45908acfd67e028cd568da98982c"
1142        );
1143    }
1144
1145    #[rstest]
1146    fn test_register_with_pre_encoded_signatures() {
1147        let mut manager = DefiDataSubscriptionManager::new();
1148
1149        // Register with pre-encoded keccak256 hashes (with 0x prefix)
1150        manager.register_dex_for_subscriptions(
1151            DexType::UniswapV3,
1152            "Swap(address,address,int256,int256,uint160,uint128,int24)",
1153            "Mint(address,address,int24,int24,uint128,uint256,uint256)",
1154            "Burn(address,int24,int24,uint128,uint256,uint256)",
1155            "Collect(address,address,int24,int24,uint128,uint128)",
1156            Some("Flash(address,address,uint256,uint256,uint256,uint256)"),
1157        );
1158
1159        // Should store them unchanged (normalized to lowercase)
1160        let swap_sig = manager
1161            .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
1162            .unwrap();
1163        let mint_sig = manager
1164            .get_dex_pool_mint_event_signature(&DexType::UniswapV3)
1165            .unwrap();
1166        let burn_sig = manager
1167            .get_dex_pool_burn_event_signature(&DexType::UniswapV3)
1168            .unwrap();
1169
1170        assert_eq!(
1171            swap_sig,
1172            "0xc42079f94a6350d7e6235f29174924f928cc2ac818eb64fed8004e115fbcca67"
1173        );
1174        assert_eq!(
1175            mint_sig,
1176            "0x7a53080ba414158be7ec69b987b5fb7d07dee101fe85488f0853ae16239d0bde"
1177        );
1178        assert_eq!(
1179            burn_sig,
1180            "0x0c396cd989a39f4459b5fa1aed6a9a8dcdbc45908acfd67e028cd568da98982c"
1181        );
1182    }
1183
1184    #[rstest]
1185    fn test_register_with_pre_encoded_signatures_no_prefix() {
1186        let mut manager = DefiDataSubscriptionManager::new();
1187
1188        // Register with pre-encoded hashes without 0x prefix
1189        manager.register_dex_for_subscriptions(
1190            DexType::UniswapV3,
1191            "Swap(address,address,int256,int256,uint160,uint128,int24)",
1192            "Mint(address,address,int24,int24,uint128,uint256,uint256)",
1193            "Burn(address,int24,int24,uint128,uint256,uint256)",
1194            "Collect(address,address,int24,int24,uint128,uint128)",
1195            Some("Flash(address,address,uint256,uint256,uint256,uint256)"),
1196        );
1197
1198        // Should add 0x prefix and normalize to lowercase
1199        let swap_sig = manager
1200            .get_dex_pool_swap_event_signature(&DexType::UniswapV3)
1201            .unwrap();
1202        let mint_sig = manager
1203            .get_dex_pool_mint_event_signature(&DexType::UniswapV3)
1204            .unwrap();
1205        let burn_sig = manager
1206            .get_dex_pool_burn_event_signature(&DexType::UniswapV3)
1207            .unwrap();
1208
1209        assert_eq!(
1210            swap_sig,
1211            "0xc42079f94a6350d7e6235f29174924f928cc2ac818eb64fed8004e115fbcca67"
1212        );
1213        assert_eq!(
1214            mint_sig,
1215            "0x7a53080ba414158be7ec69b987b5fb7d07dee101fe85488f0853ae16239d0bde"
1216        );
1217        assert_eq!(
1218            burn_sig,
1219            "0x0c396cd989a39f4459b5fa1aed6a9a8dcdbc45908acfd67e028cd568da98982c"
1220        );
1221    }
1222}
nautilus_blockchain/data/subscription.rs

nautilus_blockchain/data/
subscription.rs
