refactor(connection): make the resend of the retention solid

Pull the packets while sending them, this is a limitation of the MQTT
library.

Signed-off-by: Joshua Chapman <joshua.chapman@secomind.com>

Author: joshuachp

astarte-device-sdk-derive/src/lib.rs

@@ -301,7 +301,7 @@ pub fn astarte_aggregate_derive(input: TokenStream) -> TokenStream {
 ///
 /// ### Example
 ///
-/// To derive the trait it for an individual.
+/// To derive the trait for an individual.
 ///
 /// ```no_compile
 /// #[derive(FromEvent)]

examples/object_datastream/main.rs

@@ -18,6 +18,7 @@
 
 use std::time::Duration;
 
+use chrono::Utc;
 use serde::{Deserialize, Serialize};
 
 use astarte_device_sdk::IntoAstarteObject;
@@ -84,10 +85,11 @@ async fn main() -> eyre::Result<()> {
 
             println!("Sending {data:?}");
             client
-                .send_object(
+                .send_object_with_timestamp(
                     "org.astarte-platform.rust.examples.object-datastream.DeviceDatastream",
                     "/23",
                     data.try_into().unwrap(),
+                    Utc::now(),
                 )
                 .await?;
 

src/client/individual.rs

@@ -23,7 +23,7 @@ use tracing::{debug, trace, warn};
 use crate::client::ValidatedIndividual;
 use crate::interface::mapping::path::MappingPath;
 use crate::interface::Retention;
-use crate::state::SharedState;
+use crate::state::{SharedState, Status};
 use crate::store::StoreCapabilities;
 use crate::transport::Connection;
 use crate::{AstarteType, Error};
@@ -51,16 +51,24 @@ where
 
         debug!("sending individual type {}", validated.data.display_type());
 
-        if !self.state.status.is_connected() {
-            trace!("publish individual while connection is offline");
-
-            return Self::offline_send_individual(
-                &self.state,
-                &self.store,
-                &mut self.sender,
-                validated,
-            )
-            .await;
+        match self.state.status.connection() {
+            Status::Connected => {
+                trace!("publish individual while connection is online");
+            }
+            Status::Disconnected => {
+                trace!("publish individual while connection is offline");
+
+                return Self::offline_send_individual(
+                    &self.state,
+                    &self.store,
+                    &mut self.sender,
+                    validated,
+                )
+                .await;
+            }
+            Status::Closed => {
+                return Err(Error::Disconnected);
+            }
         }
 
         match mapping.retention() {

src/client/object.rs

@@ -25,7 +25,7 @@ use crate::client::ValidatedObject;
 use crate::error::AggregationError;
 use crate::interface::mapping::path::MappingPath;
 use crate::interface::{Aggregation, Retention};
-use crate::state::SharedState;
+use crate::state::{SharedState, Status};
 use crate::store::StoreCapabilities;
 use crate::transport::Connection;
 use crate::Error;
@@ -66,16 +66,24 @@ where
 
         debug!("sending object {}{}", interface_name, path);
 
-        if !self.state.status.is_connected() {
-            trace!("publish object while connection is offline");
-
-            return Self::offline_send_object(
-                &self.state,
-                &self.store,
-                &mut self.sender,
-                validated,
-            )
-            .await;
+        match self.state.status.connection() {
+            Status::Connected => {
+                trace!("publish object while connection is connected");
+            }
+            Status::Disconnected => {
+                trace!("publish object while connection is offline");
+
+                return Self::offline_send_object(
+                    &self.state,
+                    &self.store,
+                    &mut self.sender,
+                    validated,
+                )
+                .await;
+            }
+            Status::Closed => {
+                return Err(Error::Disconnected);
+            }
         }
 
         match validated.retention {

src/client/property.rs

@@ -23,6 +23,7 @@ use tracing::{debug, error, trace};
 use crate::interface::mapping::path::MappingPath;
 use crate::interface::reference::PropertyRef;
 use crate::interface::Ownership;
+use crate::state::Status;
 use crate::store::{PropertyMapping, PropertyStore, StoredProp};
 use crate::transport::Connection;
 use crate::validate::{ValidatedProperty, ValidatedUnset};
@@ -70,13 +71,21 @@ where
             mapping.interface().version_major()
         );
 
-        if self.state.status.is_connected() {
-            self.sender.send_property(validated).await?;
+        match self.state.status.connection() {
+            Status::Connected => {
+                self.sender.send_property(validated).await?;
 
-            trace!(
-                "property sent {interface_name}{path}:{}",
-                mapping.interface().version_major()
-            );
+                trace!(
+                    "property sent {interface_name}{path}:{}",
+                    mapping.interface().version_major()
+                );
+            }
+            Status::Disconnected => {
+                trace!("property not sent since offline")
+            }
+            Status::Closed => {
+                return Err(Error::Disconnected);
+            }
         }
 
         Ok(())
@@ -147,14 +156,20 @@ where
         let property_mapping = (&validated).into();
         self.store.unset_prop(&property_mapping).await?;
 
-        if self.state.status.is_connected() {
-            self.sender.unset(validated.clone()).await?;
+        match self.state.status.connection() {
+            Status::Connected => {
+                self.sender.unset(validated.clone()).await?;
 
-            debug!("deleting property {interface_name}{path} from store",);
+                debug!("deleting property {interface_name}{path} from store");
 
-            // TODO: this should be done when the package has been acknowledged, but it's hard
-            //       for the MQTT implementation at the moment so we delete it here to cleanup
-            self.store.delete_prop(&property_mapping).await?;
+                self.store.delete_prop(&property_mapping).await?;
+            }
+            Status::Disconnected => {
+                trace!("not deleting property from store, since disconnected");
+            }
+            Status::Closed => {
+                return Err(Error::Disconnected);
+            }
         }
 
         Ok(())

src/connection/incoming.rs

@@ -23,8 +23,8 @@ use crate::error::{AggregationError, InterfaceTypeError};
 use crate::interface::mapping::path::MappingPath;
 use crate::interface::{Aggregation, InterfaceTypeDef};
 use crate::store::{PropertyMapping, PropertyStore, StoredProp};
-use crate::transport::{Connection, Receive, ReceivedEvent, TransportError};
-use crate::{DeviceEvent, Error, Interface, Value};
+use crate::transport::{Connection, Receive, TransportError};
+use crate::{Error, Interface, Value};
 
 use super::DeviceConnection;
 
@@ -33,7 +33,7 @@ where
     C: Connection,
 {
     #[instrument(skip(self, payload))]
-    async fn handle_event(
+    pub(crate) async fn handle_event(
         &self,
         interface: &str,
         path: &str,
@@ -205,32 +205,4 @@ where
 
         Ok(Value::Object { data, timestamp })
     }
-
-    pub(super) async fn handle_connection_event(
-        &self,
-        event: ReceivedEvent<C::Payload>,
-    ) -> Result<(), Error>
-    where
-        C: Receive + Sync,
-    {
-        let data = match self
-            .handle_event(&event.interface, &event.path, event.payload)
-            .await
-        {
-            Ok(aggregation) => Ok(DeviceEvent {
-                interface: event.interface,
-                path: event.path,
-                data: aggregation,
-            }),
-            Err(TransportError::Recv(recv_err)) => Err(recv_err),
-            Err(TransportError::Transport(err)) => {
-                return Err(err);
-            }
-        };
-
-        self.tx
-            .send_async(data)
-            .await
-            .map_err(|_| Error::Disconnected)
-    }
 }

src/connection/mod.rs

@@ -22,16 +22,18 @@ use std::future::Future;
 use std::sync::Arc;
 
 use chrono::Utc;
-use tracing::{debug, info, warn};
+use tokio::task::JoinHandle;
+use tracing::{debug, info, trace, warn};
 
-use crate::state::SharedState;
+use crate::error::Report;
+use crate::state::{SharedState, Status};
 use crate::transport::TransportError;
 use crate::Timestamp;
 use crate::{
     client::RecvError,
     event::DeviceEvent,
     store::wrapper::StoreWrapper,
-    transport::{Connection, Disconnect, Publish, Receive, Reconnect, Register},
+    transport::{Connection, Publish, Receive, Reconnect},
     Error,
 };
 
@@ -84,6 +86,7 @@ where
     connection: C,
     sender: C::Sender,
     state: Arc<SharedState>,
+    resend: Option<JoinHandle<()>>,
 }
 
 impl<C> DeviceConnection<C>
@@ -103,6 +106,7 @@ where
             state,
             connection,
             sender,
+            resend: None,
         }
     }
 
@@ -130,19 +134,61 @@ where
             }),
         }
     }
+
+    /// Keeps polling connection events
+    pub(super) async fn poll(&mut self) -> Result<Status, TransportError>
+    where
+        C: Receive + Reconnect,
+        C::Sender: Publish + 'static,
+    {
+        let Some(event) = self.connection.next_event().await? else {
+            trace!("disconnected");
+
+            self.state.status.set_connected(false);
+
+            // This will check if the connection was closed
+            return Ok(self.state.status.connection());
+        };
+
+        let event = self
+            .handle_event(&event.interface, &event.path, event.payload)
+            .await
+            .map(|data| DeviceEvent {
+                interface: event.interface,
+                path: event.path,
+                data,
+            })?;
+
+        self.tx.send(Ok(event)).map_err(|err| {
+            debug!(error = %Report::new(err), "disconnected");
+
+            TransportError::Transport(Error::Disconnected)
+        })?;
+
+        Ok(self.state.status.connection())
+    }
 }
 
 impl<C> EventLoop for DeviceConnection<C>
 where
-    C: Connection + Reconnect + Receive + Send + Sync + 'static,
-    C::Sender: Send + Register + Publish + Disconnect + 'static,
+    C: Connection + Reconnect + Receive + 'static,
+    C::Sender: Publish + 'static,
 {
     async fn handle_events(mut self) -> Result<(), crate::Error> {
         self.init_stored_retention().await?;
 
+        // We are connected and all the stored packet have been sent
+        self.state.status.set_connected(true);
+
         loop {
-            let opt = match self.connection.next_event().await {
-                Ok(opt) => opt,
+            match self.poll().await {
+                Ok(Status::Connected) => {}
+                Ok(Status::Disconnected) => {
+                    self.reconnect_and_resend().await?;
+                }
+                Ok(Status::Closed) => {
+                    break;
+                }
                 Err(TransportError::Transport(err)) => {
                     return Err(err);
                 }
@@ -152,45 +198,8 @@ where
                         .send_async(Err(recv_err))
                         .await
                         .map_err(|_| Error::Disconnected)?;
-
-                    continue;
                 }
-            };
-
-            let Some(event_data) = opt else {
-                if self.state.status.is_closed() {
-                    debug!("connection closed");
-
-                    break;
-                }
-
-                debug!("reconnecting");
-
-                self.state.status.set_connected(false);
-
-                let interfaces = self.state.interfaces.read().await;
-
-                self.connection.reconnect(&interfaces).await?;
-
-                Self::resend_volatile_publishes(
-                    &self.state.volatile_store,
-                    &mut self.sender,
-                    &self.state.status,
-                )
-                .await?;
-                Self::resend_stored_publishes(
-                    &mut self.store,
-                    &mut self.sender,
-                    &self.state.status,
-                )
-                .await?;
-
-                self.state.status.set_connected(true);
-
-                continue;
-            };
-
-            self.handle_connection_event(event_data).await?
+            }
         }
 
         info!("connection closed successfully");