updated slot0

Author: jnoorchashm37

crates/uni-v4-common/src/shared_pools.rs

@@ -152,6 +152,7 @@ impl UniswapPools {
                         continue;
                     };
 
+                    tracing::trace!(?update, "updating slot0 for pool");
                     pool.value_mut().update_slot0(
                         update.tick,
                         update.sqrt_price_x96.into(),

crates/uni-v4-upkeeper/src/lib.rs

@@ -9,7 +9,8 @@ pub mod pool_manager_service;
 pub mod pool_manager_service_builder;
 pub mod pool_providers;
 pub mod pool_registry;
-pub mod slot0;
+pub mod slot0_pool_updater;
+pub mod slot0_stream;
 
 fn i128_to_i256(value: i128) -> I256 {
     I256::try_from(value).unwrap()
@@ -20,5 +21,5 @@ pub enum ConversionError {
     #[error("overflow from i32 to i24 {0:?}")]
     OverflowErrorI24(i32),
     #[error("overflow from I256 to I128 {0:?}")]
-    OverflowErrorI28(I256)
+    OverflowErrorI28(I256),
 }

crates/uni-v4-upkeeper/src/pool_manager_service.rs

@@ -2,30 +2,34 @@ use std::{
     collections::HashSet,
     pin::Pin,
     sync::Arc,
-    task::{Context, Poll}
+    task::{Context, Poll},
 };
 
 use alloy::{primitives::Address, providers::Provider};
 use futures::{Future, Stream, StreamExt};
 use thiserror::Error;
-use tokio::sync::mpsc;
+use tokio::sync::mpsc::{self, UnboundedSender};
 use uni_v4_common::{PoolId, PoolKey, PoolUpdate, UniswapPools};
 use uni_v4_structure::BaselinePoolState;
 
 use super::{
     baseline_pool_factory::{BaselinePoolFactory, BaselinePoolFactoryError, UpdateMessage},
-    fetch_pool_keys::set_controller_address
+    fetch_pool_keys::set_controller_address,
+};
+use crate::{
+    pool_providers::PoolEventStream,
+    slot0_pool_updater::{Slot0PoolUpdater, Slot0UpgradeCommand},
+    slot0_stream::Slot0Stream,
 };
-use crate::{pool_providers::PoolEventStream, slot0::Slot0Stream};
 
 /// Pool information combining BaselinePoolState with token metadata
 #[derive(Debug, Clone)]
 pub struct PoolInfo {
-    pub baseline_state:  BaselinePoolState,
-    pub token0:          Address,
-    pub token1:          Address,
+    pub baseline_state: BaselinePoolState,
+    pub token0: Address,
+    pub token1: Address,
     pub token0_decimals: u8,
-    pub token1_decimals: u8
+    pub token1_decimals: u8,
 }
 
 #[derive(Error, Debug)]
@@ -37,39 +41,38 @@ pub enum PoolManagerServiceError {
     #[error("Pool factory error: {0}")]
     PoolFactory(String),
     #[error("Baseline pool factory error: {0}")]
-    BaselineFactory(#[from] BaselinePoolFactoryError)
+    BaselineFactory(#[from] BaselinePoolFactoryError),
 }
 
 /// Service for managing Uniswap V4 pools with real-time block subscription
 /// updates
-pub struct PoolManagerService<P, Event, S = ()>
+pub struct PoolManagerService<P, Event>
 where
     P: Provider + Unpin + Clone + 'static,
-    Event: PoolEventStream
+    Event: PoolEventStream,
 {
-    pub(crate) factory:            BaselinePoolFactory<P>,
-    pub(crate) event_stream:       Event,
-    pub(crate) pools:              UniswapPools,
-    pub(crate) current_block:      u64,
+    pub(crate) factory: BaselinePoolFactory<P>,
+    pub(crate) event_stream: Event,
+    pub(crate) pools: UniswapPools,
+    pub(crate) current_block: u64,
     pub(crate) auto_pool_creation: bool,
-    pub(crate) slot0_stream:       Option<S>,
+    pub(crate) slot0_stream_updater_tx: Option<UnboundedSender<Slot0UpgradeCommand>>,
     // If we are loading more ticks at a block, we will queue up updates messages here
     // so that we don't hit any race conditions.
-    pending_updates:               Vec<PoolUpdate>,
+    pending_updates: Vec<PoolUpdate>,
     // Channel for sending updates instead of applying them directly
-    update_sender:                 Option<mpsc::Sender<PoolUpdate>>
+    update_sender: Option<mpsc::Sender<PoolUpdate>>,
 }
 
-impl<P, Event, S> PoolManagerService<P, Event, S>
+impl<P, Event> PoolManagerService<P, Event>
 where
     P: Provider + Clone + Unpin + 'static,
     Event: PoolEventStream,
     BaselinePoolFactory<P>: Stream<Item = UpdateMessage> + Unpin,
-    S: Slot0Stream
 {
     /// Create a new PoolManagerService
     #[allow(clippy::too_many_arguments)]
-    pub async fn new(
+    pub async fn new<S: Slot0Stream>(
         provider: Arc<P>,
         event_stream: Event,
         angstrom_address: Address,
@@ -80,10 +83,10 @@ where
         tick_edge_threshold: Option<u16>,
         filter_pool_keys: Option<HashSet<PoolKey>>,
         auto_pool_creation: bool,
-        slot0_stream: Option<S>,
+        mut slot0_stream: Option<S>,
         current_block: Option<u64>,
         ticks_per_batch: Option<usize>,
-        update_channel: Option<mpsc::Sender<PoolUpdate>>
+        update_channel: Option<mpsc::Sender<PoolUpdate>>,
     ) -> Result<Self, PoolManagerServiceError> {
         // Set the controller address for the fetch_pool_keys module
         set_controller_address(controller_address);
@@ -105,19 +108,24 @@ where
             tick_band,
             tick_edge_threshold,
             filter_pool_keys,
-            ticks_per_batch
+            ticks_per_batch,
         )
         .await;
 
+        let (slot0_stream_updater_tx, slot0_stream_updater_rx) = slot0_stream
+            .is_some()
+            .then(tokio::sync::mpsc::unbounded_channel)
+            .unzip();
+
         let mut service = Self {
             event_stream,
             factory,
             pools: UniswapPools::new(pools, current_block),
             current_block: deploy_block,
             auto_pool_creation,
-            slot0_stream,
+            slot0_stream_updater_tx,
             pending_updates: Vec::new(),
-            update_sender: update_channel
+            update_sender: update_channel,
         };
 
         service
@@ -130,10 +138,15 @@ where
         }
 
         // Subscribe all initial pools to slot0 stream if present (using angstrom IDs)
-        if let Some(slot0_stream) = &mut service.slot0_stream {
+        if let Some(mut slot0_stream) = slot0_stream.take() {
             let angstrom_pool_ids: HashSet<PoolId> =
                 service.factory.registry().public_keys().collect();
             slot0_stream.subscribe_pools(angstrom_pool_ids);
+            Slot0PoolUpdater::<S>::spawn_new(
+                slot0_stream,
+                service.pools.clone(),
+                slot0_stream_updater_rx.unwrap(),
+            );
         }
 
         // Send all initialized pools through the channel on startup
@@ -144,7 +157,7 @@ where
                 .iter()
                 .map(|entry| PoolUpdate::NewPoolState {
                     pool_id: *entry.key(),
-                    state:   entry.value().clone()
+                    state: entry.value().clone(),
                 })
                 .collect();
 
@@ -178,14 +191,14 @@ where
         block_number: u64,
         bundle_fee: u32,
         swap_fee: u32,
-        protocol_fee: u32
+        protocol_fee: u32,
     ) {
         self.factory.queue_pool_creation(
             pool_key,
             block_number,
             bundle_fee,
             swap_fee,
-            protocol_fee
+            protocol_fee,
         );
     }
 
@@ -243,7 +256,7 @@ where
                 swap_fee,
                 protocol_fee,
                 tick_spacing,
-                block
+                block,
             } => {
                 if self.auto_pool_creation {
                     // Reconstruct pool_key from the NewPool data
@@ -254,7 +267,7 @@ where
                             *block,
                             *bundle_fee,
                             *swap_fee,
-                            *protocol_fee
+                            *protocol_fee,
                         );
 
                         tracing::info!(
@@ -282,8 +295,9 @@ where
                 self.factory.remove_pool_by_id(*pool_id);
 
                 // Unsubscribe pool from slot0 stream (pool_id here is already angstrom ID)
-                if let Some(slot0_stream) = &mut self.slot0_stream {
-                    slot0_stream.unsubscribe_pools(HashSet::from([*pool_id]));
+                if let Some(slot0_stream_tx) = self.slot0_stream_updater_tx.as_ref() {
+                    let _ = slot0_stream_tx
+                        .send(Slot0UpgradeCommand::UnsubscribePools(HashSet::from([*pool_id])));
                 }
             }
             PoolUpdate::FeeUpdate { pool_id, bundle_fee, swap_fee, protocol_fee, .. } => {
@@ -316,11 +330,14 @@ where
                 self.event_stream.start_tracking_pool(*pool_id);
 
                 // Subscribe new pool to slot0 stream (using angstrom ID)
-                if let Some(slot0_stream) = &mut self.slot0_stream
+                if let Some(slot0_stream_tx) = self.slot0_stream_updater_tx.as_ref()
                     && let Some(angstrom_pool_id) =
                         self.factory.registry().public_key_from_private(pool_id)
                 {
-                    slot0_stream.subscribe_pools(HashSet::from([angstrom_pool_id]));
+                    let _ =
+                        slot0_stream_tx.send(Slot0UpgradeCommand::SubscribePools(HashSet::from([
+                            angstrom_pool_id,
+                        ])));
                 }
 
                 tracing::info!("Tracking new pool from factory: {:?}", pool_id);
@@ -334,12 +351,11 @@ where
     }
 }
 
-impl<P, Event, S> Future for PoolManagerService<P, Event, S>
+impl<P, Event> Future for PoolManagerService<P, Event>
 where
     P: Provider + Clone + Unpin + 'static,
     Event: PoolEventStream,
     BaselinePoolFactory<P>: Stream<Item = UpdateMessage> + Unpin,
-    S: Slot0Stream
 {
     type Output = ();
 
@@ -386,7 +402,7 @@ where
                     this.factory.check_and_request_ticks_if_needed(
                         *entry.key(),
                         entry.value(),
-                        Some(this.current_block)
+                        Some(this.current_block),
                     );
                 }
             }
@@ -415,7 +431,7 @@ where
                         this.factory.check_and_request_ticks_if_needed(
                             *entry.key(),
                             entry.value(),
-                            Some(this.current_block)
+                            Some(this.current_block),
                         );
                     }
                 }
@@ -424,16 +440,16 @@ where
             }
         }
 
-        if let Some(slot0_stream) = this.slot0_stream.as_mut() {
-            let mut slot0_updates = Vec::new();
-            while let Poll::Ready(Some(update)) = slot0_stream.poll_next_unpin(cx) {
-                slot0_updates.push(update);
-            }
-            for update in slot0_updates {
-                let pool_update = PoolUpdate::Slot0Update(update);
-                this.dispatch_update(pool_update);
-            }
-        }
+        // if let Some(slot0_stream) = this.slot0_stream.as_mut() {
+        //     let mut slot0_updates = Vec::new();
+        //     while let Poll::Ready(Some(update)) = slot0_stream.poll_next_unpin(cx) {
+        //         slot0_updates.push(update);
+        //     }
+        //     for update in slot0_updates {
+        //         let pool_update = PoolUpdate::Slot0Update(update);
+        //         this.dispatch_update(pool_update);
+        //     }
+        // }
 
         Poll::Pending
     }