feat: Add rudimentary request_reply api and example

Author: Leonidas Loucas

rust/lib/srpc/client/Cargo.toml

@@ -8,10 +8,12 @@ name = "srpc_client"
 crate-type = ["cdylib", "rlib"]
 
 [dependencies]
+async-trait = "0.1"
 bytes = "1"
 futures = "0.3"
 tokio = { version = "1", features = ["full"] }
 openssl = "0.10"
+serde = {version = "1", features = ["derive"]}
 serde_json = "1"
 tokio-openssl = "0.6"
 tokio-util = { version = "0.7", features = ["codec"] }

rust/lib/srpc/client/examples/python_client_example3.py

@@ -0,0 +1,70 @@
+"""
+This example demonstrates how to use the srpc_client Python bindings.
+
+To run this example:
+1. Build the Rust library: maturin build --features python
+2. Install the wheel: pip install target/wheels/srpc_client-*.whl
+3. Run this script: python examples/python_client_example.py
+"""
+
+import asyncio
+import json
+import os
+from srpc_client import SrpcClientConfig
+
+
+async def main():
+    print("Starting client..")
+
+    # Create a new ClientConfig instance
+    client = SrpcClientConfig(
+        os.environ["EXAMPLE_3_SRPC_SERVER_HOST"],
+        int(os.environ["EXAMPLE_3_SRPC_SERVER_PORT"]),
+        os.environ["EXAMPLE_3_SRPC_SERVER_ENPOINT"],
+        os.environ["EXAMPLE_3_SRPC_SERVER_CERT"],
+        os.environ["EXAMPLE_3_SRPC_SERVER_KEY"],
+    )
+
+    # Connect to the server
+    client = await client.connect()
+    print("Connected to server")
+
+    message = "Hypervisor.ListVMs\n"
+
+    # Send a message
+    print(f"Sending message: {message}")
+    await client.send_message(message)
+    print(f"Sent message: {message}")
+
+    # Receive an empty response
+    print("Waiting for empty string response...")
+    responses = await client.receive_message(expect_empty=True, should_continue=False)
+    async for response in responses:
+        print(f"Received response: {response}")
+
+    # Send a JSON message
+    payload = json.dumps({"IgnoreStateMask": 0, "OwnerGroups": [], "OwnerUsers": [], "Sort": True, "VmTagsToMatch": {}})
+    print(f"Sending payload: {payload}")
+    await client.send_json(payload)
+    print(f"Sent payload: {payload}")
+
+    # Receive an empty response
+    print("Waiting for empty string response for payload...")
+    responses = await client.receive_message(expect_empty=True, should_continue=False)
+    async for response in responses:
+        print(f"Received response: {response}")
+
+    # Receive responses
+    print("Waiting for response...")
+    responses = await client.receive_json_cb(should_continue=lambda _: False)
+    async for response in responses:
+        print(f"Received response: {json.loads(response)}")
+
+    # Use RequestReply
+    print(f"Sending request_reply: {message}")
+    res = await client.request_reply(message, json.dumps({"IgnoreStateMask": 0, "OwnerGroups": [], "OwnerUsers": [], "Sort": True, "VmTagsToMatch": {}}))
+    print(f"Sent request_reply: {message}, got reply: {res}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

rust/lib/srpc/client/examples/rust_client_example3.rs

@@ -0,0 +1,120 @@
+use std::{collections::HashMap, error::Error};
+
+use srpc_client::{ClientConfig, CustomError, ReceiveOptions, SimpleValue};
+use tracing::{error, info, level_filters::LevelFilter};
+use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    tracing_subscriber::registry()
+        .with(
+            tracing_subscriber::EnvFilter::builder()
+                .with_default_directive(LevelFilter::INFO.into())
+                .from_env_lossy(),
+        )
+        .with(tracing_subscriber::fmt::Layer::default().compact())
+        .init();
+
+    info!("Starting client...");
+
+    // Create a new ClientConfig instance
+    let config = ClientConfig::new(
+        &std::env::var("EXAMPLE_3_SRPC_SERVER_HOST")?,
+        std::env::var("EXAMPLE_3_SRPC_SERVER_PORT")?.parse()?,
+        &std::env::var("EXAMPLE_3_SRPC_SERVER_ENPOINT")?,
+        &std::env::var("EXAMPLE_3_SRPC_SERVER_CERT")?,
+        &std::env::var("EXAMPLE_3_SRPC_SERVER_KEY")?,
+    );
+
+    // Connect to the server
+    let client = config.connect().await?;
+    info!("Connected to server");
+
+    let message = "Hypervisor.ListVMs\n";
+
+    // Send a message
+    info!("Sending message: {:?}", message);
+    client.send_message(message).await?;
+    info!("Sent message: {:?}", message);
+
+    // Receive an empty response
+    info!("Waiting for empty string response...");
+    let mut rx = client
+        .receive_message(true, |_| false, &ReceiveOptions::default())
+        .await?;
+    while let Some(result) = rx.recv().await {
+        match result {
+            Ok(response) => info!("Received response: {:?}", response),
+            Err(e) => error!("Error receiving message: {:?}", e),
+        }
+    }
+
+    #[derive(Debug, serde::Serialize)]
+    struct ListVMsRequest {
+        ignore_state_mask: u32,
+        owner_groups: Vec<String>,
+        owner_users: Vec<String>,
+        sort: bool,
+        vm_tags_to_match: HashMap<String, String>,
+    }
+
+    #[derive(Debug, serde::Deserialize)]
+    struct ListVMsResponse {
+        ip_addresses: Vec<String>,
+    }
+
+    let request = ListVMsRequest {
+        ignore_state_mask: 0,
+        owner_groups: vec![],
+        owner_users: vec![],
+        sort: false,
+        vm_tags_to_match: HashMap::new(),
+    };
+
+    // Send a JSON message
+    info!("Sending payload: {:?}", request);
+    client.send_json(&serde_json::to_value(&request)?).await?;
+    info!("Sent payload: {:?}", request);
+
+    // Receive an empty response
+    info!("Waiting for empty string response for payload...");
+    let mut rx = client
+        .receive_message(true, |_| false, &ReceiveOptions::default())
+        .await?;
+    while let Some(result) = rx.recv().await {
+        match result {
+            Ok(response) => info!("Received response: {:?}", response),
+            Err(e) => error!("Error receiving message: {:?}", e),
+        }
+    }
+
+    // Receive responses
+    let mut rx = client
+        .receive_json(|_| false, &ReceiveOptions::default())
+        .await?;
+    while let Some(result) = rx.recv().await {
+        match result
+            .and_then(|response| {
+                serde_json::from_value::<ListVMsResponse>(response)
+                    .map_err(|e| Box::new(CustomError(e.to_string())) as Box<dyn Error + Send>)
+            })
+            .map_err(|e| Box::new(CustomError(e.to_string())) as Box<dyn Error>)
+        {
+            Ok(response) => info!("Received response: {:?}", response),
+            Err(e) => error!("Error receiving message: {:?}", e),
+        }
+    }
+
+    info!("Sending request_reply: {}", message);
+    let res = client
+        .request_reply::<SimpleValue>(message, serde_json::to_value(&request)?)
+        .await
+        .map_err(|e| Box::new(CustomError(e.to_string())) as Box<dyn Error>)?;
+    info!(
+        "Sent request_reply: {}, got reply: {:?}",
+        message,
+        serde_json::to_string(&res)?
+    );
+
+    Ok(())
+}

