risingwave_simulation/
cluster.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#![cfg_attr(not(madsim), allow(unused_imports))]
16
17use std::cmp::max;
18use std::collections::HashMap;
19use std::future::Future;
20use std::io::Write;
21use std::path::{Path, PathBuf};
22use std::sync::Arc;
23use std::time::Duration;
24
25use anyhow::{Result, anyhow, bail};
26use cfg_or_panic::cfg_or_panic;
27use clap::Parser;
28use futures::channel::{mpsc, oneshot};
29use futures::future::join_all;
30use futures::{SinkExt, StreamExt};
31use itertools::Itertools;
32#[cfg(madsim)]
33use madsim::runtime::{Handle, NodeHandle};
34use rand::RngExt as _;
35use rand::seq::IteratorRandom;
36use risingwave_common::util::tokio_util::sync::CancellationToken;
37use risingwave_common::util::worker_util::DEFAULT_RESOURCE_GROUP;
38#[cfg(madsim)]
39use risingwave_object_store::object::sim::SimServer as ObjectStoreSimServer;
40use risingwave_pb::common::WorkerNode;
41use sqllogictest::AsyncDB;
42use tempfile::NamedTempFile;
43#[cfg(not(madsim))]
44use tokio::runtime::Handle;
45use uuid::Uuid;
46
47use crate::client::RisingWave;
48
49/// The path to the configuration file for the cluster.
50#[derive(Clone, Debug)]
51pub enum ConfigPath {
52    /// A regular path pointing to a external configuration file.
53    Regular(String),
54    /// A temporary path pointing to a configuration file created at runtime.
55    Temp(Arc<tempfile::TempPath>),
56}
57
58impl ConfigPath {
59    pub fn as_str(&self) -> &str {
60        match self {
61            ConfigPath::Regular(s) => s,
62            ConfigPath::Temp(p) => p.as_os_str().to_str().unwrap(),
63        }
64    }
65}
66
67/// RisingWave cluster configuration.
68#[derive(Debug, Clone)]
69pub struct Configuration {
70    /// The path to configuration file.
71    ///
72    /// Empty string means using the default config.
73    pub config_path: ConfigPath,
74
75    /// The number of frontend nodes.
76    pub frontend_nodes: usize,
77
78    /// The number of compute nodes.
79    pub compute_nodes: usize,
80
81    /// The number of meta nodes.
82    pub meta_nodes: usize,
83
84    /// The number of compactor nodes.
85    pub compactor_nodes: usize,
86
87    /// The number of CPU cores for each compute node.
88    ///
89    /// This determines `worker_node_parallelism`.
90    pub compute_node_cores: usize,
91
92    /// Queries to run per session.
93    pub per_session_queries: Arc<Vec<String>>,
94
95    /// Resource groups for compute nodes.
96    pub compute_resource_groups: HashMap<usize, String>,
97
98    /// Roles for compute nodes (1-indexed). If not set, defaults to "both".
99    /// Values: "serving", "streaming", "both".
100    pub compute_node_roles: HashMap<usize, String>,
101}
102
103impl Default for Configuration {
104    fn default() -> Self {
105        let config_path = {
106            let mut file =
107                tempfile::NamedTempFile::new().expect("failed to create temp config file");
108
109            let config_data = r#"
110[server]
111telemetry_enabled = false
112metrics_level = "Disabled"
113"#
114            .to_owned();
115            file.write_all(config_data.as_bytes())
116                .expect("failed to write config file");
117            file.into_temp_path()
118        };
119
120        Configuration {
121            config_path: ConfigPath::Temp(config_path.into()),
122            frontend_nodes: 1,
123            compute_nodes: 1,
124            meta_nodes: 1,
125            compactor_nodes: 1,
126            compute_node_cores: 1,
127            per_session_queries: vec![].into(),
128            compute_resource_groups: Default::default(),
129            compute_node_roles: Default::default(),
130        }
131    }
132}
133
134impl Configuration {
135    /// Returns the configuration for scale test.
136    pub fn for_scale() -> Self {
137        // Embed the config file and create a temporary file at runtime. The file will be deleted
138        // automatically when it's dropped.
139        let config_path = {
140            let mut file =
141                tempfile::NamedTempFile::new().expect("failed to create temp config file");
142            file.write_all(include_bytes!("risingwave-scale.toml"))
143                .expect("failed to write config file");
144            file.into_temp_path()
145        };
146
147        Configuration {
148            config_path: ConfigPath::Temp(config_path.into()),
149            frontend_nodes: 2,
150            compute_nodes: 3,
151            meta_nodes: 1,
152            compactor_nodes: 2,
153            compute_node_cores: 2,
154            per_session_queries: vec![].into(),
155            ..Default::default()
156        }
157    }
158
159    /// Provides a configuration for scale test which ensures that the arrangement backfill is disabled,
160    /// so table scan will use `no_shuffle`.
161    pub fn for_scale_no_shuffle() -> Self {
162        let mut conf = Self::for_scale();
163        conf.per_session_queries = vec![
164            "SET STREAMING_USE_ARRANGEMENT_BACKFILL = false;".into(),
165            "SET STREAMING_USE_SNAPSHOT_BACKFILL = false;".into(),
166        ]
167        .into();
168        conf
169    }
170
171    pub fn for_scale_shared_source() -> Self {
172        let mut conf = Self::for_scale();
173        conf.per_session_queries = vec!["SET STREAMING_USE_SHARED_SOURCE = true;".into()].into();
174        conf
175    }
176
177    pub fn for_auto_parallelism(
178        max_heartbeat_interval_secs: u64,
179        enable_auto_parallelism: bool,
180    ) -> Self {
181        let disable_automatic_parallelism_control = !enable_auto_parallelism;
182
183        let config_path = {
184            let mut file =
185                tempfile::NamedTempFile::new().expect("failed to create temp config file");
186
187            let config_data = format!(
188                r#"[meta]
189max_heartbeat_interval_secs = {max_heartbeat_interval_secs}
190disable_automatic_parallelism_control = {disable_automatic_parallelism_control}
191parallelism_control_trigger_first_delay_sec = 0
192parallelism_control_batch_size = 10
193parallelism_control_trigger_period_sec = 10
194
195[system]
196barrier_interval_ms = 250
197checkpoint_frequency = 4
198
199[server]
200telemetry_enabled = false
201metrics_level = "Disabled"
202"#
203            );
204            file.write_all(config_data.as_bytes())
205                .expect("failed to write config file");
206            file.into_temp_path()
207        };
208
209        Configuration {
210            config_path: ConfigPath::Temp(config_path.into()),
211            frontend_nodes: 1,
212            compute_nodes: 3,
213            meta_nodes: 1,
214            compactor_nodes: 1,
215            compute_node_cores: 2,
216            per_session_queries: vec![
217                "create view if not exists table_parallelism as select t.name, tf.parallelism from rw_tables t, rw_table_fragments tf where t.id = tf.table_id;".into(),
218                "create view if not exists mview_parallelism as select m.name, tf.parallelism from rw_materialized_views m, rw_table_fragments tf where m.id = tf.table_id;".into(),
219            ]
220                .into(),
221            ..Default::default()
222        }
223    }
224
225    pub fn for_default_parallelism(default_parallelism: usize) -> Self {
226        let config_path = {
227            let mut file =
228                tempfile::NamedTempFile::new().expect("failed to create temp config file");
229
230            let config_data = format!(
231                r#"
232[server]
233telemetry_enabled = false
234metrics_level = "Disabled"
235[meta]
236default_parallelism = {default_parallelism}
237"#
238            );
239            file.write_all(config_data.as_bytes())
240                .expect("failed to write config file");
241            file.into_temp_path()
242        };
243
244        Configuration {
245            config_path: ConfigPath::Temp(config_path.into()),
246            frontend_nodes: 1,
247            compute_nodes: 1,
248            meta_nodes: 1,
249            compactor_nodes: 1,
250            compute_node_cores: default_parallelism * 2,
251            per_session_queries: vec![].into(),
252            compute_resource_groups: Default::default(),
253            compute_node_roles: Default::default(),
254        }
255    }
256
257    /// Returns the config for backfill test.
258    pub fn for_backfill() -> Self {
259        // Embed the config file and create a temporary file at runtime. The file will be deleted
260        // automatically when it's dropped.
261        let config_path = {
262            let mut file =
263                tempfile::NamedTempFile::new().expect("failed to create temp config file");
264            file.write_all(include_bytes!("backfill.toml"))
265                .expect("failed to write config file");
266            file.into_temp_path()
267        };
268
269        Configuration {
270            config_path: ConfigPath::Temp(config_path.into()),
271            frontend_nodes: 1,
272            compute_nodes: 1,
273            meta_nodes: 1,
274            compactor_nodes: 1,
275            compute_node_cores: 4,
276            ..Default::default()
277        }
278    }
279
280    pub fn for_arrangement_backfill() -> Self {
281        // Embed the config file and create a temporary file at runtime. The file will be deleted
282        // automatically when it's dropped.
283        let config_path = {
284            let mut file =
285                tempfile::NamedTempFile::new().expect("failed to create temp config file");
286            file.write_all(include_bytes!("arrangement_backfill.toml"))
287                .expect("failed to write config file");
288            file.into_temp_path()
289        };
290
291        Configuration {
292            config_path: ConfigPath::Temp(config_path.into()),
293            frontend_nodes: 1,
294            compute_nodes: 3,
295            meta_nodes: 1,
296            compactor_nodes: 1,
297            compute_node_cores: 1,
298            per_session_queries: vec![
299                "SET STREAMING_USE_ARRANGEMENT_BACKFILL = true;".into(),
300                "SET STREAMING_USE_SNAPSHOT_BACKFILL = false;".into(),
301            ]
302            .into(),
303            ..Default::default()
304        }
305    }
306
307    pub fn for_background_ddl() -> Self {
308        // Embed the config file and create a temporary file at runtime. The file will be deleted
309        // automatically when it's dropped.
310        let config_path = {
311            let mut file =
312                tempfile::NamedTempFile::new().expect("failed to create temp config file");
313            file.write_all(include_bytes!("background_ddl.toml"))
314                .expect("failed to write config file");
315            file.into_temp_path()
316        };
317
318        Configuration {
319            config_path: ConfigPath::Temp(config_path.into()),
320            // NOTE(kwannoel): The cancel test depends on `processlist`,
321            // which will cancel a stream job within the process.
322            // so we cannot have multiple frontend node, since a new session spawned
323            // to cancel the job could be routed to a different frontend node,
324            // in a different process.
325            frontend_nodes: 1,
326            compute_nodes: 3,
327            meta_nodes: 1,
328            compactor_nodes: 2,
329            compute_node_cores: 2,
330            ..Default::default()
331        }
332    }
333
334    pub fn enable_arrangement_backfill() -> Self {
335        let config_path = {
336            let mut file =
337                tempfile::NamedTempFile::new().expect("failed to create temp config file");
338            file.write_all(include_bytes!("disable_arrangement_backfill.toml"))
339                .expect("failed to write config file");
340            file.into_temp_path()
341        };
342        Configuration {
343            config_path: ConfigPath::Temp(config_path.into()),
344            frontend_nodes: 1,
345            compute_nodes: 1,
346            meta_nodes: 1,
347            compactor_nodes: 1,
348            compute_node_cores: 1,
349            per_session_queries: vec![].into(),
350            ..Default::default()
351        }
352    }
353
354    /// Returns the total number of cores for streaming compute nodes.
355    pub fn total_streaming_cores(&self) -> u32 {
356        (self.compute_nodes * self.compute_node_cores) as u32
357    }
358}
359
360/// A risingwave cluster.
361///
362/// # Nodes
363///
364/// | Name             | IP            |
365/// | ---------------- | ------------- |
366/// | meta-x           | 192.168.1.x   |
367/// | frontend-x       | 192.168.2.x   |
368/// | compute-x        | 192.168.3.x   |
369/// | compactor-x      | 192.168.4.x   |
370/// | kafka-broker     | 192.168.11.1  |
371/// | kafka-producer   | 192.168.11.2  |
372/// | object_store_sim | 192.168.12.1  |
373/// | client           | 192.168.100.1 |
374/// | ctl              | 192.168.101.1 |
375pub struct Cluster {
376    config: Configuration,
377    handle: Handle,
378    #[cfg(madsim)]
379    pub(crate) client: NodeHandle,
380    #[cfg(madsim)]
381    pub(crate) ctl: NodeHandle,
382    #[cfg(madsim)]
383    pub(crate) sqlite_file_handle: NamedTempFile,
384}
385
386impl Cluster {
387    /// Filesystem path of the `SQLite` metastore backing this simulated cluster.
388    #[cfg_or_panic(madsim)]
389    pub fn meta_sqlite_path(&self) -> &Path {
390        self.sqlite_file_handle.path()
391    }
392
393    /// Start a RisingWave cluster for testing.
394    ///
395    /// This function should be called exactly once in a test.
396    #[cfg_or_panic(madsim)]
397    pub async fn start(conf: Configuration) -> Result<Self> {
398        use madsim::net::ipvs::*;
399
400        let handle = madsim::runtime::Handle::current();
401        println!("seed = {}", handle.seed());
402        println!("{:#?}", conf);
403
404        // TODO: support mutil meta nodes
405        assert_eq!(conf.meta_nodes, 1);
406
407        // setup DNS and load balance
408        let net = madsim::net::NetSim::current();
409        for i in 1..=conf.meta_nodes {
410            net.add_dns_record(
411                &format!("meta-{i}"),
412                format!("192.168.1.{i}").parse().unwrap(),
413            );
414        }
415
416        net.add_dns_record("frontend", "192.168.2.0".parse().unwrap());
417        net.add_dns_record("message_queue", "192.168.11.1".parse().unwrap());
418        net.global_ipvs().add_service(
419            ServiceAddr::Tcp("192.168.2.0:4566".into()),
420            Scheduler::RoundRobin,
421        );
422        for i in 1..=conf.frontend_nodes {
423            net.global_ipvs().add_server(
424                ServiceAddr::Tcp("192.168.2.0:4566".into()),
425                &format!("192.168.2.{i}:4566"),
426            )
427        }
428
429        // kafka broker
430        handle
431            .create_node()
432            .name("kafka-broker")
433            .ip("192.168.11.1".parse().unwrap())
434            .init(move || async move {
435                rdkafka::SimBroker::default()
436                    .serve("0.0.0.0:29092".parse().unwrap())
437                    .await
438            })
439            .build();
440
441        // object_store_sim
442        handle
443            .create_node()
444            .name("object_store_sim")
445            .ip("192.168.12.1".parse().unwrap())
446            .init(move || async move {
447                ObjectStoreSimServer::builder()
448                    .serve("0.0.0.0:9301".parse().unwrap())
449                    .await
450            })
451            .build();
452
453        // wait for the service to be ready
454        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
455
456        let mut meta_addrs = vec![];
457        for i in 1..=conf.meta_nodes {
458            meta_addrs.push(format!("http://meta-{i}:5690"));
459        }
460        unsafe { std::env::set_var("RW_META_ADDR", meta_addrs.join(",")) };
461
462        let sqlite_file_handle: NamedTempFile = NamedTempFile::new().unwrap();
463        let file_path = sqlite_file_handle.path().display().to_string();
464        tracing::info!(?file_path, "sqlite_file_path");
465        let sql_endpoint = format!("sqlite://{}?mode=rwc", file_path);
466        let backend_args = vec!["--backend", "sql", "--sql-endpoint", &sql_endpoint];
467
468        // meta node
469        for i in 1..=conf.meta_nodes {
470            let args = [
471                "meta-node",
472                "--config-path",
473                conf.config_path.as_str(),
474                "--listen-addr",
475                "0.0.0.0:5690",
476                "--advertise-addr",
477                &format!("meta-{i}:5690"),
478                "--state-store",
479                "hummock+sim://hummockadmin:hummockadmin@192.168.12.1:9301/hummock001",
480                "--data-directory",
481                "hummock_001",
482                "--temp-secret-file-dir",
483                &format!("./secrets/meta-{i}"),
484            ];
485            let args = args.into_iter().chain(backend_args.clone().into_iter());
486            let opts = risingwave_meta_node::MetaNodeOpts::parse_from(args);
487            handle
488                .create_node()
489                .name(format!("meta-{i}"))
490                .ip([192, 168, 1, i as u8].into())
491                .init(move || {
492                    risingwave_meta_node::start(
493                        opts.clone(),
494                        CancellationToken::new(), // dummy
495                    )
496                })
497                .build();
498        }
499
500        // wait for the service to be ready
501        tokio::time::sleep(std::time::Duration::from_secs(15)).await;
502
503        // frontend node
504        for i in 1..=conf.frontend_nodes {
505            let opts = risingwave_frontend::FrontendOpts::parse_from([
506                "frontend-node",
507                "--config-path",
508                conf.config_path.as_str(),
509                "--listen-addr",
510                "0.0.0.0:4566",
511                "--health-check-listener-addr",
512                "0.0.0.0:6786",
513                "--advertise-addr",
514                &format!("192.168.2.{i}:4566"),
515                "--temp-secret-file-dir",
516                &format!("./secrets/frontend-{i}"),
517            ]);
518            handle
519                .create_node()
520                .name(format!("frontend-{i}"))
521                .ip([192, 168, 2, i as u8].into())
522                .init(move || {
523                    risingwave_frontend::start(
524                        opts.clone(),
525                        CancellationToken::new(), // dummy
526                    )
527                })
528                .build();
529        }
530
531        // compute node
532        for i in 1..=conf.compute_nodes {
533            let opts = risingwave_compute::ComputeNodeOpts::parse_from([
534                "compute-node",
535                "--config-path",
536                conf.config_path.as_str(),
537                "--listen-addr",
538                "0.0.0.0:5688",
539                "--advertise-addr",
540                &format!("192.168.3.{i}:5688"),
541                "--total-memory-bytes",
542                "6979321856",
543                "--parallelism",
544                &conf.compute_node_cores.to_string(),
545                "--temp-secret-file-dir",
546                &format!("./secrets/compute-{i}"),
547                "--resource-group",
548                &conf
549                    .compute_resource_groups
550                    .get(&i)
551                    .cloned()
552                    .unwrap_or(DEFAULT_RESOURCE_GROUP.to_string()),
553                "--role",
554                &conf
555                    .compute_node_roles
556                    .get(&i)
557                    .cloned()
558                    .unwrap_or("both".to_string()),
559            ]);
560            handle
561                .create_node()
562                .name(format!("compute-{i}"))
563                .ip([192, 168, 3, i as u8].into())
564                .cores(conf.compute_node_cores)
565                .init(move || {
566                    risingwave_compute::start(
567                        opts.clone(),
568                        CancellationToken::new(), // dummy
569                    )
570                })
571                .build();
572        }
573
574        // compactor node
575        for i in 1..=conf.compactor_nodes {
576            let opts = risingwave_compactor::CompactorOpts::parse_from([
577                "compactor-node",
578                "--config-path",
579                conf.config_path.as_str(),
580                "--listen-addr",
581                "0.0.0.0:6660",
582                "--advertise-addr",
583                &format!("192.168.4.{i}:6660"),
584            ]);
585            handle
586                .create_node()
587                .name(format!("compactor-{i}"))
588                .ip([192, 168, 4, i as u8].into())
589                .init(move || {
590                    risingwave_compactor::start(
591                        opts.clone(),
592                        CancellationToken::new(), // dummy
593                    )
594                })
595                .build();
596        }
597
598        // wait for the service to be ready
599        tokio::time::sleep(Duration::from_secs(15)).await;
600
601        // client
602        let client = handle
603            .create_node()
604            .name("client")
605            .ip([192, 168, 100, 1].into())
606            .build();
607
608        // risectl
609        let ctl = handle
610            .create_node()
611            .name("ctl")
612            .ip([192, 168, 101, 1].into())
613            .build();
614
615        Ok(Self {
616            config: conf,
617            handle,
618            client,
619            ctl,
620            sqlite_file_handle,
621        })
622    }
623
624    #[cfg_or_panic(madsim)]
625    fn per_session_queries(&self) -> Arc<Vec<String>> {
626        self.config.per_session_queries.clone()
627    }
628
629    /// Start a SQL session on the client node.
630    #[cfg_or_panic(madsim)]
631    pub fn start_session(&mut self) -> Session {
632        let (query_tx, mut query_rx) = mpsc::channel::<SessionRequest>(0);
633        let per_session_queries = self.per_session_queries();
634
635        self.client.spawn(async move {
636            let mut client = RisingWave::connect("frontend".into(), "dev".into()).await?;
637
638            for sql in per_session_queries.as_ref() {
639                client.run(sql).await?;
640            }
641            drop(per_session_queries);
642
643            while let Some((sql, tx)) = query_rx.next().await {
644                let result = client
645                    .run(&sql)
646                    .await
647                    .map(|output| match output {
648                        sqllogictest::DBOutput::Rows { rows, .. } => rows
649                            .into_iter()
650                            .map(|row| {
651                                row.into_iter()
652                                    .map(|v| v.to_string())
653                                    .collect::<Vec<_>>()
654                                    .join(" ")
655                            })
656                            .collect::<Vec<_>>()
657                            .join("\n"),
658                        _ => "".to_string(),
659                    })
660                    .map_err(Into::into);
661
662                let _ = tx.send(result);
663            }
664
665            Ok::<_, anyhow::Error>(())
666        });
667
668        Session { query_tx }
669    }
670
671    /// Run a SQL query on a **new** session of the client node.
672    ///
673    /// This is a convenience method that creates a new session and runs the query on it. If you
674    /// want to run multiple queries on the same session, use `start_session` and `Session::run`.
675    pub async fn run(&mut self, sql: impl Into<String>) -> Result<String> {
676        self.start_session().run(sql).await
677    }
678
679    /// Run a future on the client node.
680    #[cfg_or_panic(madsim)]
681    pub async fn run_on_client<F>(&self, future: F) -> F::Output
682    where
683        F: Future + Send + 'static,
684        F::Output: Send + 'static,
685    {
686        self.client.spawn(future).await.unwrap()
687    }
688
689    pub async fn get_random_worker_nodes(&self, n: usize) -> Result<Vec<WorkerNode>> {
690        let worker_nodes = self.get_cluster_info().await?.get_worker_nodes().clone();
691        if worker_nodes.len() < n {
692            return Err(anyhow!("cannot remove more nodes than present"));
693        }
694        let rand_nodes = worker_nodes.iter().sample(&mut rand::rng(), n).clone();
695        Ok(rand_nodes.iter().cloned().cloned().collect_vec())
696    }
697
698    /// Run a SQL query from the client and wait until the condition is met.
699    pub async fn wait_until(
700        &mut self,
701        sql: impl Into<String> + Clone,
702        mut p: impl FnMut(&str) -> bool,
703        interval: Duration,
704        timeout: Duration,
705    ) -> Result<String> {
706        let fut = async move {
707            let mut interval = tokio::time::interval(interval);
708            loop {
709                interval.tick().await;
710                let result = self.run(sql.clone()).await?;
711                if p(&result) {
712                    return Ok::<_, anyhow::Error>(result);
713                }
714            }
715        };
716
717        match tokio::time::timeout(timeout, fut).await {
718            Ok(r) => Ok(r?),
719            Err(_) => bail!("wait_until timeout"),
720        }
721    }
722
723    /// Run a SQL query from the client and wait until the return result is not empty.
724    pub async fn wait_until_non_empty(
725        &mut self,
726        sql: &str,
727        interval: Duration,
728        timeout: Duration,
729    ) -> Result<String> {
730        self.wait_until(sql, |r| !r.trim().is_empty(), interval, timeout)
731            .await
732    }
733
734    /// Generate a list of random worker nodes to kill by `opts`, then call `kill_nodes` to kill and
735    /// restart them.
736    pub async fn kill_node(&self, opts: &KillOpts) {
737        let mut nodes = vec![];
738        if opts.kill_meta {
739            let rand = rand::rng().random_range(0..3);
740            for i in 1..=self.config.meta_nodes {
741                match rand {
742                    0 => break,                                     // no killed
743                    1 => {}                                         // all killed
744                    _ if !rand::rng().random_bool(0.5) => continue, // random killed
745                    _ => {}
746                }
747                nodes.push(format!("meta-{}", i));
748            }
749            // don't kill all meta services
750            if nodes.len() == self.config.meta_nodes {
751                nodes.truncate(1);
752            }
753        }
754        if opts.kill_frontend {
755            let rand = rand::rng().random_range(0..3);
756            for i in 1..=self.config.frontend_nodes {
757                match rand {
758                    0 => break,                                     // no killed
759                    1 => {}                                         // all killed
760                    _ if !rand::rng().random_bool(0.5) => continue, // random killed
761                    _ => {}
762                }
763                nodes.push(format!("frontend-{}", i));
764            }
765        }
766        if opts.kill_compute {
767            let rand = rand::rng().random_range(0..3);
768            for i in 1..=self.config.compute_nodes {
769                match rand {
770                    0 => break,                                     // no killed
771                    1 => {}                                         // all killed
772                    _ if !rand::rng().random_bool(0.5) => continue, // random killed
773                    _ => {}
774                }
775                nodes.push(format!("compute-{}", i));
776            }
777        }
778        if opts.kill_compactor {
779            let rand = rand::rng().random_range(0..3);
780            for i in 1..=self.config.compactor_nodes {
781                match rand {
782                    0 => break,                                     // no killed
783                    1 => {}                                         // all killed
784                    _ if !rand::rng().random_bool(0.5) => continue, // random killed
785                    _ => {}
786                }
787                nodes.push(format!("compactor-{}", i));
788            }
789        }
790
791        self.kill_nodes(nodes, opts.restart_delay_secs).await
792    }
793
794    /// Kill the given nodes by their names and restart them in 2s + `restart_delay_secs` with a
795    /// probability of 0.1.
796    #[cfg_or_panic(madsim)]
797    pub async fn kill_nodes(
798        &self,
799        nodes: impl IntoIterator<Item = impl AsRef<str>>,
800        restart_delay_secs: u32,
801    ) {
802        join_all(nodes.into_iter().map(|name| async move {
803            let name = name.as_ref();
804            let t = rand::rng().random_range(Duration::from_secs(0)..Duration::from_secs(1));
805            tokio::time::sleep(t).await;
806            tracing::info!("kill {name}");
807            Handle::current().kill(name);
808
809            let mut t = rand::rng().random_range(Duration::from_secs(0)..Duration::from_secs(1));
810            // has a small chance to restart after a long time
811            // so that the node is expired and removed from the cluster
812            if rand::rng().random_bool(0.1) {
813                // max_heartbeat_interval_secs = 15
814                t += Duration::from_secs(restart_delay_secs as u64);
815            }
816            tokio::time::sleep(t).await;
817            tracing::info!("restart {name}");
818            Handle::current().restart(name);
819        }))
820        .await;
821    }
822
823    #[cfg_or_panic(madsim)]
824    pub async fn kill_nodes_and_restart(
825        &self,
826        nodes: impl IntoIterator<Item = impl AsRef<str>>,
827        restart_delay_secs: u32,
828    ) {
829        join_all(nodes.into_iter().map(|name| async move {
830            let name = name.as_ref();
831            tracing::info!("kill {name}");
832            Handle::current().kill(name);
833            tokio::time::sleep(Duration::from_secs(restart_delay_secs as u64)).await;
834            tracing::info!("restart {name}");
835            Handle::current().restart(name);
836        }))
837        .await;
838    }
839
840    #[cfg_or_panic(madsim)]
841    pub async fn simple_kill_nodes(&self, nodes: impl IntoIterator<Item = impl AsRef<str>>) {
842        join_all(nodes.into_iter().map(|name| async move {
843            let name = name.as_ref();
844            tracing::info!("kill {name}");
845            Handle::current().kill(name);
846        }))
847        .await;
848    }
849
850    #[cfg_or_panic(madsim)]
851    pub async fn simple_restart_nodes(&self, nodes: impl IntoIterator<Item = impl AsRef<str>>) {
852        join_all(nodes.into_iter().map(|name| async move {
853            let name = name.as_ref();
854            tracing::info!("restart {name}");
855            Handle::current().restart(name);
856        }))
857        .await;
858    }
859
860    /// Create a node for kafka producer and prepare data.
861    #[cfg_or_panic(madsim)]
862    pub async fn create_kafka_producer(&self, datadir: &str) {
863        self.handle
864            .create_node()
865            .name("kafka-producer")
866            .ip("192.168.11.2".parse().unwrap())
867            .build()
868            .spawn(crate::kafka::producer(
869                "192.168.11.1:29092",
870                datadir.to_string(),
871            ))
872            .await
873            .unwrap();
874    }
875
876    /// Create a kafka topic.
877    #[cfg_or_panic(madsim)]
878    pub fn create_kafka_topics(&self, topics: HashMap<String, i32>) {
879        self.handle
880            .create_node()
881            .name("kafka-topic-create")
882            .ip("192.168.11.3".parse().unwrap())
883            .build()
884            .spawn(crate::kafka::create_topics("192.168.11.1:29092", topics));
885    }
886
887    pub fn config(&self) -> Configuration {
888        self.config.clone()
889    }
890
891    pub fn handle(&self) -> &Handle {
892        &self.handle
893    }
894
895    /// Graceful shutdown all RisingWave nodes.
896    #[cfg_or_panic(madsim)]
897    pub async fn graceful_shutdown(&self) {
898        let mut nodes = vec![];
899        let mut metas = vec![];
900        for i in 1..=self.config.meta_nodes {
901            metas.push(format!("meta-{i}"));
902        }
903        for i in 1..=self.config.frontend_nodes {
904            nodes.push(format!("frontend-{i}"));
905        }
906        for i in 1..=self.config.compute_nodes {
907            nodes.push(format!("compute-{i}"));
908        }
909        for i in 1..=self.config.compactor_nodes {
910            nodes.push(format!("compactor-{i}"));
911        }
912
913        tracing::info!("graceful shutdown");
914        let waiting_time = Duration::from_secs(10);
915        // shutdown frontends, computes, compactors
916        for node in &nodes {
917            self.handle.send_ctrl_c(node);
918        }
919        tokio::time::sleep(waiting_time).await;
920        // shutdown metas
921        for meta in &metas {
922            self.handle.send_ctrl_c(meta);
923        }
924        tokio::time::sleep(waiting_time).await;
925
926        // check all nodes are exited
927        for node in nodes.iter().chain(metas.iter()) {
928            if !self.handle.is_exit(node) {
929                panic!("failed to graceful shutdown {node} in {waiting_time:?}");
930            }
931        }
932    }
933
934    pub async fn wait_for_recovery(&mut self) -> Result<()> {
935        let timeout = Duration::from_secs(200);
936        let mut session = self.start_session();
937        tokio::time::timeout(timeout, async {
938            loop {
939                if let Ok(result) = session.run("select rw_recovery_status()").await
940                    && result == "RUNNING"
941                {
942                    break;
943                }
944                tokio::time::sleep(Duration::from_nanos(10)).await;
945            }
946        })
947        .await?;
948        Ok(())
949    }
950
951    /// This function only works if all actors in your cluster are following adaptive scaling.
952    pub async fn wait_for_scale(&mut self, parallelism: usize) -> Result<()> {
953        let timeout = Duration::from_secs(200);
954        let mut session = self.start_session();
955        tokio::time::timeout(timeout, async {
956            loop {
957                let parallelism_sql = format!(
958                    "select count(parallelism) filter (where parallelism != {parallelism})\
959                from (select count(*) parallelism from rw_actors group by fragment_id);"
960                );
961                if let Ok(result) = session.run(&parallelism_sql).await
962                    && result == "0"
963                {
964                    break;
965                }
966                tokio::time::sleep(Duration::from_nanos(10)).await;
967            }
968        })
969        .await?;
970        Ok(())
971    }
972}
973
974type SessionRequest = (
975    String,                          // query sql
976    oneshot::Sender<Result<String>>, // channel to send result back
977);
978
979/// A SQL session on the simulated client node.
980#[derive(Debug, Clone)]
981pub struct Session {
982    query_tx: mpsc::Sender<SessionRequest>,
983}
984
985impl Session {
986    /// Run the given SQLs on the session.
987    pub async fn run_all(&mut self, sqls: Vec<impl Into<String>>) -> Result<Vec<String>> {
988        let mut results = Vec::with_capacity(sqls.len());
989        for sql in sqls {
990            let result = self.run(sql).await?;
991            results.push(result);
992        }
993        Ok(results)
994    }
995
996    /// Run the given SQL query on the session.
997    pub async fn run(&mut self, sql: impl Into<String>) -> Result<String> {
998        let (tx, rx) = oneshot::channel();
999        self.query_tx.send((sql.into(), tx)).await?;
1000        rx.await?
1001    }
1002
1003    /// Run `FLUSH` on the session.
1004    pub async fn flush(&mut self) -> Result<()> {
1005        self.run("FLUSH").await?;
1006        Ok(())
1007    }
1008
1009    pub async fn is_arrangement_backfill_enabled(&mut self) -> Result<bool> {
1010        let result = self.run("show streaming_use_arrangement_backfill").await?;
1011        Ok(result == "true")
1012    }
1013}
1014
1015/// Options for killing nodes.
1016#[derive(Debug, Clone, Copy, PartialEq)]
1017pub struct KillOpts {
1018    pub kill_rate: f32,
1019    pub kill_meta: bool,
1020    pub kill_frontend: bool,
1021    pub kill_compute: bool,
1022    pub kill_compactor: bool,
1023    pub restart_delay_secs: u32,
1024}
1025
1026impl KillOpts {
1027    /// Killing all kind of nodes.
1028    pub const ALL: Self = KillOpts {
1029        kill_rate: 1.0,
1030        kill_meta: false, // FIXME: make it true when multiple meta nodes are supported
1031        kill_frontend: true,
1032        kill_compute: true,
1033        kill_compactor: true,
1034        restart_delay_secs: 20,
1035    };
1036    pub const ALL_FAST: Self = KillOpts {
1037        kill_rate: 1.0,
1038        kill_meta: false, // FIXME: make it true when multiple meta nodes are supported
1039        kill_frontend: true,
1040        kill_compute: true,
1041        kill_compactor: true,
1042        restart_delay_secs: 2,
1043    };
1044}
risingwave_simulation/cluster.rs

risingwave_simulation/
cluster.rs
