Enable dumping logs via streaming, cancellable via ctrl-c

Author: agg23

src/adb.rs

@@ -1,10 +1,12 @@
+use crate::logs::LineBuffer;
 use crate::{InstallerError, Result};
 use adb_client::{ADBDeviceExt, ADBServer, ADBUSBDevice};
+use std::io::Write;
 use std::net::{Ipv4Addr, SocketAddrV4};
 use std::path::Path;
 
 pub struct AdbManager {
-    device: Box<dyn ADBDeviceExt>,
+    device: Box<dyn ADBDeviceExt + Send>,
 }
 
 impl AdbManager {
@@ -126,6 +128,18 @@ impl AdbManager {
         Ok(output_str.trim().to_string())
     }
 
+    pub fn shell_stream<T>(&mut self, command: &str, writer: T) -> Result<()>
+    where
+        T: Write,
+    {
+        let cmd_parts: Vec<&str> = command.split_whitespace().collect();
+        let mut line_buffer = LineBuffer::new(writer);
+
+        self.device
+            .shell_command(&cmd_parts, &mut line_buffer)
+            .map_err(|e| InstallerError::Adb(format!("Failed to run shell command: {}", e)))
+    }
+
     pub async fn push_file(&mut self, local: &Path, remote: &str) -> Result<()> {
         let mut file = std::fs::File::open(local)
             .map_err(|e| InstallerError::Adb(format!("Failed to open file: {}", e)))?;

src/logs.rs

@@ -1,28 +1,130 @@
+use std::fs::File;
+use std::io::{self, Write};
 use std::time::{SystemTime, UNIX_EPOCH};
 
-use tokio::fs;
+use tokio::signal::ctrl_c;
+use tokio::sync::watch::{self, Sender};
+use tokio::task::spawn_blocking;
 
-use crate::error::Result;
-use crate::AdbManager;
+use crate::{AdbManager, InstallerError};
 
-// TODO: Implement streaming
-pub async fn dump_logcat(stream: bool, remote_auth_url: Option<String>) -> Result<()> {
-    let mut adb = AdbManager::connect(remote_auth_url).await?;
+// Taken from adb_client LogFilter
+pub struct LineBuffer<W: Write> {
+    writer: W,
+    buffer: Vec<u8>,
+}
+
+impl<W: Write> LineBuffer<W> {
+    pub fn new(writer: W) -> Self {
+        LineBuffer {
+            writer,
+            buffer: Vec::new(),
+        }
+    }
+
+    fn should_write(&self, _line: &[u8]) -> bool {
+        true
+    }
+}
+
+impl<W: Write> Write for LineBuffer<W> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.buffer.extend_from_slice(buf);
+
+        let buf_clone = self.buffer.clone();
+        let mut lines = buf_clone.split_inclusive(|&byte| byte == b'\n').peekable();
+
+        while let Some(line) = lines.next() {
+            if lines.peek().is_some() {
+                if self.should_write(line) {
+                    self.writer.write_all(line)?;
+                }
+            } else {
+                // This is the last (unfinished) element, we keep it for next round
+                self.buffer = line.to_vec();
+                break;
+            }
+        }
+
+        Ok(buf.len())
+    }
 
-    let result = adb.shell("logcat -d").await?;
+    fn flush(&mut self) -> io::Result<()> {
+        self.writer.flush()
+    }
+}
+
+struct PrintFileWriter {
+    file: File,
+    line_count: usize,
+    tx: Sender<usize>,
+}
+
+impl Write for PrintFileWriter {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        // This isn't correct if there needs to be a retry, but assume it just works
+        let result = self.file.write(buf);
+        if let Ok(string) = String::from_utf8(buf.into()) {
+            println!("{string}");
+        }
+        self.line_count += 1;
+        let _ = self.tx.send(self.line_count);
 
-    let line_count = result.split("\n").count();
+        result
+    }
 
+    fn flush(&mut self) -> io::Result<()> {
+        todo!()
+    }
+}
+
+pub async fn dump_logcat_and_exit(stream: bool, remote_auth_url: Option<String>) {
     let timestamp = SystemTime::now()
         .duration_since(UNIX_EPOCH)
         .unwrap()
         .as_millis();
 
     let filename = format!("penumbra_log_dump_{timestamp}.log");
 
-    fs::write(&filename, result).await?;
+    let inner_filename = filename.clone();
+    match tokio::spawn(async move {
+        let mut adb = AdbManager::connect(remote_auth_url).await?;
+
+        let mut file = File::create(inner_filename)?;
+
+        if stream {
+            let (tx, rx) = watch::channel(0);
+            spawn_blocking(move || {
+                let mut writer = PrintFileWriter {
+                    file,
+                    line_count: 0,
+                    tx,
+                };
 
-    println!("Wrote {line_count} lines to {filename}");
+                let _ = adb.shell_stream("logcat", &mut writer);
+            });
 
-    Ok(())
+            let _ = ctrl_c().await;
+            let value = *rx.borrow();
+            Ok::<usize, InstallerError>(value)
+        } else {
+            let result = adb.shell("logcat -d").await?;
+            let line_count = result.split("\n").count();
+            file.write_all(result.as_bytes())?;
+            Ok(line_count)
+        }
+    })
+    .await
+    .unwrap()
+    {
+        Ok(line_count) => {
+            println!("\n\nWrote {line_count} lines to {filename}");
+            // Forcibly close process due to blocking call staying open
+            std::process::exit(0);
+        }
+        Err(err) => {
+            println!("Error: {err}");
+            std::process::exit(1);
+        }
+    };
 }

src/main.rs

@@ -4,7 +4,7 @@ use std::path::PathBuf;
 use tokio;
 
 use penumbra_installer::{
-    logs::dump_logcat, ConfigLoader, InstallationEngine, InstallerError, Result,
+    logs::dump_logcat_and_exit, ConfigLoader, InstallationEngine, InstallerError, Result,
 };
 
 #[derive(Parser)]
@@ -200,10 +200,7 @@ async fn main() -> Result<()> {
         Commands::DumpLogs {
             stream,
             remote_auth_url,
-        } => match dump_logcat(stream, remote_auth_url).await {
-            Ok(_) => {}
-            Err(err) => println!("Dump error: {err}"),
-        },
+        } => dump_logcat_and_exit(stream, remote_auth_url).await,
     }
 
     Ok(())