improve reconnect logic

examples/simple_command/src/main.rs

@@ -4,6 +4,7 @@ use api::data_type::DataType;
 use api::messages::command::{
     Command, CommandDefinition, CommandParameterDefinition, CommandResponse,
 };
+use log::info;
 use std::error::Error;
 use std::sync::Arc;
 use tokio_util::sync::CancellationToken;
@@ -26,15 +27,70 @@ fn handle_command(command: Command) -> anyhow::Result<String> {
         .ok_or(anyhow!("Parameter 'c' Missing"))?)
     .try_into()?;
 
-    println!("Command Received:\n  timestamp: {timestamp}\n  a: {a}\n  b: {b}\n  c: {c}");
+    info!("Command Received:\n  timestamp: {timestamp}\n  a: {a}\n  b: {b}\n  c: {c}");
 
     Ok(format!(
         "Successfully Received Command! timestamp: {timestamp} a: {a} b: {b} c: {c}"
     ))
 }
 
+struct CommandHandle {
+    cancellation_token: CancellationToken,
+}
+
+impl CommandHandle {
+    pub async fn register(
+        client: Arc<Client>,
+        command_definition: CommandDefinition,
+        mut callback: impl FnMut(Command) -> anyhow::Result<String> + Send + 'static,
+    ) -> anyhow::Result<Self> {
+        let cancellation_token = CancellationToken::new();
+        let result = Self {
+            cancellation_token: cancellation_token.clone(),
+        };
+
+        tokio::spawn(async move {
+            while !cancellation_token.is_cancelled() {
+                // This would only fail if the sender closed while trying to insert data
+                // It would wait until space is made
+                let Ok(mut rx) = client
+                    .register_callback_channel(command_definition.clone())
+                    .await
+                else {
+                    continue;
+                };
+
+                while let Some((cmd, responder)) = rx.recv().await {
+                    let response = match callback(cmd) {
+                        Ok(response) => CommandResponse {
+                            success: true,
+                            response,
+                        },
+                        Err(err) => CommandResponse {
+                            success: false,
+                            response: err.to_string(),
+                        },
+                    };
+                    // This should only err if we had an error elsewhere
+                    let _ = responder.send(response);
+                }
+            }
+        });
+
+        Ok(result)
+    }
+}
+
+impl Drop for CommandHandle {
+    fn drop(&mut self) {
+        self.cancellation_token.cancel();
+    }
+}
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn Error>> {
+    env_logger::init();
+
     let cancellation_token = CancellationToken::new();
     {
         let cancellation_token = cancellation_token.clone();
@@ -46,39 +102,32 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
     let client = Arc::new(Client::connect("ws://[::1]:8080/backend")?);
 
-    client
-        .register_callback_fn(
-            CommandDefinition {
-                name: "simple_command/a".to_string(),
-                parameters: vec![
-                    CommandParameterDefinition {
-                        name: "a".to_string(),
-                        data_type: DataType::Float32,
-                    },
-                    CommandParameterDefinition {
-                        name: "b".to_string(),
-                        data_type: DataType::Float64,
-                    },
-                    CommandParameterDefinition {
-                        name: "c".to_string(),
-                        data_type: DataType::Boolean,
-                    },
-                ],
-            },
-            |command| match handle_command(command) {
-                Ok(response) => CommandResponse {
-                    success: true,
-                    response,
+    let handle = CommandHandle::register(
+        client,
+        CommandDefinition {
+            name: "simple_command/a".to_string(),
+            parameters: vec![
+                CommandParameterDefinition {
+                    name: "a".to_string(),
+                    data_type: DataType::Float32,
                 },
-                Err(error) => CommandResponse {
-                    success: false,
-                    response: error.to_string(),
+                CommandParameterDefinition {
+                    name: "b".to_string(),
+                    data_type: DataType::Float64,
                 },
-            },
-        )
-        .await?;
+                CommandParameterDefinition {
+                    name: "c".to_string(),
+                    data_type: DataType::Boolean,
+                },
+            ],
+        },
+        handle_command,
+    )
+    .await?;
 
     cancellation_token.cancelled().await;
 
+    drop(handle);
+
     Ok(())
 }