improve reconnect logic

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