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

risedev/task/
mod.rs
