Catch violations in times(...) expectations (#2766)

* Catch violations in times(...) expectations

Adds a new quit() method to universe that causes all spawned actors
to quit. It also collects actors exit codes that can be used to verify
that all actors quit gracefully, which in turn allows to fail the test
if mock objects expectations are not met.

Relates to #2754

* Make clippy and rustfmt happy

* Fix test_indexer_trigger_on_memory_limit

* Update quickwit/quickwit-actors/src/actor_handle.rs

Improve panic handling

Co-authored-by: Paul Masurel <[email protected]>

* Add universe.assert_quit() method

Also adds a check to the universe.drop() to make sure that
universe.assert_quit() was indeed called at the end of the test.

* Stop test_ingest_api_source tests from hanging.

* Fix compilation errors when testing a single crate

Seems like rust-lang/cargo#8379 is the issue
here.

* Update assert_quits after merge

* Avoid using map with side-effects.

Co-authored-by: Paul Masurel <[email protected]>

* Switch ActorJoinHandle to Shared

* Explain why we join ingest source first

* There should be only one universe.

* Fix format

---------

Co-authored-by: Paul Masurel <[email protected]>

Author: imotov

quickwit/quickwit-actors/src/actor_handle.rs

@@ -21,20 +21,20 @@ use std::fmt;
 
 use serde::Serialize;
 use tokio::sync::{oneshot, watch};
-use tokio::task::JoinHandle;
 use tracing::error;
 
 use crate::actor_state::ActorState;
 use crate::command::Observe;
 use crate::mailbox::Priority;
 use crate::observation::ObservationType;
+use crate::registry::ActorJoinHandle;
 use crate::{Actor, ActorContext, ActorExitStatus, Command, Mailbox, Observation};
 
 /// An Actor Handle serves as an address to communicate with an actor.
 pub struct ActorHandle<A: Actor> {
     actor_context: ActorContext<A>,
     last_state: watch::Receiver<A::ObservableState>,
-    join_handle: JoinHandle<ActorExitStatus>,
+    join_handle: ActorJoinHandle,
 }
 
 /// Describes the health of a given actor.
@@ -98,7 +98,7 @@ impl<A: Actor> Supervisable for ActorHandle<A> {
 impl<A: Actor> ActorHandle<A> {
     pub(crate) fn new(
         last_state: watch::Receiver<A::ObservableState>,
-        join_handle: JoinHandle<ActorExitStatus>,
+        join_handle: ActorJoinHandle,
         actor_context: ActorContext<A>,
     ) -> Self {
         ActorHandle {
@@ -208,13 +208,7 @@ impl<A: Actor> ActorHandle<A> {
 
     /// Waits until the actor exits by itself. This is the equivalent of `Thread::join`.
     pub async fn join(self) -> (ActorExitStatus, A::ObservableState) {
-        let exit_status = self.join_handle.await.unwrap_or_else(|join_err| {
-            if join_err.is_panic() {
-                ActorExitStatus::Panicked
-            } else {
-                ActorExitStatus::Killed
-            }
-        });
+        let exit_status = self.join_handle.join().await;
         let observation = self.last_state.borrow().clone();
         (exit_status, observation)
     }

quickwit/quickwit-actors/src/mailbox.rs

@@ -530,6 +530,7 @@ mod tests {
         assert!(backpressure_micros_counter.get() < 500);
         processed.await.unwrap();
         assert!(backpressure_micros_counter.get() < 500);
+        universe.assert_quit().await;
     }
 
     #[tokio::test]
@@ -563,6 +564,7 @@ mod tests {
             .await
             .unwrap();
         assert!(backpressure_micros_counter.get() > 1_000u64);
+        universe.assert_quit().await;
     }
 
     #[tokio::test]
@@ -584,6 +586,7 @@ mod tests {
         let elapsed = start.elapsed();
         assert!(elapsed.as_micros() > 1000);
         assert_eq!(backpressure_micros_counter.get(), 0);
+        universe.assert_quit().await;
     }
 
     #[tokio::test]

quickwit/quickwit-actors/src/registry.rs

@@ -19,21 +19,25 @@
 
 use std::any::{Any, TypeId};
 use std::collections::HashMap;
+use std::pin::Pin;
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
 
 use async_trait::async_trait;
-use futures::future;
+use futures::future::{self, Shared};
+use futures::{Future, FutureExt};
 use serde::Serialize;
 use serde_json::Value as JsonValue;
+use tokio::task::JoinHandle;
 
 use crate::command::Observe;
 use crate::mailbox::WeakMailbox;
-use crate::{Actor, Mailbox};
+use crate::{Actor, ActorExitStatus, Command, Mailbox};
 
 struct TypedJsonObservable<A: Actor> {
     actor_instance_id: String,
     weak_mailbox: WeakMailbox<A>,
+    join_handle: ActorJoinHandle,
 }
 
 #[async_trait]
@@ -42,6 +46,8 @@ trait JsonObservable: Sync + Send {
     fn any(&self) -> &dyn Any;
     fn actor_instance_id(&self) -> &str;
     async fn observe(&self) -> Option<JsonValue>;
+    async fn quit(&self) -> ActorExitStatus;
+    async fn join(&self) -> ActorExitStatus;
 }
 
 #[async_trait]
@@ -64,6 +70,17 @@ impl<A: Actor> JsonObservable for TypedJsonObservable<A> {
         let state: <A as Actor>::ObservableState = oneshot_rx.await.ok()?;
         serde_json::to_value(&state).ok()
     }
+
+    async fn quit(&self) -> ActorExitStatus {
+        if let Some(mailbox) = self.weak_mailbox.upgrade() {
+            let _ = mailbox.send_message_with_high_priority(Command::Quit);
+        }
+        self.join().await
+    }
+
+    async fn join(&self) -> ActorExitStatus {
+        self.join_handle.join().await
+    }
 }
 
 #[derive(Default, Clone)]
@@ -104,7 +121,7 @@ pub struct ActorObservation {
 }
 
 impl ActorRegistry {
-    pub fn register<A: Actor>(&self, mailbox: &Mailbox<A>) {
+    pub fn register<A: Actor>(&self, mailbox: &Mailbox<A>, join_handle: ActorJoinHandle) {
         let typed_id = TypeId::of::<A>();
         let actor_instance_id = mailbox.actor_instance_id().to_string();
         let weak_mailbox = mailbox.downgrade();
@@ -117,6 +134,7 @@ impl ActorRegistry {
             .push(Arc::new(TypedJsonObservable {
                 weak_mailbox,
                 actor_instance_id,
+                join_handle,
             }));
     }
 
@@ -162,6 +180,31 @@ impl ActorRegistry {
             registry_for_type.gc();
         }
     }
+
+    pub async fn quit(&self) -> Vec<ActorExitStatus> {
+        let mut obs_futures = Vec::new();
+        for registry_for_type in self.actors.read().unwrap().values() {
+            for obs in &registry_for_type.observables {
+                let obs_clone = obs.clone();
+                obs_futures.push(async move { obs_clone.quit().await });
+            }
+        }
+        let res = future::join_all(obs_futures).await;
+        res.into_iter().collect()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.actors
+            .read()
+            .unwrap()
+            .values()
+            .all(|registry_for_type| {
+                registry_for_type
+                    .observables
+                    .iter()
+                    .all(|obs| obs.is_disconnected())
+            })
+    }
 }
 
 fn get_iter<A: Actor>(
@@ -181,6 +224,37 @@ fn get_iter<A: Actor>(
         .filter(|mailbox| !mailbox.is_disconnected())
 }
 
+/// This structure contains an optional exit handle. The handle is present
+/// until the join() method is called.
+#[derive(Clone)]
+pub(crate) struct ActorJoinHandle {
+    holder: Shared<Pin<Box<dyn Future<Output = ActorExitStatus> + Send>>>,
+}
+
+impl ActorJoinHandle {
+    pub(crate) fn new(join_handle: JoinHandle<ActorExitStatus>) -> Self {
+        ActorJoinHandle {
+            holder: Self::inner_join(join_handle).boxed().shared(),
+        }
+    }
+
+    async fn inner_join(join_handle: JoinHandle<ActorExitStatus>) -> ActorExitStatus {
+        join_handle.await.unwrap_or_else(|join_err| {
+            if join_err.is_panic() {
+                ActorExitStatus::Panicked
+            } else {
+                ActorExitStatus::Killed
+            }
+        })
+    }
+
+    /// Joins the actor and returns its exit status on the frist invocation.
+    /// Returns None afterwards.
+    pub(crate) async fn join(&self) -> ActorExitStatus {
+        self.holder.clone().await
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::time::Duration;
@@ -194,6 +268,7 @@ mod tests {
         let universe = Universe::with_accelerated_time();
         let (_mailbox, _handle) = universe.spawn_builder().spawn(test_actor);
         let _actor_mailbox = universe.get_one::<PingReceiverActor>().unwrap();
+        universe.assert_quit().await;
     }
 
     #[tokio::test]
@@ -222,5 +297,6 @@ mod tests {
         let (_mailbox, _handle) = universe.spawn_builder().spawn(test_actor);
         let obs = universe.observe(Duration::from_millis(1000)).await;
         assert_eq!(obs.len(), 1);
+        universe.assert_quit().await;
     }
 }

quickwit/quickwit-actors/src/scheduler.rs

@@ -407,6 +407,7 @@ mod tests {
         assert_eq!(count.load(Ordering::SeqCst), 0);
         universe.sleep(Duration::from_millis(15)).await;
         assert_eq!(count.load(Ordering::SeqCst), 1);
+        universe.assert_quit().await;
     }
 
     #[tokio::test]
@@ -425,5 +426,6 @@ mod tests {
         let elapsed = start.elapsed();
         // The whole point is to accelerate time.
         assert!(elapsed.as_millis() < 50);
+        universe.assert_quit().await;
     }
 }

quickwit/quickwit-actors/src/spawn_builder.rs

@@ -24,7 +24,7 @@ use tracing::{debug, error, info};
 
 use crate::envelope::Envelope;
 use crate::mailbox::{create_mailbox, Inbox};
-use crate::registry::ActorRegistry;
+use crate::registry::{ActorJoinHandle, ActorRegistry};
 use crate::scheduler::{NoAdvanceTimeGuard, SchedulerClient};
 use crate::supervisor::Supervisor;
 use crate::{
@@ -161,11 +161,11 @@ impl<A: Actor> SpawnBuilder<A> {
         let (ctx, inbox, state_rx) = self.create_actor_context_and_inbox(&actor);
         debug!(actor_id = %ctx.actor_instance_id(), "spawn-actor");
         let mailbox = ctx.mailbox().clone();
-        ctx.registry().register(&mailbox);
         let ctx_clone = ctx.clone();
         let loop_async_actor_future =
             async move { actor_loop(actor, inbox, no_advance_time_guard, ctx).await };
-        let join_handle = runtime_handle.spawn(loop_async_actor_future);
+        let join_handle = ActorJoinHandle::new(runtime_handle.spawn(loop_async_actor_future));
+        ctx_clone.registry().register(&mailbox, join_handle.clone());
         let actor_handle = ActorHandle::new(state_rx, join_handle, ctx_clone);
         (mailbox, actor_handle)
     }

quickwit/quickwit-actors/src/supervisor.rs

@@ -285,6 +285,10 @@ mod tests {
                 num_kills: 0
             }
         );
+        assert!(!matches!(
+            supervisor_handle.quit().await.0,
+            ActorExitStatus::Panicked
+        ));
     }
 
     #[tokio::test]
@@ -313,6 +317,10 @@ mod tests {
                 num_kills: 0
             }
         );
+        assert!(!matches!(
+            supervisor_handle.quit().await.0,
+            ActorExitStatus::Panicked
+        ));
     }
 
     #[tokio::test]
@@ -354,6 +362,10 @@ mod tests {
                 num_kills: 1
             }
         );
+        assert!(!matches!(
+            supervisor_handle.quit().await.0,
+            ActorExitStatus::Panicked
+        ));
     }
 
     #[tokio::test]

quickwit/quickwit-actors/src/tests.rs

@@ -283,6 +283,7 @@ async fn test_timeouting_actor() {
     assert_eq!(buggy_handle.harvest_health(), Health::Healthy);
     universe.sleep(crate::HEARTBEAT * 2).await;
     assert_eq!(buggy_handle.harvest_health(), Health::FailureOrUnhealthy);
+    buggy_handle.kill().await;
 }
 
 #[tokio::test]
@@ -305,6 +306,7 @@ async fn test_pause_actor() {
         .is_ok());
     let end_state = ping_handle.process_pending_and_observe().await.state;
     assert_eq!(end_state, 1000);
+    universe.assert_quit().await;
 }
 
 #[tokio::test]
@@ -320,6 +322,7 @@ async fn test_actor_running_states() {
     assert_eq!(*obs, 10);
     universe.sleep(Duration::from_millis(1)).await;
     assert!(ping_handle.state() == ActorState::Idle);
+    universe.assert_quit().await;
 }
 
 #[derive(Clone, Debug, Default, Serialize)]
@@ -551,6 +554,7 @@ async fn test_actor_return_response() -> anyhow::Result<()> {
     let plus_two_plus_four = mailbox.send_message(AddOperand(4)).await?;
     assert_eq!(plus_two.await.unwrap(), 2);
     assert_eq!(plus_two_plus_four.await.unwrap(), 6);
+    universe.assert_quit().await;
     Ok(())
 }
 
@@ -631,4 +635,5 @@ async fn test_drain_is_called() {
             drain_calls_count: 2
         }
     );
+    universe.assert_quit().await;
 }