improve telemetry ergonomics

api/src/client/telemetry.rs

@@ -1,77 +1,104 @@
 use crate::client::error::MessageError;
 use crate::client::Client;
-use crate::data_type::DataType;
 use crate::data_value::DataValue;
 use crate::messages::telemetry_definition::TelemetryDefinitionRequest;
 use crate::messages::telemetry_entry::TelemetryEntry;
+use api_core::data_type::{DataType, ToDataType};
 use chrono::{DateTime, Utc};
+use std::marker::PhantomData;
 use std::sync::Arc;
 use tokio::sync::{oneshot, RwLock};
 use tokio_util::sync::CancellationToken;
 use uuid::Uuid;
 
-pub struct Telemetry;
+pub struct TelemetryRegistry {
+    client: Arc<Client>,
+}
 
-impl Telemetry {
-    pub async fn register(
-        client: Arc<Client>,
-        name: String,
+impl TelemetryRegistry {
+    pub fn new(client: Arc<Client>) -> Self {
+        Self { client }
+    }
+
+    #[inline]
+    pub async fn register_generic(
+        &self,
+        name: impl Into<String>,
         data_type: DataType,
-    ) -> TelemetryHandle {
-        let cancellation_token = CancellationToken::new();
-        let cancel_token = cancellation_token.clone();
-        let stored_client = client.clone();
+    ) -> GenericTelemetryHandle {
+        // inner for compilation performance
+        async fn inner(
+            client: Arc<Client>,
+            name: String,
+            data_type: DataType,
+        ) -> GenericTelemetryHandle {
+            let cancellation_token = CancellationToken::new();
+            let cancel_token = cancellation_token.clone();
+            let stored_client = client.clone();
 
-        let response_uuid = Arc::new(RwLock::new(Uuid::nil()));
+            let response_uuid = Arc::new(RwLock::new(Uuid::nil()));
 
-        let response_uuid_inner = response_uuid.clone();
-        let (tx, rx) = oneshot::channel();
+            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;
+            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;
                 }
+            });
 
-                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);
+
+            GenericTelemetryHandle {
+                cancellation_token,
+                uuid: response_uuid,
+                client: stored_client,
             }
-        });
-
-        // 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);
-
-        TelemetryHandle {
-            cancellation_token,
-            uuid: response_uuid,
-            client: stored_client,
         }
+        inner(self.client.clone(), name.into(), data_type).await
+    }
+
+    #[inline]
+    pub async fn register<T: ToDataType>(&self, name: impl Into<String>) -> TelemetryHandle<T> {
+        self.register_generic(name, T::DATA_TYPE).await.coerce()
     }
 }
 
-pub struct TelemetryHandle {
+impl Drop for GenericTelemetryHandle {
+    fn drop(&mut self) {
+        self.cancellation_token.cancel();
+    }
+}
+
+pub struct GenericTelemetryHandle {
     cancellation_token: CancellationToken,
     uuid: Arc<RwLock<Uuid>>,
     client: Arc<Client>,
 }
 
-impl TelemetryHandle {
+impl GenericTelemetryHandle {
     pub async fn publish(
         &self,
         value: DataValue,
@@ -97,10 +124,42 @@ impl TelemetryHandle {
 
         Ok(())
     }
-}
 
-impl Drop for TelemetryHandle {
-    fn drop(&mut self) {
-        self.cancellation_token.cancel();
+    #[inline]
+    pub async fn publish_now(&self, value: DataValue) -> Result<(), MessageError> {
+        self.publish(value, Utc::now()).await
+    }
+
+    fn coerce<T: Into<DataValue>>(self) -> TelemetryHandle<T> {
+        TelemetryHandle::<T> {
+            generic_handle: self,
+            _phantom: PhantomData,
+        }
+    }
+}
+
+pub struct TelemetryHandle<T> {
+    generic_handle: GenericTelemetryHandle,
+    _phantom: PhantomData<T>,
+}
+
+impl<T> TelemetryHandle<T> {
+    pub fn to_generic(self) -> GenericTelemetryHandle {
+        self.generic_handle
+    }
+    pub fn as_generic(&self) -> &GenericTelemetryHandle {
+        &self.generic_handle
+    }
+}
+
+impl<T: Into<DataValue>> TelemetryHandle<T> {
+    #[inline]
+    pub async fn publish(&self, value: T, timestamp: DateTime<Utc>) -> Result<(), MessageError> {
+        self.as_generic().publish(value.into(), timestamp).await
+    }
+
+    #[inline]
+    pub async fn publish_now(&self, value: T) -> Result<(), MessageError> {
+        self.publish(value, Utc::now()).await
     }
 }