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::Rng;
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
695            .iter()
696            .choose_multiple(&mut rand::rng(), n)
697            .clone();
698        Ok(rand_nodes.iter().cloned().cloned().collect_vec())
699    }
700
701    /// Run a SQL query from the client and wait until the condition is met.
702    pub async fn wait_until(
703        &mut self,
704        sql: impl Into<String> + Clone,
705        mut p: impl FnMut(&str) -> bool,
706        interval: Duration,
707        timeout: Duration,
708    ) -> Result<String> {
709        let fut = async move {
710            let mut interval = tokio::time::interval(interval);
711            loop {
712                interval.tick().await;
713                let result = self.run(sql.clone()).await?;
714                if p(&result) {
715                    return Ok::<_, anyhow::Error>(result);
716                }
717            }
718        };
719
720        match tokio::time::timeout(timeout, fut).await {
721            Ok(r) => Ok(r?),
722            Err(_) => bail!("wait_until timeout"),
723        }
724    }
725
726    /// Run a SQL query from the client and wait until the return result is not empty.
727    pub async fn wait_until_non_empty(
728        &mut self,
729        sql: &str,
730        interval: Duration,
731        timeout: Duration,
732    ) -> Result<String> {
733        self.wait_until(sql, |r| !r.trim().is_empty(), interval, timeout)
734            .await
735    }
736
737    /// Generate a list of random worker nodes to kill by `opts`, then call `kill_nodes` to kill and
738    /// restart them.
739    pub async fn kill_node(&self, opts: &KillOpts) {
740        let mut nodes = vec![];
741        if opts.kill_meta {
742            let rand = rand::rng().random_range(0..3);
743            for i in 1..=self.config.meta_nodes {
744                match rand {
745                    0 => break,                                     // no killed
746                    1 => {}                                         // all killed
747                    _ if !rand::rng().random_bool(0.5) => continue, // random killed
748                    _ => {}
749                }
750                nodes.push(format!("meta-{}", i));
751            }
752            // don't kill all meta services
753            if nodes.len() == self.config.meta_nodes {
754                nodes.truncate(1);
755            }
756        }
757        if opts.kill_frontend {
758            let rand = rand::rng().random_range(0..3);
759            for i in 1..=self.config.frontend_nodes {
760                match rand {
761                    0 => break,                                     // no killed
762                    1 => {}                                         // all killed
763                    _ if !rand::rng().random_bool(0.5) => continue, // random killed
764                    _ => {}
765                }
766                nodes.push(format!("frontend-{}", i));
767            }
768        }
769        if opts.kill_compute {
770            let rand = rand::rng().random_range(0..3);
771            for i in 1..=self.config.compute_nodes {
772                match rand {
773                    0 => break,                                     // no killed
774                    1 => {}                                         // all killed
775                    _ if !rand::rng().random_bool(0.5) => continue, // random killed
776                    _ => {}
777                }
778                nodes.push(format!("compute-{}", i));
779            }
780        }
781        if opts.kill_compactor {
782            let rand = rand::rng().random_range(0..3);
783            for i in 1..=self.config.compactor_nodes {
784                match rand {
785                    0 => break,                                     // no killed
786                    1 => {}                                         // all killed
787                    _ if !rand::rng().random_bool(0.5) => continue, // random killed
788                    _ => {}
789                }
790                nodes.push(format!("compactor-{}", i));
791            }
792        }
793
794        self.kill_nodes(nodes, opts.restart_delay_secs).await
795    }
796
797    /// Kill the given nodes by their names and restart them in 2s + `restart_delay_secs` with a
798    /// probability of 0.1.
799    #[cfg_or_panic(madsim)]
800    pub async fn kill_nodes(
801        &self,
802        nodes: impl IntoIterator<Item = impl AsRef<str>>,
803        restart_delay_secs: u32,
804    ) {
805        join_all(nodes.into_iter().map(|name| async move {
806            let name = name.as_ref();
807            let t = rand::rng().random_range(Duration::from_secs(0)..Duration::from_secs(1));
808            tokio::time::sleep(t).await;
809            tracing::info!("kill {name}");
810            Handle::current().kill(name);
811
812            let mut t = rand::rng().random_range(Duration::from_secs(0)..Duration::from_secs(1));
813            // has a small chance to restart after a long time
814            // so that the node is expired and removed from the cluster
815            if rand::rng().random_bool(0.1) {
816                // max_heartbeat_interval_secs = 15
817                t += Duration::from_secs(restart_delay_secs as u64);
818            }
819            tokio::time::sleep(t).await;
820            tracing::info!("restart {name}");
821            Handle::current().restart(name);
822        }))
823        .await;
824    }
825
826    #[cfg_or_panic(madsim)]
827    pub async fn kill_nodes_and_restart(
828        &self,
829        nodes: impl IntoIterator<Item = impl AsRef<str>>,
830        restart_delay_secs: u32,
831    ) {
832        join_all(nodes.into_iter().map(|name| async move {
833            let name = name.as_ref();
834            tracing::info!("kill {name}");
835            Handle::current().kill(name);
836            tokio::time::sleep(Duration::from_secs(restart_delay_secs as u64)).await;
837            tracing::info!("restart {name}");
838            Handle::current().restart(name);
839        }))
840        .await;
841    }
842
843    #[cfg_or_panic(madsim)]
844    pub async fn simple_kill_nodes(&self, nodes: impl IntoIterator<Item = impl AsRef<str>>) {
845        join_all(nodes.into_iter().map(|name| async move {
846            let name = name.as_ref();
847            tracing::info!("kill {name}");
848            Handle::current().kill(name);
849        }))
850        .await;
851    }
852
853    #[cfg_or_panic(madsim)]
854    pub async fn simple_restart_nodes(&self, nodes: impl IntoIterator<Item = impl AsRef<str>>) {
855        join_all(nodes.into_iter().map(|name| async move {
856            let name = name.as_ref();
857            tracing::info!("restart {name}");
858            Handle::current().restart(name);
859        }))
860        .await;
861    }
862
863    /// Create a node for kafka producer and prepare data.
864    #[cfg_or_panic(madsim)]
865    pub async fn create_kafka_producer(&self, datadir: &str) {
866        self.handle
867            .create_node()
868            .name("kafka-producer")
869            .ip("192.168.11.2".parse().unwrap())
870            .build()
871            .spawn(crate::kafka::producer(
872                "192.168.11.1:29092",
873                datadir.to_string(),
874            ))
875            .await
876            .unwrap();
877    }
878
879    /// Create a kafka topic.
880    #[cfg_or_panic(madsim)]
881    pub fn create_kafka_topics(&self, topics: HashMap<String, i32>) {
882        self.handle
883            .create_node()
884            .name("kafka-topic-create")
885            .ip("192.168.11.3".parse().unwrap())
886            .build()
887            .spawn(crate::kafka::create_topics("192.168.11.1:29092", topics));
888    }
889
890    pub fn config(&self) -> Configuration {
891        self.config.clone()
892    }
893
894    pub fn handle(&self) -> &Handle {
895        &self.handle
896    }
897
898    /// Graceful shutdown all RisingWave nodes.
899    #[cfg_or_panic(madsim)]
900    pub async fn graceful_shutdown(&self) {
901        let mut nodes = vec![];
902        let mut metas = vec![];
903        for i in 1..=self.config.meta_nodes {
904            metas.push(format!("meta-{i}"));
905        }
906        for i in 1..=self.config.frontend_nodes {
907            nodes.push(format!("frontend-{i}"));
908        }
909        for i in 1..=self.config.compute_nodes {
910            nodes.push(format!("compute-{i}"));
911        }
912        for i in 1..=self.config.compactor_nodes {
913            nodes.push(format!("compactor-{i}"));
914        }
915
916        tracing::info!("graceful shutdown");
917        let waiting_time = Duration::from_secs(10);
918        // shutdown frontends, computes, compactors
919        for node in &nodes {
920            self.handle.send_ctrl_c(node);
921        }
922        tokio::time::sleep(waiting_time).await;
923        // shutdown metas
924        for meta in &metas {
925            self.handle.send_ctrl_c(meta);
926        }
927        tokio::time::sleep(waiting_time).await;
928
929        // check all nodes are exited
930        for node in nodes.iter().chain(metas.iter()) {
931            if !self.handle.is_exit(node) {
932                panic!("failed to graceful shutdown {node} in {waiting_time:?}");
933            }
934        }
935    }
936
937    pub async fn wait_for_recovery(&mut self) -> Result<()> {
938        let timeout = Duration::from_secs(200);
939        let mut session = self.start_session();
940        tokio::time::timeout(timeout, async {
941            loop {
942                if let Ok(result) = session.run("select rw_recovery_status()").await
943                    && result == "RUNNING"
944                {
945                    break;
946                }
947                tokio::time::sleep(Duration::from_nanos(10)).await;
948            }
949        })
950        .await?;
951        Ok(())
952    }
953
954    /// This function only works if all actors in your cluster are following adaptive scaling.
955    pub async fn wait_for_scale(&mut self, parallelism: usize) -> Result<()> {
956        let timeout = Duration::from_secs(200);
957        let mut session = self.start_session();
958        tokio::time::timeout(timeout, async {
959            loop {
960                let parallelism_sql = format!(
961                    "select count(parallelism) filter (where parallelism != {parallelism})\
962                from (select count(*) parallelism from rw_actors group by fragment_id);"
963                );
964                if let Ok(result) = session.run(&parallelism_sql).await
965                    && result == "0"
966                {
967                    break;
968                }
969                tokio::time::sleep(Duration::from_nanos(10)).await;
970            }
971        })
972        .await?;
973        Ok(())
974    }
975}
976
977type SessionRequest = (
978    String,                          // query sql
979    oneshot::Sender<Result<String>>, // channel to send result back
980);
981
982/// A SQL session on the simulated client node.
983#[derive(Debug, Clone)]
984pub struct Session {
985    query_tx: mpsc::Sender<SessionRequest>,
986}
987
988impl Session {
989    /// Run the given SQLs on the session.
990    pub async fn run_all(&mut self, sqls: Vec<impl Into<String>>) -> Result<Vec<String>> {
991        let mut results = Vec::with_capacity(sqls.len());
992        for sql in sqls {
993            let result = self.run(sql).await?;
994            results.push(result);
995        }
996        Ok(results)
997    }
998
999    /// Run the given SQL query on the session.
1000    pub async fn run(&mut self, sql: impl Into<String>) -> Result<String> {
1001        let (tx, rx) = oneshot::channel();
1002        self.query_tx.send((sql.into(), tx)).await?;
1003        rx.await?
1004    }
1005
1006    /// Run `FLUSH` on the session.
1007    pub async fn flush(&mut self) -> Result<()> {
1008        self.run("FLUSH").await?;
1009        Ok(())
1010    }
1011
1012    pub async fn is_arrangement_backfill_enabled(&mut self) -> Result<bool> {
1013        let result = self.run("show streaming_use_arrangement_backfill").await?;
1014        Ok(result == "true")
1015    }
1016}
1017
1018/// Options for killing nodes.
1019#[derive(Debug, Clone, Copy, PartialEq)]
1020pub struct KillOpts {
1021    pub kill_rate: f32,
1022    pub kill_meta: bool,
1023    pub kill_frontend: bool,
1024    pub kill_compute: bool,
1025    pub kill_compactor: bool,
1026    pub restart_delay_secs: u32,
1027}
1028
1029impl KillOpts {
1030    /// Killing all kind of nodes.
1031    pub const ALL: Self = KillOpts {
1032        kill_rate: 1.0,
1033        kill_meta: false, // FIXME: make it true when multiple meta nodes are supported
1034        kill_frontend: true,
1035        kill_compute: true,
1036        kill_compactor: true,
1037        restart_delay_secs: 20,
1038    };
1039    pub const ALL_FAST: Self = KillOpts {
1040        kill_rate: 1.0,
1041        kill_meta: false, // FIXME: make it true when multiple meta nodes are supported
1042        kill_frontend: true,
1043        kill_compute: true,
1044        kill_compactor: true,
1045        restart_delay_secs: 2,
1046    };
1047}
risingwave_simulation/cluster.rs

risingwave_simulation/
cluster.rs
