risedev/task/
task_tcp_ready_check.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
15use anyhow::{Context, Result};
16
17use super::{ExecuteContext, Task};
18
19/// Check if a TCP port can be connected to.
20///
21/// Note that accepting a connection does not always mean the service is ready.
22pub struct TcpReadyCheckTask {
23    advertise_address: String,
24    port: u16,
25    user_managed: bool,
26}
27
28impl TcpReadyCheckTask {
29    pub fn new(advertise_address: String, port: u16, user_managed: bool) -> Result<Self> {
30        Ok(Self {
31            advertise_address,
32            port,
33            user_managed,
34        })
35    }
36}
37
38impl Task for TcpReadyCheckTask {
39    fn execute(&mut self, ctx: &mut ExecuteContext<impl std::io::Write>) -> anyhow::Result<()> {
40        let Some(id) = ctx.id.clone() else {
41            panic!("Service should be set before executing TcpReadyCheckTask");
42        };
43        let address = format!("{}:{}", self.advertise_address, self.port);
44
45        if self.user_managed {
46            ctx.pb.set_message(
47                "waiting for user-managed service online... (see `risedev.log` for cli args)",
48            );
49            ctx.wait_tcp_user(&address).with_context(|| {
50                format!("failed to wait for user-managed service `{id}` to be online")
51            })?;
52        } else {
53            ctx.pb.set_message("waiting for online...");
54            ctx.wait_tcp(&address)
55                .with_context(|| format!("failed to wait for service `{id}` to be online"))?;
56        }
57
58        ctx.complete_spin();
59
60        Ok(())
61    }
62}
risedev/task/task_tcp_ready_check.rs

risedev/task/
task_tcp_ready_check.rs
