Merge pull request #67 from Dridus/rmm/windows-io-hang

work around hang starting database on windows

Author: brianheineman

postgresql_commands/src/traits.rs

@@ -160,28 +160,103 @@ impl CommandExecutor for std::process::Command {
 impl AsyncCommandExecutor for tokio::process::Command {
     /// Execute the command and return the stdout and stderr
     async fn execute(&mut self, timeout: Option<Duration>) -> Result<(String, String)> {
+        #[tracing::instrument(level = "debug", skip(reader))]
+        async fn read_process(
+            mut reader: Option<impl tokio::io::AsyncRead + Unpin>,
+            mut exit_anyway_broadcast_receiver: tokio::sync::broadcast::Receiver<()>,
+        ) -> Result<String> {
+            let Some(reader) = reader.as_mut() else {
+                return Ok(String::new());
+            };
+            let mut vec = Vec::new();
+            loop {
+                use tokio::io::AsyncReadExt as _;
+                tokio::select! {
+                    n = reader.read_buf(&mut vec) => {
+                        if n? == 0 {
+                            return Ok(String::from_utf8_lossy(&vec).into_owned());
+                        }
+                    },
+                    _ = exit_anyway_broadcast_receiver.recv() => {
+                        return Ok(String::from_utf8_lossy(&vec).into_owned());
+                    },
+                }
+            }
+        }
+
         debug!("Executing command: {}", self.to_command_string());
-        let output = match timeout {
-            Some(duration) => tokio::time::timeout(duration, self.output()).await?,
-            None => self.output().await,
-        }?;
 
-        let stdout = String::from_utf8_lossy(&output.stdout).into_owned();
-        let stderr = String::from_utf8_lossy(&output.stderr).into_owned();
-        debug!(
-            "Result: {}\nstdout: {}\nstderr: {}",
-            output
-                .status
-                .code()
-                .map_or("None".to_string(), |c| c.to_string()),
-            stdout,
-            stderr
-        );
+        let res_fut = async {
+            let mut child = self
+                .stdout(std::process::Stdio::piped())
+                .stderr(std::process::Stdio::piped())
+                .spawn()?;
+
+            let stdout = child.stdout.take();
+            let stderr = child.stderr.take();
+
+            // on windows, pg_ctl start will appear to hang if you try to read out all of stdout
+            // and stderr. so, on windows do a horrible hack and forcibly end reading of stdout
+            // and stderr 50ms after the process exits. on not-windows, this early exit mechanism
+            //  is set up but never sent to, resulting in the same behavior as `read_to_end`.
+
+            let (exit_anyway_broadcast_sender, exit_anyway_broadcast_receiver_stdout) =
+                tokio::sync::broadcast::channel(1);
+            let exit_anyway_broadcast_receiver_stderr = exit_anyway_broadcast_sender.subscribe();
+            let stdout = tokio::spawn(async {
+                read_process(stdout, exit_anyway_broadcast_receiver_stdout).await
+            });
+            let stderr = tokio::spawn(async {
+                read_process(stderr, exit_anyway_broadcast_receiver_stderr).await
+            });
+            let exit_status = child.wait().await;
+            #[cfg(target_os = "windows")]
+            {
+                tokio::time::sleep(Duration::from_millis(50)).await;
+                let _ = exit_anyway_broadcast_sender.send(());
+            }
+            let (stdout, stderr) = tokio::join!(stdout, stderr);
+            std::mem::drop(exit_anyway_broadcast_sender);
+
+            let exit_status = exit_status?;
+            fn debug_render(
+                which: &'static str,
+                res: &std::result::Result<Result<String>, tokio::task::JoinError>,
+            ) -> String {
+                match res {
+                    Ok(Ok(s)) => s.into(),
+                    Ok(Err(io_err)) => format!("<failed to read {}: {:?}>", which, io_err),
+                    Err(join_err) => format!("<failed to read {}: {:?}>", which, join_err),
+                }
+            }
+            debug!(
+                "Result: {}\nstdout: {}\nstderr: {}",
+                exit_status
+                    .code()
+                    .map_or("None".to_string(), |c| c.to_string()),
+                debug_render("stdout", &stdout),
+                debug_render("stderr", &stderr)
+            );
+
+            fn unwrap2_or_empty_string<E, E2>(
+                r: std::result::Result<std::result::Result<String, E>, E2>,
+            ) -> String {
+                r.map_or_else(|_| String::new(), |r| r.unwrap_or_else(|_| String::new()))
+            }
 
-        if output.status.success() {
-            Ok((stdout, stderr))
-        } else {
-            Err(Error::CommandError { stdout, stderr })
+            let stdout = unwrap2_or_empty_string(stdout);
+            let stderr = unwrap2_or_empty_string(stderr);
+
+            if exit_status.success() {
+                Ok((stdout, stderr))
+            } else {
+                Err(Error::CommandError { stdout, stderr })
+            }
+        };
+
+        match timeout {
+            Some(duration) => tokio::time::timeout(duration, res_fut).await?,
+            None => res_fut.await,
         }
     }
 }