feat: support control command

Cargo.toml

@@ -4,23 +4,26 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
-bytes = "1"
-clap = { version = "4.5", default-features = false, features = [
+bytes = "1.12"
+clap = { version = "4.6", default-features = false, features = [
   "std",
   "help",
   "usage",
   "error-context",
   "derive",
 ] }
-dashmap = "6"
+dashmap = "6.2"
+futures-util = { version = "0.3", default-features = false, features = [
+  "sink",
+] }
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 snafu = "0.9"
-tokio = { version = "1", features = ["full"] }
+tokio = { version = "1.52", features = ["full"] }
 tokio-stream = "0.1"
 tokio-util = { version = "0.7", features = ["codec"] }
 toml = "0.8"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 tracing-appender = "0.2"
-uuid = { version = "1", features = ["v4"] }
+uuid = { version = "1.23", features = ["v4"] }

src/client.rs

@@ -1,14 +1,21 @@
+use futures_util::SinkExt;
 use snafu::ResultExt;
 use tokio::{
-    io::{self, AsyncReadExt, AsyncWriteExt},
+    io::{self, AsyncWriteExt},
     net::{
         TcpStream,
         tcp::{OwnedReadHalf, OwnedWriteHalf},
     },
 };
+use tokio_stream::StreamExt;
+use tokio_util::codec::{FramedRead, FramedWrite};
 use tracing::{debug, info, warn};
 
-use crate::error::Result;
+use crate::error::{PlainTextSnafu, Result};
+use crate::protocol::{
+    ControlMessage, InstanceStatus, LspFrameDecoder, LspFrameStream, RadFrameCocdec,
+    RadFrameStream, RadMessage, ServerStatus,
+};
 use crate::{config::DEFAULT_ADDR, error::IoSnafu};
 
 pub struct Options {
@@ -34,10 +41,10 @@ pub async fn run(opts: Options) -> Result<()> {
 
     info!(server_addr, "client proxy connected to rad server");
 
-    let (read, write) = stream.into_split();
+    let (r, w) = stream.into_split();
 
-    let upstream = tokio::spawn(stdin_to_server(write));
-    let downstream = tokio::spawn(server_to_stdout(read));
+    let upstream = tokio::spawn(stdin_to_server(w));
+    let downstream = tokio::spawn(server_to_stdout(r));
 
     tokio::select! {
         _ = upstream => {}
@@ -49,52 +56,157 @@ pub async fn run(opts: Options) -> Result<()> {
     Ok(())
 }
 
+pub async fn status(opts: Options) -> Result<()> {
+    let Options { server_addr } = opts;
+
+    let stream = TcpStream::connect(&server_addr)
+        .await
+        .with_context(|_| IoSnafu {
+            detail: format!("failed to connect to red server, server addr: {server_addr}"),
+        })?;
+
+    let msg = RadMessage::control(ControlMessage::StatusRequest);
+    let (r, w) = stream.into_split();
+
+    let mut sink = FramedWrite::new(w, RadFrameCocdec);
+    sink.send(msg).await?;
+
+    let mut stream = FramedRead::new(r, RadFrameCocdec);
+    let Some(msg) = stream.next().await else {
+        return PlainTextSnafu {
+            msg: "rad server closed connection before status response",
+        }
+        .fail();
+    };
+
+    match msg? {
+        RadMessage::Control(ControlMessage::StatusResponse { status }) => {
+            print!("{}", format_status(status));
+            Ok(())
+        }
+        RadMessage::Control(ControlMessage::Error { message }) => {
+            PlainTextSnafu { msg: message }.fail()
+        }
+        RadMessage::Control(ControlMessage::StatusRequest) => PlainTextSnafu {
+            msg: "unexpected status request from rad server".to_string(),
+        }
+        .fail(),
+        RadMessage::Lsp(_) => PlainTextSnafu {
+            msg: "unexpected lsp message from rad server".to_string(),
+        }
+        .fail(),
+    }
+}
+
+fn format_status(mut status: ServerStatus) -> String {
+    if status.instances.is_empty() {
+        return "no lsp instances\n".to_string();
+    }
+
+    status.instances.sort_by(|a, b| {
+        a.workspace
+            .cmp(&b.workspace)
+            .then_with(|| a.pid.cmp(&b.pid))
+    });
+
+    let mut out = String::new();
+    for (idx, instance) in status.instances.iter().enumerate() {
+        if idx > 0 {
+            out.push('\n');
+        }
+        out.push_str(&format_instance_status(instance));
+    }
+    out
+}
+
+fn format_instance_status(instance: &InstanceStatus) -> String {
+    format!(
+        "workspace: {}\n  pid:      {}\n  clients:  {}\n  idle:     {}\n  healthy:  {}\n",
+        instance.workspace,
+        instance.pid,
+        instance.client_count,
+        format_duration_secs(instance.idle_secs),
+        if instance.healthy { "yes" } else { "no" },
+    )
+}
+
+fn format_duration_secs(secs: i64) -> String {
+    let secs = secs.max(0);
+    if secs < 60 {
+        return format!("{secs}s");
+    }
+
+    let minutes = secs / 60;
+    let seconds = secs % 60;
+    if minutes < 60 {
+        return format!("{minutes}m {seconds}s");
+    }
+
+    let hours = minutes / 60;
+    let minutes = minutes % 60;
+    format!("{hours}h {minutes}m")
+}
+
 async fn stdin_to_server(mut write: OwnedWriteHalf) {
-    let mut stdin = io::stdin();
-    let mut buf = vec![0; 8192];
+    let stdin = io::stdin();
+    let mut stream = LspFrameStream::new(stdin, LspFrameDecoder);
 
-    loop {
-        let n = match stdin.read(&mut buf).await {
-            Ok(0) => {
-                debug!("stdin reached eof");
+    while let Some(frame) = stream.next().await {
+        let frame = match frame {
+            Ok(frame) => frame,
+            Err(e) => {
+                warn!(error = ?e, "failed to decode lsp frame from stdin");
                 break;
             }
-            Ok(n) => n,
+        };
+        let bytes = match RadMessage::lsp(frame).to_bytes() {
+            Ok(bytes) => bytes,
             Err(e) => {
-                warn!(error = ?e, "failed to read from stdin");
+                warn!(error = ?e, "failed to encode rad lsp message");
                 break;
             }
         };
 
-        if let Err(e) = write.write_all(&buf[..n]).await {
+        if let Err(e) = write.write_all(&bytes).await {
             warn!(error = ?e, "failed to write to rad server");
             break;
         }
     }
 
+    debug!("stdin reached eof");
+
     if let Err(e) = write.shutdown().await {
         warn!(error = ?e, "failed to shutdown write");
     }
 }
 
-async fn server_to_stdout(mut read: OwnedReadHalf) {
+async fn server_to_stdout(read: OwnedReadHalf) {
     let mut stdout = io::stdout();
-    let mut buf = vec![0; 8192];
+    let mut stream = RadFrameStream::new(read, RadFrameCocdec);
 
-    loop {
-        let n = match read.read(&mut buf).await {
-            Ok(0) => {
-                debug!("rad server closed connection");
+    while let Some(msg) = stream.next().await {
+        let msg = match msg {
+            Ok(msg) => msg,
+            Err(e) => {
+                warn!(error = ?e, "failed to read rad message from server");
                 break;
             }
-            Ok(n) => n,
+        };
+
+        let RadMessage::Lsp(frame) = msg else {
+            warn!("ignoring unexpected control message from rad server");
+            continue;
+        };
+
+        let bytes = match frame.to_bytes() {
+            Ok(bytes) => bytes,
             Err(e) => {
-                warn!(error = ?e, "failed to read from rad server");
+                warn!(error = ?e, "failed to encode lsp frame for stdout");
                 break;
             }
         };
 
-        if let Err(e) = stdout.write_all(&buf[..n]).await {
+        if let Err(e) = stdout.write_all(&bytes).await {
             warn!(error = ?e, "failed to write to stdout");
             break;
         }
@@ -104,4 +216,67 @@ async fn server_to_stdout(mut read: OwnedReadHalf) {
             break;
         }
     }
+
+    debug!("rad server closed connection");
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn formats_empty_status() {
+        let status = ServerStatus {
+            instances: Vec::new(),
+        };
+
+        assert_eq!("no lsp instances\n", format_status(status));
+    }
+
+    #[test]
+    fn formats_status_by_workspace() {
+        let status = ServerStatus {
+            instances: vec![
+                InstanceStatus {
+                    workspace: "file:///z".to_string(),
+                    pid: 20,
+                    client_count: 2,
+                    idle_secs: 75,
+                    healthy: false,
+                },
+                InstanceStatus {
+                    workspace: "file:///a".to_string(),
+                    pid: 10,
+                    client_count: 1,
+                    idle_secs: 5,
+                    healthy: true,
+                },
+            ],
+        };
+
+        assert_eq!(
+            concat!(
+                "workspace: file:///a\n",
+                "  pid:      10\n",
+                "  clients:  1\n",
+                "  idle:     5s\n",
+                "  healthy:  yes\n",
+                "\n",
+                "workspace: file:///z\n",
+                "  pid:      20\n",
+                "  clients:  2\n",
+                "  idle:     1m 15s\n",
+                "  healthy:  no\n",
+            ),
+            format_status(status)
+        );
+    }
+
+    #[test]
+    fn formats_duration_secs() {
+        assert_eq!("0s", format_duration_secs(-1));
+        assert_eq!("59s", format_duration_secs(59));
+        assert_eq!("1m 0s", format_duration_secs(60));
+        assert_eq!("1h 1m", format_duration_secs(3661));
+    }
 }