risedev/task/
task_log_ready_check.rs

1// Copyright 2024 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
15use std::io::{Read as _, Seek as _, SeekFrom};
16use std::time::Duration;
17
18use anyhow::{Context, Result, bail};
19use fs_err::File;
20
21use super::{ExecuteContext, Task};
22use crate::wait::wait;
23
24const DEFAULT_LOG_READY_TIMEOUT: Duration = Duration::from_secs(60);
25
26/// Check if all log patterns are found in the log output indicating the service is ready.
27pub struct LogReadyCheckTask {
28    patterns: Vec<String>,
29    timeout: Duration,
30}
31
32impl LogReadyCheckTask {
33    pub fn new(pattern: impl Into<String>) -> Result<Self> {
34        Ok(Self {
35            patterns: vec![pattern.into()],
36            timeout: DEFAULT_LOG_READY_TIMEOUT,
37        })
38    }
39
40    pub fn new_all(patterns: impl IntoIterator<Item = impl Into<String>>) -> Result<Self> {
41        Ok(Self {
42            patterns: patterns.into_iter().map(Into::into).collect(),
43            timeout: DEFAULT_LOG_READY_TIMEOUT,
44        })
45    }
46
47    pub fn with_timeout(mut self, timeout: Duration) -> Self {
48        self.timeout = timeout;
49        self
50    }
51}
52
53impl Task for LogReadyCheckTask {
54    fn execute(&mut self, ctx: &mut ExecuteContext<impl std::io::Write>) -> anyhow::Result<()> {
55        let Some(id) = ctx.id.clone() else {
56            panic!("Service should be set before executing LogReadyCheckTask");
57        };
58
59        ctx.pb.set_message("waiting for ready...");
60        ctx.wait_log_contains(&self.patterns, self.timeout)
61            .with_context(|| format!("failed to wait for service `{id}` to be ready"))?;
62
63        ctx.complete_spin();
64
65        Ok(())
66    }
67}
68
69impl<W> ExecuteContext<W>
70where
71    W: std::io::Write,
72{
73    fn wait_log_contains(&mut self, patterns: &[String], timeout: Duration) -> anyhow::Result<()> {
74        let log_path = self.log_path().to_path_buf();
75
76        let mut content = String::new();
77        let mut offset = 0;
78
79        wait(
80            || {
81                let mut file = File::open(&log_path).context("log file does not exist")?;
82                file.seek(SeekFrom::Start(offset as u64))?;
83                offset += file.read_to_string(&mut content)?;
84
85                // Always going through the whole log file could be stupid, but it's reliable.
86                for pattern in patterns {
87                    if !content.contains(pattern) {
88                        bail!("pattern \"{}\" not found in log", pattern)
89                    }
90                }
91
92                Ok(())
93            },
94            &mut self.log,
95            self.status_file.as_ref().unwrap(),
96            self.id.as_ref().unwrap(),
97            Some(timeout),
98            true,
99        )?;
100
101        Ok(())
102    }
103}
risedev/task/task_log_ready_check.rs

risedev/task/
task_log_ready_check.rs
