fix: flush pipeline aggregator cache on shutdown with optional actor shutdown trait

Author: m0ar

actor/src/lib.rs

@@ -61,6 +61,16 @@ pub trait Actor {
     }
 }
 
+/// Optional trait for actors that need to perform cleanup operations during shutdown.
+/// Only actors that need cleanup should implement this trait.
+#[async_trait]
+pub trait Shutdown {
+    /// Called when the actor is shutting down to allow for cleanup operations.
+    /// This method is called after all pending messages have been processed
+    /// but before the actor loop exits.
+    async fn shutdown(&mut self);
+}
+
 /// Wrapper of any T with its tracing span context.
 #[derive(Debug)]
 pub struct Traced<T> {

actor/src/macros.rs

@@ -1,19 +1,75 @@
 /// Constructs an envelope enumeration that contains all messages for an actor.
 ///
+/// # Usage
+///
+/// Basic usage (no shutdown support):
+/// ```ignore
+/// actor_envelope! {
+///     MyEnvelope,
+///     MyActor,
+///     MyRecorder,
+///     Message1 => Msg1Type,
+///     Message2 => Msg2Type,
+/// }
+/// ```
+///
+/// With shutdown support (actor must implement `Shutdown` trait):
+/// ```ignore
+/// actor_envelope! {
+///     MyEnvelope,
+///     MyActor,
+///     MyRecorder,
+///     with_shutdown,
+///     Message1 => Msg1Type,
+///     Message2 => Msg2Type,
+/// }
+/// ```
+///
 /// The first identifier is the name of the enum.
 /// The second identifier is the name of a trait specific to the actor.
 /// The third identifier is the name of a trait for recording message events.
+/// Optional `with_shutdown` parameter adds shutdown support for actors that implement the Shutdown trait.
 /// The remaining pairs are the variants of the envelope indicating the messages the actor handles.
 ///
 /// The constructed actor trait is a union of the [`crate::Handler`] traits for each message along with the [`crate::Actor`] trait.
 ///
 /// The constructed recorder trait is a union of the [`ceramic_metrics::Recorder`] traits for each message.
 #[macro_export]
 macro_rules! actor_envelope {
+    // Version with shutdown support
+    (
+        $enum_name:ident,
+        $actor_trait:ident,
+        $recorder_trait:ident,
+        with_shutdown,
+        $(
+            $variant_name:ident => $message_type:ty,
+        )*
+    ) => {
+        $crate::actor_envelope_impl!($enum_name, $actor_trait, $recorder_trait, with_shutdown, $($variant_name => $message_type,)*);
+    };
+
+    // Default version without shutdown
+    (
+        $enum_name:ident,
+        $actor_trait:ident,
+        $recorder_trait:ident,
+        $(
+            $variant_name:ident => $message_type:ty,
+        )*
+    ) => {
+        $crate::actor_envelope_impl!($enum_name, $actor_trait, $recorder_trait, without_shutdown, $($variant_name => $message_type,)*);
+    };
+}
+
+/// Internal implementation macro that handles both with/without shutdown cases
+#[macro_export]
+macro_rules! actor_envelope_impl {
     (
         $enum_name:ident,
         $actor_trait:ident,
         $recorder_trait:ident,
+        $shutdown_mode:ident,
         $(
             $variant_name:ident => $message_type:ty,
         )*
@@ -51,11 +107,7 @@ macro_rules! actor_envelope {
                 }
             }
         }
-        #[doc = std::stringify!($actor_trait)]
-        #[doc = " is an [`ceramic_actor::Actor`] and [`ceramic_actor::Handler`] for each message type in the actor envelope "]
-        #[doc = stringify!($enum_name)]
-        #[doc = "."]
-        pub trait $actor_trait : $crate::Actor<Envelope = $enum_name> $( + $crate::Handler<$message_type> )* + ::std::marker::Send + 'static { }
+        $crate::actor_trait_def!($actor_trait, $enum_name, $shutdown_mode, $($message_type,)*);
 
         #[doc = std::stringify!($recorder_trait)]
         #[doc = " is an [`ceramic_metrics::Recorder`] for each message type in the actor envelope "]
@@ -64,6 +116,7 @@ macro_rules! actor_envelope {
         pub trait $recorder_trait : $(::ceramic_metrics::Recorder<$crate::MessageEvent<$message_type>> +)*
             ::std::fmt::Debug + ::std::marker::Send + ::std::marker::Sync + 'static { }
 
+
         impl $enum_name {
             /// Runs the actor handling messages as they arrive.
             pub async fn run<A>(mut actor: A, mut receiver: $crate::Receiver<A::Envelope>, mut shutdown: impl ::std::future::Future<Output=()> + ::std::marker::Send + 'static)
@@ -111,6 +164,7 @@ macro_rules! actor_envelope {
                         }
                     }
                 }
+                $crate::actor_shutdown_call!($shutdown_mode, actor);
             }
         }
         $(
@@ -132,3 +186,33 @@ macro_rules! actor_envelope {
         )*
     };
 }
