Make sure background sync is correctly started for chains being listened to

linera-client/src/benchmark.rs

@@ -133,7 +133,7 @@ impl<Env: Environment> Benchmark<Env> {
         let barrier = Arc::new(Barrier::new(num_chains + 1));
 
         let chain_listener_future = chain_listener
-            .run(true) // Enabling background sync for benchmarks
+            .run()
             .await
             .map_err(|_| BenchmarkError::ChainListenerStartupError)?;
         let chain_listener_handle = tokio::spawn(chain_listener_future.in_current_span());

linera-client/src/chain_listener.rs

@@ -164,6 +164,8 @@ struct ListeningClient<C: ClientContext> {
     timeout: Timestamp,
     /// The mode of listening to this chain.
     listening_mode: ListeningMode,
+    /// The cancellation token for the background sync process, if started.
+    maybe_sync_cancellation_token: Option<CancellationToken>,
 }
 
 impl<C: ClientContext> ListeningClient<C> {
@@ -173,6 +175,7 @@ impl<C: ClientContext> ListeningClient<C> {
         join_handle: NonBlockingFuture<()>,
         notification_stream: NotificationStream,
         listening_mode: ListeningMode,
+        maybe_sync_cancellation_token: Option<CancellationToken>,
     ) -> Self {
         Self {
             client,
@@ -182,12 +185,16 @@ impl<C: ClientContext> ListeningClient<C> {
             notification_stream: Arc::new(Mutex::new(notification_stream)),
             timeout: Timestamp::from(u64::MAX),
             listening_mode,
+            maybe_sync_cancellation_token,
         }
     }
 
     async fn stop(self) {
         // TODO(#4965): this is unnecessary: the join handle now also acts as an abort handle
         drop(self.abort_handle);
+        if let Some(cancellation_token) = self.maybe_sync_cancellation_token {
+            cancellation_token.cancel();
+        }
         self.join_handle.await;
     }
 }
@@ -213,6 +220,8 @@ pub struct ChainListener<C: ClientContext> {
     cancellation_token: CancellationToken,
     /// The channel through which the listener can receive commands.
     command_receiver: UnboundedReceiver<ListenerCommand>,
+    /// Whether to fully sync chains in the background.
+    enable_background_sync: bool,
 }
 
 impl<C: ClientContext + 'static> ChainListener<C> {
@@ -223,6 +232,7 @@ impl<C: ClientContext + 'static> ChainListener<C> {
         storage: <C::Environment as Environment>::Storage,
         cancellation_token: CancellationToken,
         command_receiver: UnboundedReceiver<ListenerCommand>,
+        enable_background_sync: bool,
     ) -> Self {
         Self {
             storage,
@@ -232,15 +242,13 @@ impl<C: ClientContext + 'static> ChainListener<C> {
             event_subscribers: Default::default(),
             cancellation_token,
             command_receiver,
+            enable_background_sync,
         }
     }
 
     /// Runs the chain listener.
     #[instrument(skip(self))]
-    pub async fn run(
-        mut self,
-        enable_background_sync: bool,
-    ) -> Result<impl Future<Output = Result<(), Error>>, Error> {
+    pub async fn run(mut self) -> Result<impl Future<Output = Result<(), Error>>, Error> {
         let chain_ids = {
             let guard = self.context.lock().await;
             let admin_chain_id = guard.admin_chain();
@@ -278,33 +286,6 @@ impl<C: ClientContext + 'static> ChainListener<C> {
             chain_ids
         };
 
-        // Start background tasks to sync received certificates for each chain,
-        // if enabled.
-        if enable_background_sync {
-            let context = Arc::clone(&self.context);
-            let cancellation_token = self.cancellation_token.clone();
-            for (chain_id, mode) in chain_ids.iter() {
-                if mode != &ListeningMode::FullChain {
-                    continue;
-                }
-                let context = Arc::clone(&context);
-                let cancellation_token = cancellation_token.clone();
-                let chain_id = *chain_id;
-                linera_base::task::spawn(async move {
-                    if let Err(e) = Self::background_sync_received_certificates(
-                        context,
-                        chain_id,
-                        cancellation_token,
-                    )
-                    .await
-                    {
-                        warn!("Background sync failed for chain {chain_id}: {e}");
-                    }
-                })
-                .forget();
-            }
-        }
-
         Ok(async {
             self.listen_recursively(chain_ids).await?;
             loop {
@@ -495,6 +476,10 @@ impl<C: ClientContext + 'static> ChainListener<C> {
                 .get(&chain_id)
                 .map(|existing_client| existing_client.listening_mode.clone()),
         );
+        // Start background tasks to sync received certificates, if enabled.
+        let maybe_sync_cancellation_token = self
+            .start_background_sync(chain_id, listening_mode.clone())
+            .await;
         let client = self
             .context
             .lock()
@@ -510,13 +495,43 @@ impl<C: ClientContext + 'static> ChainListener<C> {
             join_handle,
             notification_stream,
             listening_mode,
+            maybe_sync_cancellation_token,
         );
         self.listening.insert(chain_id, listening_client);
         let publishing_chains = self.update_event_subscriptions(chain_id).await?;
         self.maybe_process_inbox(chain_id).await?;
         Ok(publishing_chains)
     }
 
+    async fn start_background_sync(
+        &mut self,
+        chain_id: ChainId,
+        mode: ListeningMode,
+    ) -> Option<CancellationToken> {
+        if !self.enable_background_sync {
+            return None;
+        }
+        if mode != ListeningMode::FullChain {
+            return None;
+        }
+        let context = Arc::clone(&self.context);
+        let cancellation_token = CancellationToken::new();
+        let cancellation_token_clone = cancellation_token.clone();
+        linera_base::task::spawn(async move {
+            if let Err(e) = Self::background_sync_received_certificates(
+                context,
+                chain_id,
+                cancellation_token_clone,
+            )
+            .await
+            {
+                warn!("Background sync failed for chain {chain_id}: {e}");
+            }
+        })
+        .forget();
+        Some(cancellation_token)
+    }
+
     /// Updates the event subscribers map, and returns all publishing chains we need to listen to.
     async fn update_event_subscriptions(
         &mut self,

linera-client/src/unit_tests/chain_listener.rs

@@ -159,8 +159,9 @@ async fn test_chain_listener() -> anyhow::Result<()> {
         storage,
         child_token,
         tokio::sync::mpsc::unbounded_channel().1,
+        false, // Unit test doesn't need background sync
     )
-    .run(false) // Unit test doesn't need background sync
+    .run()
     .await
     .unwrap();
 
@@ -273,8 +274,9 @@ async fn test_chain_listener_follow_only() -> anyhow::Result<()> {
         storage.clone(),
         child_token,
         tokio::sync::mpsc::unbounded_channel().1,
+        false, // Unit test doesn't need background sync
     )
-    .run(false) // Unit test doesn't need background sync
+    .run()
     .await
     .unwrap();
 
@@ -400,8 +402,9 @@ async fn test_chain_listener_admin_chain() -> anyhow::Result<()> {
         storage.clone(),
         child_token,
         tokio::sync::mpsc::unbounded_channel().1,
+        false, // Unit test doesn't need background sync
     )
-    .run(false) // Unit test doesn't need background sync
+    .run()
     .await
     .unwrap();
 

linera-faucet/server/src/lib.rs

@@ -963,8 +963,9 @@ where
             self.storage,
             cancellation_token.clone(),
             tokio::sync::mpsc::unbounded_channel().1,
+            false, // Faucet doesn't receive messages, so no need for background sync
         )
-        .run(false) // Faucet doesn't receive messages, so no need for background sync
+        .run()
         .await?;
         let batch_processor_task = batch_processor.run(cancellation_token.clone());
         let tcp_listener =

linera-service/src/cli/main.rs

@@ -810,6 +810,7 @@ impl Runnable for Job {
                         storage.clone(),
                         shutdown_notifier.clone(),
                         mpsc::unbounded_channel().1,
+                        true, // Enabling background sync for benchmarks
                     );
                     linera_client::benchmark::Benchmark::run_benchmark(
                         bps,

linera-service/src/node_service.rs

@@ -983,8 +983,9 @@ where
             storage,
             cancellation_token.clone(),
             command_receiver,
+            true,
         )
-        .run(true)
+        .run()
         .await?;
         let mut chain_listener = Box::pin(chain_listener).fuse();
         let tcp_listener =

web/@linera/client/src/lib.rs

@@ -103,8 +103,9 @@ impl Client {
             storage,
             tokio_util::sync::CancellationToken::new(),
             tokio::sync::mpsc::unbounded_channel().1,
+            true, // Enable background sync
         )
-        .run(true) // Enable background sync
+        .run()
         .boxed_local()
         .await?
         .boxed_local();