rust/lib/srpc/client/src/lib.rs

@@ -1,3 +1,4 @@
+use async_trait::async_trait;
 use chunk_limiter::ChunkLimiter;
 use futures::StreamExt;
 use openssl::ssl::{Ssl, SslConnector, SslMethod, SslVerifyMode};
@@ -20,7 +21,7 @@ mod tests;
 
 // Custom error type
 #[derive(Debug)]
-struct CustomError(String);
+pub struct CustomError(pub String);
 
 impl fmt::Display for CustomError {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -82,6 +83,70 @@ where
     stream: Arc<Mutex<T>>,
 }
 
+#[async_trait]
+pub trait RequestReply {
+    type Request;
+    type Reply;
+
+    async fn request_reply<T>(
+        client: &ConnectedClient<T>,
+        payload: Self::Request,
+    ) -> Result<Self::Reply, Box<dyn Error + Send>>
+    where
+        T: AsyncRead + AsyncWrite + Unpin + Send + 'static;
+}
+
+pub struct SimpleValue;
+
+#[async_trait]
+impl RequestReply for SimpleValue {
+    type Request = serde_json::Value;
+    type Reply = serde_json::Value;
+
+    async fn request_reply<T>(
+        client: &ConnectedClient<T>,
+        payload: Self::Request,
+    ) -> Result<Self::Reply, Box<dyn Error + Send>>
+    where
+        T: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+    {
+        client.send_json_and_check(&payload).await?;
+
+        let mut rx = client
+            .receive_json(|_| false, &ReceiveOptions::default())
+            .await
+            .map_err(|e| Box::new(CustomError(e.to_string())) as Box<dyn Error + Send>)?;
+        let json_value = rx.recv().await.ok_or_else(|| {
+            Box::new(CustomError("Expected JSON value".to_string())) as Box<dyn Error + Send>
+        })??;
+        Ok(json_value)
+    }
+}
+
+pub struct StreamValue;
+
+#[async_trait]
+impl RequestReply for StreamValue {
+    type Request = serde_json::Value;
+    type Reply = mpsc::Receiver<Result<serde_json::Value, Box<dyn Error + Send>>>;
+
+    async fn request_reply<T>(
+        client: &ConnectedClient<T>,
+        payload: Self::Request,
+    ) -> Result<Self::Reply, Box<dyn Error + Send>>
+    where
+        T: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+    {
+        client.send_json_and_check(&payload).await?;
+
+        let rx = client
+            .receive_json(|_| true, &ReceiveOptions::default())
+            .await
+            .map_err(|e| Box::new(CustomError(e.to_string())) as Box<dyn Error + Send>)?;
+        Ok(rx)
+    }
+}
+
 impl ClientConfig {
     pub fn new(host: &str, port: u16, path: &str, cert: &str, key: &str) -> Self {
         ClientConfig {
@@ -190,6 +255,36 @@ where
         }
     }
 
+    pub async fn request_reply<R>(
+        &self,
+        method: &str,
+        payload: R::Request,
+    ) -> Result<R::Reply, Box<dyn Error + Send>>
+    where
+        R: RequestReply,
+    {
+        self.send_message(method)
+            .await
+            .map_err(|e| Box::new(CustomError(e.to_string())) as Box<dyn Error + Send>)?;
+        let mut rx = self
+            .receive_message(true, |_| false, &ReceiveOptions::default())
+            .await
+            .map_err(|e| Box::new(CustomError(e.to_string())) as Box<dyn Error + Send>)?;
+        if rx
+            .recv()
+            .await
+            .ok_or_else(|| {
+                Box::new(CustomError("Expected response".to_string())) as Box<dyn Error + Send>
+            })??
+            .is_empty()
+        {
+        } else {
+            return Err(Box::new(CustomError("Expected empty line".to_string())));
+        }
+
+        R::request_reply(self, payload).await
+    }
+
     pub async fn send_message(&self, message: &str) -> Result<(), Box<dyn Error>> {
         let stream = self.stream.lock().await;
         let mut pinned = Pin::new(stream);

rust/lib/srpc/client/src/python_bindings.rs

@@ -1,4 +1,4 @@
-use crate::{ClientConfig, ConnectedClient, ReceiveOptions};
+use crate::{ClientConfig, ConnectedClient, ReceiveOptions, SimpleValue};
 use futures::{Stream, StreamExt};
 use pyo3::exceptions::{PyRuntimeError, PyStopAsyncIteration};
 use pyo3::prelude::*;
@@ -277,4 +277,25 @@ impl ConnectedSrpcClient {
             }))
         })
     }
+
+    pub fn request_reply<'p>(
+        &self,
+        py: Python<'p>,
+        method: String,
+        payload: String,
+    ) -> PyResult<Bound<'p, PyAny>> {
+        let client = self.0.clone();
+
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            let value: Value = serde_json::from_str(&payload)
+                .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+            let response = client
+                .lock()
+                .await
+                .request_reply::<SimpleValue>(&method, value)
+                .await
+                .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+            Ok(response.to_string())
+        })
+    }
 }

rust/lib/srpc/client/srpc_client.pyi

@@ -1,5 +1,6 @@
 from typing import Callable, List
 
+type JsonStr = str
 
 class SrpcClientConfig:
     def __init__(self, host: str, port: int, path: str, cert: str, key: str) -> None: ...
@@ -12,3 +13,4 @@ class ConnectedSrpcClient:
     async def send_json(self, payload: str) -> None: ...
     async def receive_json(self, should_continue: bool) -> List[str]: ...
     async def receive_json_cb(self, should_continue: Callable[[str], bool]) -> List[str]: ...
+    async def request_reply(self, message: str, payload: JsonStr) -> JsonStr: ...