fix simulation

Signed-off-by: Bugen Zhao <[email protected]>
risingwavelabs · Jul 3, 2024 · a931b1e · a931b1e
1 parent 64d8560
commit a931b1e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/tests/simulation/src/cluster.rs b/src/tests/simulation/src/cluster.rs
@@ -32,6 +32,7 @@ use itertools::Itertools;
 use madsim::runtime::{Handle, NodeHandle};
 use rand::seq::IteratorRandom;
 use rand::Rng;
+use risingwave_common::util::tokio_util::sync::CancellationToken;
 #[cfg(madsim)]
 use risingwave_object_store::object::sim::SimServer as ObjectStoreSimServer;
 use risingwave_pb::common::WorkerNode;
@@ -500,7 +501,7 @@ impl Cluster {
                 .name(format!("compute-{i}"))
                 .ip([192, 168, 3, i as u8].into())
                 .cores(conf.compute_node_cores)
-                .init(move || risingwave_compute::start(opts.clone()))
+                .init(move || risingwave_compute::start(opts.clone(), CancellationToken::new()))
                 .build();
         }
 

diff --git a/src/tests/simulation/tests/integration_tests/batch/mod.rs b/src/tests/simulation/tests/integration_tests/batch/mod.rs
@@ -18,6 +18,7 @@ use std::io::Write;
 
 use clap::Parser;
 use itertools::Itertools;
+use risingwave_common::util::tokio_util::sync::CancellationToken;
 use risingwave_simulation::cluster::{Cluster, ConfigPath, Configuration, Session};
 use tokio::time::Duration;
 
@@ -44,7 +45,7 @@ fn create_compute_node(cluster: &Cluster, idx: usize, role: &str) {
         .name(format!("compute-{idx}"))
         .ip([192, 168, 3, idx as u8].into())
         .cores(config.compute_node_cores)
-        .init(move || risingwave_compute::start(opts.clone()))
+        .init(move || risingwave_compute::start(opts.clone(), CancellationToken::new()))
         .build();
 }