+
+/// Helper macro to define actor trait with or without shutdown requirement
+#[macro_export]
+macro_rules! actor_trait_def {
+    ($actor_trait:ident, $enum_name:ident, with_shutdown, $($message_type:ty,)*) => {
+        #[doc = std::stringify!($actor_trait)]
+        #[doc = " is an [`ceramic_actor::Actor`] and [`ceramic_actor::Handler`] for each message type in the actor envelope "]
+        #[doc = stringify!($enum_name)]
+        #[doc = ". This actor must also implement [`ceramic_actor::Shutdown`]."]
+        pub trait $actor_trait : $crate::Actor<Envelope = $enum_name> $( + $crate::Handler<$message_type> )* + $crate::Shutdown + ::std::marker::Send + 'static { }
+    };
+    ($actor_trait:ident, $enum_name:ident, without_shutdown, $($message_type:ty,)*) => {
+        #[doc = std::stringify!($actor_trait)]
+        #[doc = " is an [`ceramic_actor::Actor`] and [`ceramic_actor::Handler`] for each message type in the actor envelope "]
+        #[doc = stringify!($enum_name)]
+        #[doc = "."]
+        pub trait $actor_trait : $crate::Actor<Envelope = $enum_name> $( + $crate::Handler<$message_type> )* + ::std::marker::Send + 'static { }
+    };
+}
+
+/// Helper macro to conditionally call shutdown
+#[macro_export]
+macro_rules! actor_shutdown_call {
+    (with_shutdown, $actor:ident) => {
+        $actor.shutdown().await;
+    };
+    (without_shutdown, $actor:ident) => {
+        // No shutdown call for actors that don't implement shutdown
+    };
+}

pipeline/src/aggregator/mock.rs

@@ -1,6 +1,6 @@
 //! Provides a mock implmentation of the aggregator actor.
 use async_trait::async_trait;
-use ceramic_actor::{Actor, Handler, Message};
+use ceramic_actor::{Actor, Handler, Message, Shutdown};
 use mockall::mock;
 use prometheus_client::registry::Registry;
 
@@ -69,6 +69,13 @@ impl Actor for MockAggregator {
 }
 impl AggregatorActor for MockAggregator {}
 
