From a931b1ef002f84ae65ec131b8a66880afdba96aa Mon Sep 17 00:00:00 2001
From: Bugen Zhao <i@bugenzhao.com>
Date: Wed, 3 Jul 2024 10:26:48 +0800
Subject: [PATCH] fix simulation

Signed-off-by: Bugen Zhao <i@bugenzhao.com>
---
 src/tests/simulation/src/cluster.rs                       | 3 ++-
 src/tests/simulation/tests/integration_tests/batch/mod.rs | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/tests/simulation/src/cluster.rs b/src/tests/simulation/src/cluster.rs
index 8caf8b52931ca..351e45279fd5e 100644
--- 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
index 25c690cfdcbf8..1ee5132884a2a 100644
--- 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();
 }