Fix attempt

Author: squadgazzz

.github/workflows/deploy-profiling.yaml

@@ -9,6 +9,7 @@ on:
   push:
     branches:
       - jemalloc-profiling-support
+      - ilya/mem-leak-experiment
 
 jobs:
   deploy-profiling:

crates/driver/src/infra/tokens.rs

@@ -39,10 +39,7 @@ impl Fetcher {
             cache: RwLock::new(HashMap::new()),
             requests: BoxRequestSharing::labelled("token_info".into()),
         });
-        tokio::task::spawn(
-            update_task(block_stream, Arc::downgrade(&inner))
-                .instrument(tracing::info_span!("token_fetcher")),
-        );
+        tokio::task::spawn(update_task(block_stream, Arc::downgrade(&inner)));
         Self(inner)
     }
 
@@ -61,14 +58,23 @@ impl Fetcher {
 /// fetcher is dropped.
 async fn update_task(blocks: CurrentBlockWatcher, inner: std::sync::Weak<Inner>) {
     let mut stream = block_stream::into_stream(blocks);
-    while stream.next().await.is_some() {
-        let inner = match inner.upgrade() {
-            Some(inner) => inner,
-            // Fetcher was dropped, stop update task.
-            None => break,
-        };
-        if let Err(err) = update_balances(inner).await {
-            tracing::warn!(?err, "error updating token cache");
+    while let Some(block) = stream.next().await {
+        let span = tracing::info_span!("token_fetcher", block = block.number);
+        let should_break = async {
+            let inner = match inner.upgrade() {
+                Some(inner) => inner,
+                // Fetcher was dropped, stop update task.
+                None => return true,
+            };
+            if let Err(err) = update_balances(inner).await {
+                tracing::warn!(?err, "error updating token cache");
+            }
+            false
+        }
+        .instrument(span)
+        .await;
+        if should_break {
+            break;
         }
     }
 }

crates/ethrpc/src/block_stream/mod.rs

@@ -172,7 +172,7 @@ pub async fn current_block_ws_stream(
     let first_block = BlockInfo::try_from(first_block).context("failed to parse initial block")?;
 
     let (sender, receiver) = watch::channel(first_block);
-    let update_future = async move {
+    tokio::task::spawn(async move {
         // Keep WebSocket provider alive to maintain connection
         let _ws_provider = ws_provider;
         let mut previous_block = first_block;
@@ -188,6 +188,9 @@ pub async fn current_block_ws_stream(
                 }
             };
 
+            let _span =
+                tracing::info_span!("current_block_stream", block = block_info.number).entered();
+
             update_current_block_metrics(block_info.number);
 
             // If the block is exactly the same as the previous one, ignore it.
@@ -217,9 +220,7 @@ pub async fn current_block_ws_stream(
 
         // If we reach here, the stream ended permanently
         tracing::error!("block stream ended after max reconnection attempts");
-    };
-
-    tokio::task::spawn(update_future.instrument(tracing::info_span!("current_block_stream")));
+    });
     Ok(receiver)
 }
 

crates/shared/src/account_balances/cached.rs

@@ -116,42 +116,47 @@ impl Balances {
         let cache = self.balance_cache.clone();
         let mut stream = into_stream(block_stream);
 
-        let task = async move {
+        tokio::spawn(async move {
             while let Some(block) = stream.next().await {
-                let balances_to_update = {
+                let span = tracing::info_span!("balance_cache", block = block.number);
+                async {
+                    let balances_to_update = {
+                        let mut cache = cache.lock().unwrap();
+                        cache.last_seen_block = block.number;
+                        cache
+                            .data
+                            .iter()
+                            .filter_map(|(query, entry)| {
+                                // Only update balances that have been requested recently.
+                                let oldest_allowed_request =
+                                    cache.last_seen_block.saturating_sub(EVICTION_TIME);
+                                (entry.requested_at >= oldest_allowed_request)
+                                    .then_some(query.clone())
+                            })
+                            .collect_vec()
+                    };
+
+                    let results = inner.get_balances(&balances_to_update).await;
+
                     let mut cache = cache.lock().unwrap();
-                    cache.last_seen_block = block.number;
-                    cache
-                        .data
-                        .iter()
-                        .filter_map(|(query, entry)| {
-                            // Only update balances that have been requested recently.
-                            let oldest_allowed_request =
-                                cache.last_seen_block.saturating_sub(EVICTION_TIME);
-                            (entry.requested_at >= oldest_allowed_request).then_some(query.clone())
-                        })
-                        .collect_vec()
-                };
-
-                let results = inner.get_balances(&balances_to_update).await;
-
-                let mut cache = cache.lock().unwrap();
-                balances_to_update
-                    .into_iter()
-                    .zip(results)
-                    .for_each(|(query, result)| {
-                        if let Ok(balance) = result {
-                            cache.update_balance(&query, balance, block.number);
-                        }
+                    balances_to_update
+                        .into_iter()
+                        .zip(results)
+                        .for_each(|(query, result)| {
+                            if let Ok(balance) = result {
+                                cache.update_balance(&query, balance, block.number);
+                            }
+                        });
+                    cache.data.retain(|_, value| {
+                        // Only keep balances where we know we have the most recent data.
+                        value.updated_at >= block.number
                     });
-                cache.data.retain(|_, value| {
-                    // Only keep balances where we know we have the most recent data.
-                    value.updated_at >= block.number
-                });
+                }
+                .instrument(span)
+                .await;
             }
             tracing::error!("block stream terminated unexpectedly");
-        };
-        tokio::spawn(task.instrument(tracing::info_span!("balance_cache")));
+        });
     }
 }
 

crates/shared/src/recent_block_cache.rs

@@ -200,21 +200,28 @@ where
         block_stream: CurrentBlockWatcher,
         label: String,
     ) {
-        tokio::task::spawn(
-            async move {
-                let mut stream = ethrpc::block_stream::into_stream(block_stream);
-                while let Some(block) = stream.next().await {
+        tokio::task::spawn(async move {
+            let mut stream = ethrpc::block_stream::into_stream(block_stream);
+            while let Some(block) = stream.next().await {
+                let span =
+                    tracing::info_span!("cache_maintenance", cache = %label, block = block.number);
+                let should_break = async {
                     let Some(inner) = inner.upgrade() else {
                         tracing::debug!("cache no longer in use; terminate GC task");
-                        break;
+                        return true;
                     };
                     if let Err(err) = inner.update_cache_at_block(block.number).await {
                         tracing::warn!(?err, "failed to update cache");
                     }
+                    false
+                }
+                .instrument(span)
+                .await;
+                if should_break {
+                    break;
                 }
             }
-            .instrument(tracing::info_span!("cache_maintenance", cache = label)),
-        );
+        });
     }
 }