+#[async_trait]
+impl Shutdown for MockAggregator {
+    async fn shutdown(&mut self) {
+        // Mock implementation - no cleanup needed
+    }
+}
+
 impl MockAggregator {
     /// Spawn a mock aggregator actor.
     pub fn spawn(mock_actor: MockAggregator) -> AggregatorHandle {

pipeline/src/aggregator/mod.rs

@@ -35,7 +35,7 @@ use arrow::{
 };
 use arrow_schema::{DataType, SchemaRef};
 use async_trait::async_trait;
-use ceramic_actor::{actor_envelope, Actor, Handler, Message};
+use ceramic_actor::{actor_envelope, Actor, Handler, Message, Shutdown as ActorShutdown};
 use ceramic_core::{StreamId, StreamIdType};
 use cid::Cid;
 use datafusion::{
@@ -67,7 +67,7 @@ use model_instance_patch::ModelInstancePatch;
 use model_patch::ModelPatch;
 use shutdown::{Shutdown, ShutdownSignal};
 use tokio::{select, sync::broadcast, task::JoinHandle};
-use tracing::{debug, error, instrument};
+use tracing::{debug, error, info, instrument};
 
 use crate::{
     cache_table::CacheTable,
@@ -802,38 +802,47 @@ impl Aggregator {
             .await
             .context("writing to mem table data")?;
 
-        let count = self
-            .ctx
-            .table(EVENT_STATES_MEM_TABLE)
-            .await?
-            .count()
-            .await
-            .context("count mem table")?;
+        let count = self.count_cache().await?;
         // If we have enough data cached in memory write it out to persistent store
         if count >= self.max_cached_rows {
-            self.ctx
-                .table(EVENT_STATES_MEM_TABLE)
-                .await?
-                .write_table(
-                    EVENT_STATES_PERSISTENT_TABLE,
-                    DataFrameWriteOptions::new()
-                        .with_partition_by(vec!["event_cid_partition".to_owned()]),
-                )
-                .await
-                .context("writing to persistent table")?;
-            // Clear all data in the memory batch, by writing an empty batch
-            self.ctx
-                .read_batch(RecordBatch::new_empty(schemas::event_states_partitioned()))
-                .context("reading empty batch")?
-                .write_table(
-                    EVENT_STATES_MEM_TABLE,
-                    DataFrameWriteOptions::new().with_insert_operation(InsertOp::Overwrite),
-                )
-                .await
-                .context("clearing mem table")?;
+            self.flush_cache().await?;
         }
         ordered.collect().await.context("collecting ordered events")
     }
+
+    async fn count_cache(&self) -> Result<usize> {
+        self.ctx
+            .table(EVENT_STATES_MEM_TABLE)
+            .await?
+            .count()
+            .await
+            .context("count mem table")
+    }
+
+    async fn flush_cache(&self) -> Result<usize> {
+        let cnt = self.count_cache().await?;
+        self.ctx
+            .table(EVENT_STATES_MEM_TABLE)
+            .await?
+            .write_table(
+                EVENT_STATES_PERSISTENT_TABLE,
+                DataFrameWriteOptions::new()
+                    .with_partition_by(vec!["event_cid_partition".to_owned()]),
+            )
+            .await
+            .context("writing to persistent table")?;
+        // Clear all data in the memory batch, by writing an empty batch
+        self.ctx
+            .read_batch(RecordBatch::new_empty(schemas::event_states_partitioned()))
+            .context("reading empty batch")?
+            .write_table(
+                EVENT_STATES_MEM_TABLE,
+                DataFrameWriteOptions::new().with_insert_operation(InsertOp::Overwrite),
+            )
+            .await
+            .context("clearing mem table")?;
+        Ok(cnt)
+    }
 }
 
 // Construct a column for the field in the UNNAMED_TABLE.
@@ -888,13 +897,33 @@ actor_envelope! {
     AggregatorEnvelope,
     AggregatorActor,
     AggregatorRecorder,
+    with_shutdown,
     SubscribeSince => SubscribeSinceMsg,
     NewConclusionEvents => NewConclusionEventsMsg,
     // TODO: Remove this message and use the analogous message on the Resolver.
     // This way the canonical stream state is provided via the API
     StreamState => StreamStateMsg,
 }
 
+#[async_trait]
+impl ActorShutdown for Aggregator {
+    async fn shutdown(&mut self) {
+        info!("Aggregator shutdown: flushing cached data to persistent storage...");
+
+        match self.flush_cache().await {
+            Ok(rows) => {
+                debug!(
+                    "Aggregator shutdown: successfully flushed {} rows to persistent storage",
+                    rows
+                );
+            }
+            Err(e) => {
+                error!("Aggregator shutdown: failed to flush cached data: {}", e);
+            }
+        }
+    }
+}
+
 #[async_trait]
 impl Handler<SubscribeSinceMsg> for Aggregator {
     async fn handle(