Remove hacky sleep when port forwarding

- Instead of waiting until port forwarding, we read from stderr until we know for sure port forwarding has started
- Renamed `create_kubectl_port_forwarder` to `create_pg_wire_port_forwarder` since we'll be port forwarding more than once
- Removed the retry mechanism given we don't see the port forwarding disconnecting too often and we can't do it AND initial acknowledgement together easily
- Tie child process to a struct instead of a tokio spawn.

Author: SangJunBak

src/mz-debug/src/kubectl_port_forwarder.rs

@@ -19,14 +19,12 @@ use anyhow::Result;
 use k8s_openapi::api::core::v1::Service;
 use kube::api::ListParams;
 use kube::{Api, Client};
+use tokio::io::AsyncBufReadExt;
 
-use std::time::Duration;
-
-use mz_ore::retry::{self, RetryResult};
-use tracing::{info, warn};
+use tracing::info;
 
 use crate::SelfManagedDebugMode;
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct KubectlPortForwarder {
     pub namespace: String,
     pub service_name: String,
@@ -36,89 +34,73 @@ pub struct KubectlPortForwarder {
     pub context: Option<String>,
 }
 
-impl KubectlPortForwarder {
-    /// Port forwards a given k8s service via Kubectl.
-    /// The process will retry if the port-forwarding fails and
-    /// will terminate once the port forwarding reaches the max number of retries.
-    /// We retry since kubectl port-forward is flaky.
-    pub async fn port_forward(&self) {
-        if let Err(err) = retry::Retry::default()
-            .max_duration(Duration::from_secs(60))
-            .retry_async(|retry_state| {
-                let k8s_context = self.context.clone();
-                let namespace = self.namespace.clone();
-                let service_name = self.service_name.clone();
-                let local_address = self.local_address.clone();
-                let local_port = self.local_port;
-                let target_port = self.target_port;
-
-                info!(
-                    "Spawning port forwarding process for {} from ports {}:{} -> {}",
-                    service_name, local_address, local_port, target_port
-                );
+pub struct PortForwardConnection {
+    // tokio process that's killed on drop
+    pub _port_forward_process: tokio::process::Child,
+}
 
-                async move {
-                    let port_arg_str = format!("{}:{}", &local_port, &target_port);
-                    let service_name_arg_str = format!("services/{}", &service_name);
-                    let mut args = vec![
-                        "port-forward",
-                        &service_name_arg_str,
-                        &port_arg_str,
-                        "-n",
-                        &namespace,
-                        "--address",
-                        &local_address,
-                    ];
-
-                    if let Some(k8s_context) = &k8s_context {
-                        args.extend(["--context", k8s_context]);
-                    }
+impl KubectlPortForwarder {
+    /// Spawns a port forwarding process that resolves when
+    /// the port forward is established.
+    pub async fn spawn_port_forward(&mut self) -> Result<PortForwardConnection, anyhow::Error> {
+        let port_arg_str = format!("{}:{}", &self.local_port, &self.target_port);
+        let service_name_arg_str = format!("services/{}", &self.service_name);
+        let mut args = vec![
+            "port-forward",
+            &service_name_arg_str,
+            &port_arg_str,
+            "-n",
+            &self.namespace,
+            "--address",
+            &self.local_address,
+        ];
+
+        if let Some(k8s_context) = &self.context {
+            args.extend(["--context", k8s_context]);
+        }
 
-                    match tokio::process::Command::new("kubectl")
-                        .args(args)
-                        // Silence stdout/stderr
-                        .stdout(std::process::Stdio::null())
-                        .stderr(std::process::Stdio::null())
-                        .kill_on_drop(true)
-                        .output()
-                        .await
-                    {
-                        Ok(output) => {
-                            if !output.status.success() {
-                                let retry_err_msg = format!(
-                                    "Failed to port-forward{}: {}",
-                                    retry_state.next_backoff.map_or_else(
-                                        || "".to_string(),
-                                        |d| format!(", retrying in {:?}", d)
-                                    ),
-                                    String::from_utf8_lossy(&output.stderr)
-                                );
-                                warn!("{}", retry_err_msg);
-
-                                return RetryResult::RetryableErr(anyhow::anyhow!(retry_err_msg));
-                            }
-                        }
-                        Err(err) => {
-                            return RetryResult::RetryableErr(anyhow::anyhow!(
-                                "Failed to port-forward: {}",
-                                err
-                            ));
+        let child = tokio::process::Command::new("kubectl")
+            .args(args)
+            // Silence stdout
+            .stdout(std::process::Stdio::null())
+            .stderr(std::process::Stdio::piped())
+            .kill_on_drop(true)
+            .spawn();
+
+        if let Ok(mut child) = child {
+            if let Some(stderr) = child.stderr.take() {
+                let stderr_reader = tokio::io::BufReader::new(stderr);
+                // Wait until we know port forwarding is established
+                let timeout = tokio::time::timeout(std::time::Duration::from_secs(5), async {
+                    let mut lines = stderr_reader.lines();
+                    while let Ok(Some(line)) = lines.next_line().await {
+                        if line.contains("Forwarding from") {
+                            break;
                         }
                     }
-                    // The kubectl subprocess's future will only resolve on error, thus the
-                    // code here is unreachable. We return RetryResult::Ok to satisfy
-                    // the type checker.
-                    RetryResult::Ok(())
+                })
+                .await;
+
+                if timeout.is_err() {
+                    return Err(anyhow::anyhow!("Port forwarding timed out after 5 seconds"));
                 }
-            })
-            .await
-        {
-            warn!("{}", err);
+
+                info!(
+                    "Port forwarding established for {} from ports {}:{} -> {}",
+                    &self.service_name, &self.local_address, &self.local_port, &self.target_port
+                );
+
+                return Ok(PortForwardConnection {
+                    _port_forward_process: child,
+                });
+            }
         }
+        Err(anyhow::anyhow!("Failed to spawn port forwarding process"))
     }
 }
 
-pub async fn create_kubectl_port_forwarder(
+/// Creates a port forwarder for the external pg wire port of balancerd.
+pub async fn create_pg_wire_port_forwarder(
     client: &Client,
     args: &SelfManagedDebugMode,
 ) -> Result<KubectlPortForwarder, anyhow::Error> {

src/mz-debug/src/main.rs

@@ -19,15 +19,14 @@ use kube::{Client, Config};
 use mz_build_info::{BuildInfo, build_info};
 use mz_ore::cli::{self, CliConfig};
 use mz_ore::error::ErrorExt;
-use mz_ore::task;
 use tracing::{error, info, warn};
 use tracing_subscriber::EnvFilter;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::util::SubscriberInitExt;
 
 use crate::docker_dumper::DockerDumper;
 use crate::k8s_dumper::K8sDumper;
-use crate::kubectl_port_forwarder::create_kubectl_port_forwarder;
+use crate::kubectl_port_forwarder::create_pg_wire_port_forwarder;
 use crate::utils::{
     create_tracing_log_file, format_base_path, validate_pg_connection_string, zip_debug_folder,
 };
@@ -63,6 +62,8 @@ pub struct SelfManagedDebugMode {
     /// If true, the tool will dump the values of secrets in the Kubernetes cluster.
     #[clap(long, default_value = "false", action = clap::ArgAction::Set)]
     k8s_dump_secret_values: bool,
+    // TODO (debug_tool1): Convert port forwarding variables into a map since we'll be
+    // portforwarding multiple times
     /// If true, the tool will automatically port-forward the external SQL port in the Kubernetes cluster.
     /// If dump_k8s is false however, we will not automatically port-forward.
     #[clap(long, default_value = "true", action = clap::ArgAction::Set)]
@@ -221,8 +222,9 @@ async fn main() {
 
 async fn run(context: Context, args: Args) -> Result<(), anyhow::Error> {
     // Depending on if the user is debugging either a k8s environments or docker environment,
-    // dump the respective system's resources.
-    let container_system_dumper = match &args.debug_mode {
+    // dump the respective system's resources and spin up a port forwarder if auto port forwarding.
+    // We'd like to keep the port forwarder alive until the end of the program.
+    let (container_system_dumper, _pg_wire_port_forward_connection) = match &args.debug_mode {
         DebugMode::SelfManaged(args) => {
             if args.dump_k8s {
                 let client = match create_k8s_client(args.k8s_context.clone()).await {
@@ -232,27 +234,30 @@ async fn run(context: Context, args: Args) -> Result<(), anyhow::Error> {
                         return Err(e);
                     }
                 };
-
-                if args.auto_port_forward {
-                    let port_forwarder = create_kubectl_port_forwarder(&client, args).await?;
-                    task::spawn(|| "port-forwarding", async move {
-                        port_forwarder.port_forward().await;
-                    });
-                    // There may be a delay between when the port forwarding process starts and when it's ready
-                    // to use. We wait a few seconds to ensure that port forwarding is ready.
-                    tokio::time::sleep(std::time::Duration::from_secs(2)).await;
-                }
-
+                let port_forward_connection = if args.auto_port_forward {
+                    let mut port_forwarder = create_pg_wire_port_forwarder(&client, args).await?;
+
+                    match port_forwarder.spawn_port_forward().await {
+                        Ok(connection) => Some(connection),
+                        Err(err) => {
+                            warn!("{}", err);
+                            None
+                        }
+                    }
+                } else {
+                    None
+                };
+                // Parker: Return it as a tuple (dumper, port_forwarder)
                 let dumper = ContainerServiceDumper::new_k8s_dumper(
                     &context,
                     client,
                     args.k8s_namespaces.clone(),
                     args.k8s_context.clone(),
                     args.k8s_dump_secret_values,
                 );
-                Some(dumper)
+                (Some(dumper), port_forward_connection)
             } else {
-                None
+                (None, None)
             }
         }
         DebugMode::Emulator(args) => {
@@ -263,9 +268,9 @@ async fn run(context: Context, args: Args) -> Result<(), anyhow::Error> {
                     .expect("docker_container_id is required");
                 let dumper =
                     ContainerServiceDumper::new_docker_dumper(&context, docker_container_id);
-                Some(dumper)
+                (Some(dumper), None)
             } else {
-                None
+                (None, None)
             }
         }
     };