improve reconnect logic

examples/simple_command/Cargo.toml

@@ -6,5 +6,7 @@ edition = "2021"
 [dependencies]
 anyhow = { workspace = true }
 api = { path = "../../api" }
+env_logger = { workspace = true }
+log = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread", "signal"] }
 tokio-util = { workspace = true }

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(())
 }

examples/simple_producer/Cargo.toml

@@ -7,8 +7,10 @@ edition = "2021"
 anyhow = { workspace = true }
 api = { path = "../../api" }
 chrono = { workspace = true }
+env_logger = { workspace = true }
 futures-util = { workspace = true }
 num-traits = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread", "signal", "time", "macros"] }
 tokio-util = { workspace = true }
 uuid = { workspace = true }
+log = "0.4.29"

examples/simple_producer/src/main.rs

@@ -9,6 +9,7 @@ use futures_util::future::join_all;
 use num_traits::FloatConst;
 use std::sync::Arc;
 use std::time::Duration;
+use tokio::sync::{oneshot, RwLock};
 use tokio::time::{sleep_until, Instant};
 use tokio_util::sync::CancellationToken;
 use uuid::Uuid;
@@ -18,7 +19,8 @@ struct Telemetry {
 }
 
 struct TelemetryItemHandle {
-    uuid: Uuid,
+    cancellation_token: CancellationToken,
+    uuid: Arc<RwLock<Uuid>>,
     client: Arc<Client>,
 }
 
@@ -32,13 +34,48 @@ impl Telemetry {
         name: String,
         data_type: DataType,
     ) -> anyhow::Result<TelemetryItemHandle> {
-        let response = self
-            .client
-            .send_request(TelemetryDefinitionRequest { name, data_type })
-            .await?;
+        let cancellation_token = CancellationToken::new();
+        let cancel_token = cancellation_token.clone();
+        let client = self.client.clone();
+
+        let response_uuid = Arc::new(RwLock::new(Uuid::nil()));
+
+        let response_uuid_inner = response_uuid.clone();
+        let (tx, rx) = oneshot::channel();
+
+        tokio::spawn(async move {
+            let mut write_lock = Some(response_uuid_inner.write().await);
+            let _ = tx.send(());
+            while !cancel_token.is_cancelled() {
+                if let Ok(response) = client
+                    .send_request(TelemetryDefinitionRequest {
+                        name: name.clone(),
+                        data_type,
+                    })
+                    .await
+                {
+                    let mut lock = match write_lock {
+                        None => response_uuid_inner.write().await,
+                        Some(lock) => lock,
+                    };
+                    // Update the value in the lock
+                    *lock = response.uuid;
+                    // Set this value so the loop works
+                    write_lock = None;
+                }
+
+                client.wait_disconnected().await;
+            }
+        });
+
+        // Wait until the write lock is acquired
+        let _ = rx.await;
+        // Wait until the write lock is released for the first time
+        drop(response_uuid.read().await);
 
         Ok(TelemetryItemHandle {
-            uuid: response.uuid,
+            cancellation_token,
+            uuid: response_uuid,
             client: self.client.clone(),
         })
     }
@@ -46,9 +83,15 @@ impl Telemetry {
 
 impl TelemetryItemHandle {
     pub async fn publish(&self, value: DataValue, timestamp: DateTime<Utc>) -> anyhow::Result<()> {
+        let Ok(lock) = self.uuid.try_read() else {
+            return Ok(());
+        };
+        let uuid = *lock;
+        drop(lock);
+
         self.client
             .send_message_if_connected(TelemetryEntry {
-                uuid: self.uuid,
+                uuid,
                 value,
                 timestamp,
             })
@@ -62,8 +105,16 @@ impl TelemetryItemHandle {
     }
 }
 
+impl Drop for TelemetryItemHandle {
+    fn drop(&mut self) {
+        self.cancellation_token.cancel();
+    }
+}
+
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
+    env_logger::init();
+
     let client = Arc::new(Client::connect("ws://[::1]:8080/backend")?);
     let tlm = Telemetry::new(client);
 
@@ -141,7 +192,7 @@ async fn main() -> anyhow::Result<()> {
         let mut index = 0;
         let mut tasks = vec![];
         while !cancellation_token.is_cancelled() {
-            next_time += Duration::from_millis(10);
+            next_time += Duration::from_millis(1000);
             index += 1;
             sleep_until(next_time).await;
             let publish_time = start_time + TimeDelta::from_std(next_time - start_instant).unwrap();