risedev/task/
mod.rs

1// Copyright 2025 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
15mod clickhouse_service;
16mod compactor_service;
17mod compute_node_service;
18mod configure_tmux_service;
19mod docker_service;
20mod dummy_service;
21mod elasticsearch_service;
22mod ensure_stop_service;
23mod frontend_service;
24mod grafana_service;
25mod kafka_service;
26mod lakekeeper_service;
27mod meta_node_service;
28mod minio_service;
29mod moat_service;
30mod mongodb_service;
31mod moto_service;
32mod mqtt_service;
33mod mysql_service;
34mod nats_service;
35mod opensearch_service;
36mod oracle_service;
37mod postgres_service;
38mod prometheus_service;
39mod pubsub_service;
40mod pulsar_service;
41mod redis_service;
42mod schema_registry_service;
43mod sql_server_service;
44mod task_clickhouse_ready_check;
45mod task_configure_minio;
46mod task_db_ready_check;
47mod task_kafka_ready_check;
48mod task_log_ready_check;
49mod task_pubsub_emu_ready_check;
50mod task_redis_ready_check;
51mod task_tcp_ready_check;
52mod tempo_service;
53mod utils;
54
55use std::collections::HashMap;
56use std::env;
57use std::io::Write;
58use std::net::{TcpStream, ToSocketAddrs};
59use std::path::{Path, PathBuf};
60use std::process::{Command, Output};
61use std::sync::Arc;
62use std::time::Duration;
63
64use anyhow::{Context, Result, anyhow};
65use indicatif::ProgressBar;
66use reqwest::blocking::{Client, Response};
67use tempfile::TempDir;
68pub use utils::*;
69
70pub use self::clickhouse_service::*;
71pub use self::compactor_service::*;
72pub use self::compute_node_service::*;
73pub use self::configure_tmux_service::*;
74pub use self::dummy_service::DummyService;
75pub use self::elasticsearch_service::*;
76pub use self::ensure_stop_service::*;
77pub use self::frontend_service::*;
78pub use self::grafana_service::*;
79pub use self::kafka_service::*;
80pub use self::lakekeeper_service::*;
81pub use self::meta_node_service::*;
82pub use self::minio_service::*;
83pub use self::moat_service::*;
84pub use self::mongodb_service::*;
85pub use self::moto_service::*;
86pub use self::mqtt_service::*;
87pub use self::mysql_service::*;
88pub use self::nats_service::*;
89pub use self::opensearch_service::*;
90pub use self::oracle_service::*;
91pub use self::postgres_service::*;
92pub use self::prometheus_service::*;
93pub use self::pubsub_service::*;
94pub use self::pulsar_service::*;
95pub use self::redis_service::*;
96pub use self::schema_registry_service::SchemaRegistryService;
97pub use self::sql_server_service::*;
98pub use self::task_clickhouse_ready_check::*;
99pub use self::task_configure_minio::*;
100pub use self::task_db_ready_check::*;
101pub use self::task_kafka_ready_check::*;
102pub use self::task_log_ready_check::*;
103pub use self::task_pubsub_emu_ready_check::*;
104pub use self::task_redis_ready_check::*;
105pub use self::task_tcp_ready_check::*;
106pub use self::tempo_service::*;
107use crate::util::{begin_spin, complete_spin, get_program_args, get_program_name};
108use crate::wait::{wait, wait_tcp_available};
109
110pub trait Task: 'static + Send {
111    /// Execute the task
112    fn execute(&mut self, ctx: &mut ExecuteContext<impl std::io::Write>) -> anyhow::Result<()>;
113
114    /// Get task id used in progress bar
115    fn id(&self) -> String {
116        "<task>".into()
117    }
118}
119
120/// A context used in task execution
121pub struct ExecuteContext<W>
122where
123    W: std::io::Write,
124{
125    /// Global log file object. (aka. risedev.log)
126    pub log: W,
127
128    /// Progress bar on screen.
129    pub pb: ProgressBar,
130
131    /// The directory for checking status.
132    ///
133    /// `RiseDev` will instruct every task to output their status to a file in temporary folder. By
134    /// checking this file, we can know whether a task has early exited.
135    pub status_dir: Arc<TempDir>,
136
137    /// The current service id running in this context.
138    pub id: Option<String>,
139
140    /// The status file corresponding to the current context.
141    pub status_file: Option<PathBuf>,
142
143    /// The log file corresponding to the current context. (e.g. frontend-4566.log)
144    pub log_file: Option<PathBuf>,
145}
146
147impl<W> ExecuteContext<W>
148where
149    W: std::io::Write,
150{
151    pub fn new(log: W, pb: ProgressBar, status_dir: Arc<TempDir>) -> Self {
152        Self {
153            log,
154            pb,
155            status_dir,
156            status_file: None,
157            log_file: None,
158            id: None,
159        }
160    }
161
162    pub fn service(&mut self, task: &impl Task) {
163        let id = task.id();
164        if !id.is_empty() {
165            begin_spin(&self.pb);
166            self.pb.set_prefix(id.clone());
167            self.id = Some(id.clone());
168            self.status_file = Some(self.status_dir.path().join(format!("{}.status", id)));
169
170            // Remove the old log file if exists to avoid confusion.
171            let log_file = Path::new(&env::var("PREFIX_LOG").unwrap())
172                .join(format!("{}.log", self.id.as_ref().unwrap()));
173            fs_err::remove_file(&log_file).ok();
174            self.log_file = Some(log_file);
175        }
176    }
177
178    pub fn run_command(&mut self, mut cmd: Command) -> Result<Output> {
179        let program_name = get_program_name(&cmd);
180
181        writeln!(self.log, "> {} {}", program_name, get_program_args(&cmd))?;
182
183        // Record a service command for `risedev restart`
184        let node_id = self.id.as_ref();
185        if program_name == "tmux"
186            && let Some(node_id) = node_id
187            && node_id != "tmux-configure"
188        {
189            let cmd: HashMap<String, String> = HashMap::from_iter([(
190                node_id.clone(),
191                format!("{} {}", program_name, get_program_args(&cmd)),
192            )]);
193            let prefix_config = env::var("PREFIX_CONFIG").unwrap();
194            let path = Path::new(&prefix_config).join("risedev_commands.yaml");
195            let content = serde_yaml::to_string(&cmd)?;
196            fs_err::OpenOptions::new()
197                .create(true)
198                .append(true)
199                .open(path)?
200                .write_all(content.as_bytes())?;
201        }
202
203        let output = cmd.output()?;
204
205        let mut full_output = String::from_utf8_lossy(&output.stdout).to_string();
206        full_output.extend(String::from_utf8_lossy(&output.stderr).chars());
207
208        write!(self.log, "{}", full_output)?;
209
210        writeln!(
211            self.log,
212            "({} exited with {:?})",
213            program_name,
214            output.status.code()
215        )?;
216
217        writeln!(self.log, "---")?;
218
219        output.status.exit_ok().context(full_output)?;
220
221        Ok(output)
222    }
223
224    pub fn complete_spin(&mut self) {
225        complete_spin(&self.pb);
226    }
227
228    pub fn status_path(&self) -> PathBuf {
229        self.status_file.clone().unwrap()
230    }
231
232    pub fn log_path(&self) -> &Path {
233        self.log_file.as_ref().unwrap().as_path()
234    }
235
236    pub fn wait_tcp(&mut self, server: impl AsRef<str>) -> anyhow::Result<()> {
237        let addr = server
238            .as_ref()
239            .to_socket_addrs()?
240            .next()
241            .with_context(|| format!("failed to resolve {}", server.as_ref()))?;
242        wait(
243            || {
244                TcpStream::connect_timeout(&addr, Duration::from_secs(1)).with_context(|| {
245                    format!("failed to establish tcp connection to {}", server.as_ref())
246                })?;
247                Ok(())
248            },
249            &mut self.log,
250            self.status_file.as_ref().unwrap(),
251            self.id.as_ref().unwrap(),
252            Some(Duration::from_secs(30)),
253            true,
254        )?;
255        Ok(())
256    }
257
258    fn wait_http_with_response_cb(
259        &mut self,
260        server: impl AsRef<str>,
261        cb: impl Fn(Response) -> anyhow::Result<()>,
262    ) -> anyhow::Result<()> {
263        let server = server.as_ref();
264        wait(
265            || {
266                let resp = Client::new()
267                    .get(server)
268                    .timeout(Duration::from_secs(1))
269                    .body("")
270                    .send()?
271                    .error_for_status()
272                    .with_context(|| {
273                        format!("failed to establish http connection to {}", server)
274                    })?;
275
276                cb(resp)
277            },
278            &mut self.log,
279            self.status_file.as_ref().unwrap(),
280            self.id.as_ref().unwrap(),
281            Some(Duration::from_secs(30)),
282            true,
283        )
284    }
285
286    pub fn wait_http(&mut self, server: impl AsRef<str>) -> anyhow::Result<()> {
287        self.wait_http_with_response_cb(server, |_| Ok(()))
288    }
289
290    pub fn wait_http_with_text_cb(
291        &mut self,
292        server: impl AsRef<str>,
293        cb: impl Fn(&str) -> bool,
294    ) -> anyhow::Result<()> {
295        self.wait_http_with_response_cb(server, |resp| {
296            let data = resp.text()?;
297            if cb(&data) {
298                Ok(())
299            } else {
300                Err(anyhow!(
301                    "http health check callback failed with body: {:?}",
302                    data
303                ))
304            }
305        })
306    }
307
308    pub fn wait(&mut self, wait_func: impl FnMut() -> Result<()>) -> anyhow::Result<()> {
309        wait(
310            wait_func,
311            &mut self.log,
312            self.status_file.as_ref().unwrap(),
313            self.id.as_ref().unwrap(),
314            Some(Duration::from_secs(30)),
315            true,
316        )
317    }
318
319    /// Wait for a TCP port to close
320    pub fn wait_tcp_close(&mut self, server: impl AsRef<str>) -> anyhow::Result<()> {
321        wait_tcp_available(server, Some(Duration::from_secs(30)))?;
322        Ok(())
323    }
324
325    /// Wait for a user-managed service to be available
326    pub fn wait_tcp_user(&mut self, server: impl AsRef<str>) -> anyhow::Result<()> {
327        let addr = server
328            .as_ref()
329            .to_socket_addrs()?
330            .next()
331            .unwrap_or_else(|| panic!("failed to resolve {}", server.as_ref()));
332        wait(
333            || {
334                TcpStream::connect_timeout(&addr, Duration::from_secs(1))?;
335                Ok(())
336            },
337            &mut self.log,
338            self.status_file.as_ref().unwrap(),
339            self.id.as_ref().unwrap(),
340            None,
341            false,
342        )?;
343        Ok(())
344    }
345
346    pub fn tmux_run(&self, user_cmd: Command) -> anyhow::Result<Command> {
347        let prefix_path = env::var("PREFIX_BIN")?;
348        let mut cmd = new_tmux_command();
349        cmd.arg("new-window")
350            // Set target name
351            .arg("-t")
352            .arg(RISEDEV_NAME)
353            // Switch to background window
354            .arg("-d")
355            // Set session name for this window
356            .arg("-n")
357            .arg(self.id.as_ref().unwrap());
358
359        if let Some(dir) = user_cmd.get_current_dir() {
360            cmd.arg("-c").arg(dir);
361        }
362        for (k, v) in user_cmd.get_envs() {
363            cmd.arg("-e");
364            if let Some(v) = v {
365                cmd.arg(format!("{}={}", k.to_string_lossy(), v.to_string_lossy()));
366            } else {
367                cmd.arg(k);
368            }
369        }
370        cmd.arg(Path::new(&prefix_path).join("run_command.sh"));
371        cmd.arg(self.log_path());
372        cmd.arg(self.status_path());
373        cmd.arg(user_cmd.get_program());
374        for arg in user_cmd.get_args() {
375            cmd.arg(arg);
376        }
377
378        Ok(cmd)
379    }
380}
risedev/task/mod.rs

risedev/task/
mod.rs
