add rpc implementation

Author: Panaetius

Cargo.lock

@@ -84,6 +84,13 @@ iroh = "0.95.1"
 rand = "0.9.2"
 regex = "1.12.2"
 sha2 = "0.10.9"
+tarpc = { version = "0.37.0", features = [
+  "serde-transport",
+  "serde-transport-bincode",
+  "tcp",
+  "tokio1",
+] }
+tokio-duplex = "1.0.1"
 
 [build-dependencies]
 anyhow = "1.0.90"

coman/Cargo.toml

@@ -303,6 +303,14 @@ pub enum CscsJobCommands {
         #[clap(help="id or name of the job (name uses newest job of that name)",  add = ArgValueCompleter::new(job_id_or_name_completer))]
         job: JobIdOrName,
     },
+    #[clap(
+        alias = "ru",
+        about = "show current resource usage of the job, needs coman to be injected in the session [aliases: ru]"
+    )]
+    ResourceUsage {
+        #[clap(help="id or name of the job (name uses newest job of that name)",  add = ArgValueCompleter::new(job_id_or_name_completer))]
+        job: JobIdOrName,
+    },
 }
 fn job_id_or_name_completer(current: &std::ffi::OsStr) -> Vec<CompletionCandidate> {
     let mut completions = vec![];

coman/src/cli/app.rs

@@ -7,10 +7,13 @@ use iroh::{
     endpoint::ConnectionError,
     protocol::{ProtocolHandler, Router},
 };
+use nom::AsBytes;
 use pid1::Pid1Settings;
 use rust_supervisor::{ChildType, Supervisor, SupervisorConfig};
 use tokio::{io::AsyncWriteExt, net::TcpStream};
 
+use crate::cli::rpc::RpcHandler;
+
 const SECRET_KEY_ENV: &str = "COMAN_IROH_SECRET";
 const PORT_FORWARD_ENV: &str = "COMAN_FORWARDED_PORTS";
 const SSH_PORT: u16 = 15263;
@@ -121,6 +124,10 @@ async fn port_forward() -> Result<()> {
         builder = builder.accept(alpn.clone().into_bytes(), handler);
         println!("set up port forwarding for port {port} ({alpn})");
     }
+
+    // add rpc server
+    let rpc_handler = RpcHandler;
+    builder = builder.accept(b"/coman/rpc/".as_bytes(), rpc_handler);
     let _router = builder.spawn();
     println!("port forwarding started");
 

coman/src/cli/exec.rs

@@ -1,3 +1,4 @@
 pub mod app;
 pub mod exec;
 pub mod proxy;
+pub mod rpc;

coman/src/cli/mod.rs

@@ -0,0 +1,52 @@
+use futures::StreamExt;
+use iroh::protocol::ProtocolHandler;
+use tarpc::{
+    serde_transport as transport, server, server::Channel, tokio_serde::formats::Bincode,
+    tokio_util::codec::LengthDelimitedCodec,
+};
+use tokio_duplex::Duplex;
+
+use crate::cli::app::COMAN_VERSION;
+
+#[tarpc::service]
+pub trait ComanRPC {
+    async fn version() -> String;
+}
+#[derive(Debug, Clone)]
+struct RpcServer;
+
+impl ComanRPC for RpcServer {
+    async fn version(self, _: tarpc::context::Context) -> String {
+        COMAN_VERSION.to_string()
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct RpcHandler;
+
+impl ProtocolHandler for RpcHandler {
+    async fn accept(&self, connection: iroh::endpoint::Connection) -> Result<(), iroh::protocol::AcceptError> {
+        let endpoint_id = connection.remote_id();
+        match connection.accept_bi().await {
+            Ok((iroh_send, iroh_recv)) => {
+                println!("Accepted bidirectional stream from {endpoint_id}");
+                let codec_builder = LengthDelimitedCodec::builder();
+                let combined = Duplex::new(iroh_recv, iroh_send);
+                let framed = codec_builder.new_framed(combined);
+
+                let transport = transport::new(framed, Bincode::default());
+                let server = server::BaseChannel::with_defaults(transport);
+                tokio::spawn(server.execute(RpcServer.serve()).for_each(spawn));
+            }
+            Err(e) => {
+                println!("Failed to accept bidirectional stream to rpc: {e}");
+            }
+        }
+
+        Ok(())
+    }
+}
+
+async fn spawn(fut: impl Future<Output = ()> + Send + 'static) {
+    tokio::spawn(fut);
+}

coman/src/cli/rpc.rs

@@ -22,8 +22,8 @@ use crate::{
         api_client::{client::JobStartOptions, types::JobStatus},
         handlers::{
             cscs_file_delete, cscs_file_download, cscs_file_list, cscs_file_upload, cscs_job_cancel, cscs_job_details,
-            cscs_job_list, cscs_job_log, cscs_job_start, cscs_login, cscs_port_forward, cscs_system_list,
-            cscs_system_set,
+            cscs_job_list, cscs_job_log, cscs_job_start, cscs_login, cscs_port_forward, cscs_resource_usage,
+            cscs_system_list, cscs_system_set,
         },
     },
 };
@@ -136,6 +136,16 @@ pub(crate) async fn cli_cscs_port_forward(
     cscs_port_forward(job_id, source_port, destination_port, system).await
 }
 
+pub(crate) async fn cli_cscs_job_resource_usage(
+    job: JobIdOrName,
+    system: Option<String>,
+    platform: Option<ComputePlatform>,
+) -> Result<()> {
+    let job_id = maybe_job_id_from_name(job, system.clone(), platform.clone()).await?;
+    println!("running port forward for job {job_id}");
+    cscs_resource_usage(job_id, system).await
+}
+
 #[allow(clippy::too_many_arguments)]
 pub(crate) async fn cli_cscs_job_start(
     name: Option<String>,