refine EnsureStopService

Signed-off-by: xxchan <[email protected]>

e2e_test/source_inline/README.md

@@ -26,5 +26,5 @@ risedev slt 'e2e_test/source_inline/**/*.slt'
 To write tests, please ensure each file is self-contained and does not depend on running external scripts to setup the environment.
 
 Use `system` command to setup instead.
-For simple cases, you can directly write a bash command; 
+For simple cases, you can directly write a bash command;
 For more complex cases, you can write a test script (with any language like bash, python, zx), and invoke it in the `system` command.

src/risedevtool/src/bin/risedev-dev.rs

@@ -98,28 +98,8 @@ fn task_main(
     let mut ports = vec![];
 
     for service in services {
-        let listen_info = match service {
-            ServiceConfig::Minio(c) => Some((c.port, c.id.clone())),
-            ServiceConfig::Etcd(c) => Some((c.port, c.id.clone())),
-            ServiceConfig::Sqlite(_) => None,
-            ServiceConfig::Prometheus(c) => Some((c.port, c.id.clone())),
-            ServiceConfig::ComputeNode(c) => Some((c.port, c.id.clone())),
-            ServiceConfig::MetaNode(c) => Some((c.port, c.id.clone())),
-            ServiceConfig::Frontend(c) => Some((c.port, c.id.clone())),
-            ServiceConfig::Compactor(c) => Some((c.port, c.id.clone())),
-            ServiceConfig::Grafana(c) => Some((c.port, c.id.clone())),
-            ServiceConfig::Tempo(c) => Some((c.port, c.id.clone())),
-            ServiceConfig::Kafka(c) => Some((c.port, c.id.clone())),
-            ServiceConfig::Pubsub(c) => Some((c.port, c.id.clone())),
-            ServiceConfig::Redis(c) => Some((c.port, c.id.clone())),
-            ServiceConfig::ZooKeeper(c) => Some((c.port, c.id.clone())),
-            ServiceConfig::AwsS3(_) => None,
-            ServiceConfig::Opendal(_) => None,
-            ServiceConfig::RedPanda(_) => None,
-        };
-
-        if let Some(x) = listen_info {
-            ports.push(x);
+        if let Some(port) = service.port() {
+            ports.push((port, service.id().to_string(), service.user_managed()));
         }
     }
 

src/risedevtool/src/service_config.rs

@@ -380,6 +380,50 @@ impl ServiceConfig {
             Self::Opendal(c) => &c.id,
         }
     }
+
+    pub fn port(&self) -> Option<u16> {
+        match self {
+            Self::ComputeNode(c) => Some(c.port),
+            Self::MetaNode(c) => Some(c.port),
+            Self::Frontend(c) => Some(c.port),
+            Self::Compactor(c) => Some(c.port),
+            Self::Minio(c) => Some(c.port),
+            Self::Etcd(c) => Some(c.port),
+            Self::Sqlite(_) => None,
+            Self::Prometheus(c) => Some(c.port),
+            Self::Grafana(c) => Some(c.port),
+            Self::Tempo(c) => Some(c.port),
+            Self::AwsS3(_) => None,
+            Self::ZooKeeper(c) => Some(c.port),
+            Self::Kafka(c) => Some(c.port),
+            Self::Pubsub(c) => Some(c.port),
+            Self::Redis(c) => Some(c.port),
+            Self::RedPanda(_c) => None,
+            Self::Opendal(_) => None,
+        }
+    }
+
+    pub fn user_managed(&self) -> bool {
+        match self {
+            Self::ComputeNode(c) => c.user_managed,
+            Self::MetaNode(c) => c.user_managed,
+            Self::Frontend(c) => c.user_managed,
+            Self::Compactor(c) => c.user_managed,
+            Self::Minio(_c) => false,
+            Self::Etcd(_c) => false,
+            Self::Sqlite(_c) => false,
+            Self::Prometheus(_c) => false,
+            Self::Grafana(_c) => false,
+            Self::Tempo(_c) => false,
+            Self::AwsS3(_c) => false,
+            Self::ZooKeeper(_c) => false,
+            Self::Kafka(c) => c.user_managed,
+            Self::Pubsub(_c) => false,
+            Self::Redis(_c) => false,
+            Self::RedPanda(_c) => false,
+            Self::Opendal(_c) => false,
+        }
+    }
 }
 
 mod string {

src/risedevtool/src/task/ensure_stop_service.rs

@@ -17,11 +17,12 @@ use anyhow::Result;
 use super::{ExecuteContext, Task};
 
 pub struct EnsureStopService {
-    ports: Vec<(u16, String)>,
+    /// `(port, id, user_managed)`
+    ports: Vec<(u16, String, bool)>,
 }
 
 impl EnsureStopService {
-    pub fn new(ports: Vec<(u16, String)>) -> Result<Self> {
+    pub fn new(ports: Vec<(u16, String, bool)>) -> Result<Self> {
         Ok(Self { ports })
     }
 }
@@ -30,16 +31,16 @@ impl Task for EnsureStopService {
     fn execute(&mut self, ctx: &mut ExecuteContext<impl std::io::Write>) -> anyhow::Result<()> {
         ctx.service(self);
 
-        for (port, service) in &self.ports {
-            // Do not require stopping kafka services
-            if service.starts_with("kafka") {
+        for (port, service_id, user_managed) in &self.ports {
+            // Do not require stopping user-managed services
+            if *user_managed {
                 continue;
             }
             let address = format!("127.0.0.1:{}", port);
 
             ctx.pb.set_message(format!(
                 "waiting for port close - {} (will be used by {})",
-                address, service
+                address, service_id
             ));
             ctx.wait_tcp_close(&address)?;
         }