1use std::any::{Any, TypeId};
2use std::io::Cursor;
3use std::marker::PhantomData;
4use std::panic::AssertUnwindSafe;
5use std::sync::Arc;
6use std::time::Duration;
7
8use anyhow::{Context, Result};
9use futures::FutureExt;
10use rivetkit_core::actor::ShutdownKind;
11use rivetkit_core::error::{ActorLifecycle, ActorRuntime, action_not_found};
12use rivetkit_core::{ActorEvent, ActorEvents, ActorStart, QueueSendResult, QueueSendStatus, Reply};
13use serde::de::DeserializeOwned;
14use tokio::task::JoinHandle;
15use tokio_util::sync::CancellationToken;
16
17use crate::{
18 action::ActionSet,
19 actor::Actor,
20 context::{ConnCtx, Ctx},
21 event::RuntimeEvent,
22 queue::QueueSet,
23};
24
25#[derive(Debug)]
26pub struct Start<A: Actor> {
27 pub ctx: Ctx<A>,
28 pub input: Input<A>,
29 pub is_new: bool,
30 pub snapshot: Snapshot,
31 pub hibernated: Vec<Hibernated<A>>,
32 pub events: Events<A>,
33 #[doc(hidden)]
34 pub startup_ready: Option<tokio::sync::oneshot::Sender<Result<()>>>,
35}
36
37#[derive(Debug)]
38pub struct Input<A: Actor> {
39 bytes: Option<Vec<u8>>,
40 _p: PhantomData<fn() -> A>,
41}
42
43impl<A: Actor> Input<A> {
44 pub fn is_present(&self) -> bool {
45 self.bytes.is_some()
46 }
47
48 pub fn decode(&self) -> Result<A::Input> {
49 match self.bytes.as_deref() {
50 Some(bytes) => decode_cbor(bytes, "actor input"),
51 None if TypeId::of::<A::Input>() == TypeId::of::<()>() => {
52 let unit: Box<dyn Any> = Box::new(());
53 Ok(*unit
54 .downcast::<A::Input>()
55 .expect("unit input type id should downcast"))
56 }
57 None => Err(ActorRuntime::MissingInput.build()),
58 }
59 }
60
61 pub fn decode_or<F>(&self, f: F) -> Result<A::Input>
62 where
63 F: FnOnce() -> A::Input,
64 {
65 match self.bytes.as_deref() {
66 Some(bytes) => decode_cbor(bytes, "actor input"),
67 None => Ok(f()),
68 }
69 }
70
71 pub fn decode_or_default(&self) -> Result<A::Input>
72 where
73 A::Input: Default,
74 {
75 self.decode_or(A::Input::default)
76 }
77
78 pub fn raw(&self) -> Option<&[u8]> {
79 self.bytes.as_deref()
80 }
81}
82
83#[derive(Debug)]
84pub struct Snapshot {
85 is_new: bool,
86 bytes: Option<Vec<u8>>,
87}
88
89impl Snapshot {
90 pub fn is_new(&self) -> bool {
91 self.is_new
92 }
93
94 pub fn decode<S>(&self) -> Result<Option<S>>
95 where
96 S: DeserializeOwned,
97 {
98 let Some(bytes) = self.bytes.as_deref().filter(|bytes| !bytes.is_empty()) else {
99 return Ok(None);
100 };
101 decode_cbor(bytes, "actor snapshot").map(Some)
102 }
103
104 pub fn decode_or_default<S>(&self) -> Result<S>
105 where
106 S: DeserializeOwned + Default,
107 {
108 Ok(self.decode()?.unwrap_or_default())
109 }
110
111 pub fn raw(&self) -> Option<&[u8]> {
112 self.bytes.as_deref()
113 }
114}
115
116#[derive(Debug)]
117pub struct Hibernated<A: Actor> {
118 pub conn: ConnCtx<A>,
119}
120
121#[derive(Debug)]
122pub struct Events<A: Actor> {
123 ctx: Ctx<A>,
124 rx: ActorEvents,
125 _p: PhantomData<fn() -> A>,
126}
127
128impl<A: Actor> Events<A> {
129 pub(crate) async fn recv_raw(&mut self) -> Option<ActorEvent> {
130 self.rx.recv().await
131 }
132
133 pub async fn recv(&mut self) -> Option<RuntimeEvent<A>> {
134 loop {
135 let event = self.rx.recv().await?;
136 if let Some(event) = self.handle_runtime_event(event).await {
137 return Some(wrap_event(event));
138 }
139 }
140 }
141
142 pub fn try_recv(&mut self) -> Option<RuntimeEvent<A>> {
143 while let Some(event) = self.rx.try_recv() {
144 if let Some(event) = self.handle_runtime_event_sync(event) {
145 return Some(wrap_event(event));
146 }
147 }
148 None
149 }
150
151 async fn handle_runtime_event(&self, event: ActorEvent) -> Option<ActorEvent> {
152 match event {
153 ActorEvent::ConnectionOpen { reply, .. } => {
154 reply.send(Ok(()));
155 None
156 }
157 ActorEvent::DisconnectConn { conn_id, reply } => {
158 reply.send(self.ctx.disconnect_conn(&conn_id).await);
159 None
160 }
161 event => Some(event),
162 }
163 }
164
165 fn handle_runtime_event_sync(&self, event: ActorEvent) -> Option<ActorEvent> {
166 match event {
167 ActorEvent::ConnectionOpen { reply, .. } => {
168 reply.send(Ok(()));
169 None
170 }
171 ActorEvent::DisconnectConn { conn_id, reply } => {
172 let ctx = self.ctx.clone();
173 tokio::spawn(async move {
174 reply.send(ctx.disconnect_conn(&conn_id).await);
175 });
176 None
177 }
178 event => Some(event),
179 }
180 }
181}
182
183pub async fn run_actor<A: Actor>(start: Start<A>) -> Result<()> {
184 let Start {
185 ctx,
186 input,
187 is_new,
188 snapshot,
189 hibernated: _,
190 mut events,
191 startup_ready,
192 } = start;
193
194 let state = match snapshot.decode()? {
195 Some(state) => state,
196 None => A::create_state(&ctx, input.decode_or_default()?).await?,
199 };
200 ctx.set_state(state);
201 ctx.clear_state_dirty();
202
203 let actor = Arc::new(A::create(&ctx).await?);
204 if is_new {
205 actor.clone().on_create(ctx.clone()).await?;
206 }
207 actor.clone().on_start(ctx.clone()).await?;
208 if let Some(reply) = startup_ready {
209 let _ = reply.send(Ok(()));
210 }
211
212 let run_cancel = CancellationToken::new();
213 let run_task = spawn_run_task(actor.clone(), ctx.clone(), run_cancel.clone());
214
215 while let Some(event) = events.recv_raw().await {
216 let should_stop = handle_actor_event(actor.clone(), ctx.clone(), event).await?;
217 if should_stop {
218 break;
219 }
220 }
221
222 stop_run_task(&ctx, run_cancel, run_task).await
223}
224
225fn spawn_run_task<A: Actor>(
226 actor: Arc<A>,
227 ctx: Ctx<A>,
228 cancel: CancellationToken,
229) -> JoinHandle<Result<()>> {
230 tokio::spawn(async move {
231 let result = tokio::select! {
232 _ = cancel.cancelled() => return Ok(()),
233 result = AssertUnwindSafe(actor.run(ctx.clone())).catch_unwind() => result,
234 };
235 let result = result.unwrap_or_else(|_| Err(anyhow::anyhow!("actor run handler panicked")));
236 if let Err(error) = &result {
237 report_run_error(&ctx, error);
238 }
239 result
240 })
241}
242
243fn report_run_error<A: Actor>(ctx: &Ctx<A>, error: &anyhow::Error) {
249 if ctx.abort_signal().is_cancelled() {
250 tracing::debug!(?error, "actor run task failed during shutdown");
251 return;
252 }
253 if let Err(stop_error) = ctx.stop_with_error(format!("{error:#}")) {
254 if ctx.inner().is_destroy_requested() {
255 tracing::debug!(
256 ?stop_error,
257 run_error = ?error,
258 "actor run failed while a stop was already requested"
259 );
260 } else {
261 tracing::error!(
262 ?stop_error,
263 run_error = ?error,
264 "failed to report actor run error as errored stop"
265 );
266 }
267 }
268}
269
270async fn stop_run_task<A: Actor>(
271 ctx: &Ctx<A>,
272 cancel: CancellationToken,
273 mut task: JoinHandle<Result<()>>,
274) -> Result<()> {
275 ctx.inner().cancel_actor_abort_signal();
276 match tokio::time::timeout(Duration::from_millis(100), &mut task).await {
277 Ok(result) => result.context("join actor run task")?,
278 Err(_) => {
279 cancel.cancel();
280 task.await.context("join actor run task")?
281 }
282 }
283}
284
285async fn handle_actor_event<A: Actor>(
286 actor: Arc<A>,
287 ctx: Ctx<A>,
288 event: ActorEvent,
289) -> Result<bool> {
290 match event {
291 ActorEvent::Action {
292 name,
293 args,
294 conn,
295 reply,
296 ..
297 } => {
298 let handler_ctx = ctx.with_conn(conn.map(ConnCtx::from));
299 match <A::Actions as ActionSet<A>>::dispatch(
300 actor,
301 handler_ctx.clone(),
302 name.as_str(),
303 args.as_slice(),
304 ) {
305 Some(future) => {
306 spawn_action_reply(handler_ctx, reply, future);
307 }
308 None => {
309 reply.send(Err(action_not_found(name)));
310 }
311 }
312 }
313 ActorEvent::HttpRequest { request, reply } => {
314 reply.send(actor.on_fetch(ctx, request).await);
315 }
316 ActorEvent::QueueSend {
317 name,
318 body,
319 conn,
320 reply,
321 ..
322 } => {
323 let handler_ctx = ctx.with_conn(Some(ConnCtx::from(conn)));
324 match <A::Queue as QueueSet<A>>::dispatch(
325 actor,
326 handler_ctx.clone(),
327 name.as_str(),
328 body.as_slice(),
329 ) {
330 Some(future) => {
331 spawn_queue_reply(handler_ctx, reply, future);
332 }
333 None => {
334 reply.send(Err(ActorRuntime::NotFound {
335 resource: "queue handler".to_owned(),
336 id: name,
337 }
338 .build()));
339 }
340 }
341 }
342 ActorEvent::WebSocketOpen {
343 ws, request, reply, ..
344 } => {
345 reply.send(
346 actor
347 .on_websocket(ctx, ws, request.unwrap_or_default())
348 .await,
349 );
350 }
351 ActorEvent::ConnectionPreflight {
352 conn,
353 params,
354 reply,
355 ..
356 } => {
357 let result = async {
358 let params = decode_conn_params::<A>(¶ms)?;
359 actor
360 .clone()
361 .on_before_connect(ctx.clone(), ¶ms)
362 .await?;
363 let conn_state = actor.clone().create_conn_state(ctx.clone(), params).await?;
364 let conn = ConnCtx::from(conn);
365 conn.set_state(&conn_state)?;
366 actor.on_connect(ctx, conn).await
367 }
368 .await;
369 reply.send(result);
370 }
371 ActorEvent::ConnectionOpen { reply, .. } => {
372 reply.send(Ok(()));
373 }
374 ActorEvent::ConnectionClosed { conn } => {
375 actor.on_disconnect(ctx, ConnCtx::from(conn)).await;
376 }
377 ActorEvent::SubscribeRequest {
378 conn,
379 event_name,
380 reply,
381 } => {
382 reply.send(
383 actor
384 .on_subscribe(ctx, ConnCtx::from(conn), event_name)
385 .await,
386 );
387 }
388 ActorEvent::SerializeState { reply, .. } => {
389 let result = async {
390 if ctx.state_dirty() {
391 actor.on_state_change(ctx.clone()).await?;
392 }
393 let delta = ctx.encode_state_delta()?;
394 ctx.clear_state_dirty();
395 Ok(vec![delta])
396 }
397 .await;
398 reply.send(result);
399 }
400 ActorEvent::RunGracefulCleanup { reason, reply } => {
401 let result = match reason {
402 ShutdownKind::Sleep => actor.on_sleep(ctx).await,
403 ShutdownKind::Destroy => actor.on_destroy(ctx).await,
404 };
405 reply.send(result);
406 return Ok(true);
407 }
408 ActorEvent::DisconnectConn { conn_id, reply } => {
409 reply.send(ctx.disconnect_conn(&conn_id).await);
410 }
411 ActorEvent::WorkflowHistoryRequested { reply } => {
412 reply.send(Err(not_configured("workflow history")));
413 }
414 ActorEvent::WorkflowReplayRequested { reply, .. } => {
415 reply.send(Err(not_configured("workflow replay")));
416 }
417 }
418
419 Ok(false)
420}
421
422fn spawn_action_reply<A: Actor>(
423 ctx: Ctx<A>,
424 reply: Reply<Vec<u8>>,
425 future: crate::action::BoxActionFuture,
426) {
427 tokio::spawn(async move {
428 let abort = ctx.abort_signal();
429 tokio::select! {
430 _ = abort.cancelled() => {
431 reply.send(Err(ActorLifecycle::Stopping.build()));
432 }
433 result = future => {
434 reply.send(result);
435 }
436 }
437 });
438}
439
440fn spawn_queue_reply<A: Actor>(
441 ctx: Ctx<A>,
442 reply: Reply<QueueSendResult>,
443 future: crate::queue::BoxQueueFuture,
444) {
445 tokio::spawn(async move {
446 let abort = ctx.abort_signal();
447 let result = tokio::select! {
448 _ = abort.cancelled() => Err(ActorLifecycle::Stopping.build()),
449 result = future => result.map(|response| QueueSendResult {
450 status: QueueSendStatus::Completed,
451 response,
452 }),
453 };
454 reply.send(result);
455 });
456}
457
458fn not_configured(component: impl Into<String>) -> anyhow::Error {
459 ActorRuntime::NotConfigured {
460 component: component.into(),
461 }
462 .build()
463}
464
465#[doc(hidden)]
466pub fn wrap_start<A: Actor>(core_start: ActorStart) -> Result<Start<A>> {
467 let ActorStart {
468 ctx,
469 input,
470 is_new,
471 snapshot,
472 hibernated,
473 events,
474 startup_ready,
475 } = core_start;
476
477 let hibernated = hibernated
478 .into_iter()
479 .map(|(conn, bytes)| Hibernated {
480 conn: ConnCtx::from({
481 conn.set_state(bytes);
482 conn
483 }),
484 })
485 .collect();
486
487 let ctx = Ctx::new(ctx);
488
489 Ok(Start {
490 ctx: ctx.clone(),
491 input: Input {
492 bytes: input,
493 _p: PhantomData,
494 },
495 is_new,
496 snapshot: Snapshot {
497 is_new,
498 bytes: snapshot,
499 },
500 hibernated,
501 events: Events {
502 ctx,
503 rx: events,
504 _p: PhantomData,
505 },
506 startup_ready,
507 })
508}
509
510fn wrap_event<A: Actor>(event: ActorEvent) -> RuntimeEvent<A> {
511 RuntimeEvent::from_core(event)
512}
513
514fn decode_cbor<T: DeserializeOwned>(bytes: &[u8], label: &str) -> Result<T> {
515 ciborium::from_reader(Cursor::new(bytes)).with_context(|| format!("decode {label} from cbor"))
516}
517
518fn decode_conn_params<A: Actor>(bytes: &[u8]) -> Result<A::ConnParams> {
519 if bytes.is_empty() || bytes == [0xf6] {
520 return Ok(A::ConnParams::default());
521 }
522 decode_cbor(bytes, "connection params")
523}
524
525#[cfg(test)]
526mod tests {
527 use std::future::Future;
528 use std::pin::Pin;
529 use std::sync::OnceLock;
530
531 use async_trait::async_trait;
532 use rivetkit_core::{ConnHandle, QueueNextOpts, StateDelta, WebSocket};
533 use serde::{Deserialize, Serialize};
534 use tokio::sync::mpsc::unbounded_channel;
535 use tokio::sync::{Barrier, oneshot};
536
537 use super::*;
538 use crate::action::{self, Action, Handles, encode_positional};
539 use crate::queue::{HandlesQueue, QueueMessage};
540
541 type BoxTestFuture<T> = Pin<Box<dyn Future<Output = Result<T>> + Send>>;
542 static ACTION_BARRIER: OnceLock<Arc<Barrier>> = OnceLock::new();
543 static QUEUE_PULL_DRAINED: OnceLock<parking_lot::Mutex<Option<oneshot::Sender<Vec<u32>>>>> =
544 OnceLock::new();
545
546 struct EmptyActor;
547
548 impl Actor for EmptyActor {
549 type State = ();
550 type Input = ();
551 type Actions = ();
552 type Events = ();
553 type Queue = ();
554 type ConnParams = ();
555 type ConnState = ();
556 type Action = action::Raw;
557 }
558
559 struct UnitActor;
560
561 impl Actor for UnitActor {
562 type State = ();
563 type Input = UnitInput;
564 type Actions = ();
565 type Events = ();
566 type Queue = ();
567 type ConnParams = ();
568 type ConnState = ();
569 type Action = action::Raw;
570 }
571
572 struct LifecycleActor;
573
574 #[async_trait]
575 impl Actor for LifecycleActor {
576 type State = LifecycleState;
577 type Input = LifecycleInput;
578 type Actions = ();
579 type Events = ();
580 type Queue = ();
581 type ConnParams = ConnParams;
582 type ConnState = ConnState;
583 type Action = action::Raw;
584
585 async fn create_state(_ctx: &Ctx<Self>, input: Self::Input) -> Result<Self::State> {
586 Ok(LifecycleState {
587 count: input.count,
588 log: vec!["create_state".into()],
589 })
590 }
591
592 async fn create(ctx: &Ctx<Self>) -> Result<Self> {
593 ctx.state_mut().log.push("create".into());
594 Ok(Self)
595 }
596
597 async fn run(self: Arc<Self>, ctx: Ctx<Self>) -> Result<()> {
598 ctx.abort_signal().cancelled().await;
599 ctx.state_mut().log.push("run_aborted".into());
600 Ok(())
601 }
602
603 async fn on_create(self: Arc<Self>, ctx: Ctx<Self>) -> Result<()> {
604 ctx.state_mut().log.push("on_create".into());
605 Ok(())
606 }
607
608 async fn on_start(self: Arc<Self>, ctx: Ctx<Self>) -> Result<()> {
609 ctx.state_mut().log.push("on_start".into());
610 Ok(())
611 }
612
613 async fn on_state_change(self: Arc<Self>, ctx: Ctx<Self>) -> Result<()> {
614 ctx.state_mut().log.push("on_state_change".into());
615 Ok(())
616 }
617
618 async fn on_sleep(self: Arc<Self>, ctx: Ctx<Self>) -> Result<()> {
619 ctx.state_mut().log.push("on_sleep".into());
620 Ok(())
621 }
622
623 async fn on_destroy(self: Arc<Self>, ctx: Ctx<Self>) -> Result<()> {
624 ctx.state_mut().log.push("on_destroy".into());
625 Ok(())
626 }
627
628 async fn on_before_connect(
629 self: Arc<Self>,
630 ctx: Ctx<Self>,
631 params: &Self::ConnParams,
632 ) -> Result<()> {
633 if !params.allow {
634 anyhow::bail!("connection rejected");
635 }
636 ctx.state_mut()
637 .log
638 .push(format!("on_before_connect:{}", params.value));
639 Ok(())
640 }
641
642 async fn create_conn_state(
643 self: Arc<Self>,
644 _ctx: Ctx<Self>,
645 params: Self::ConnParams,
646 ) -> Result<Self::ConnState> {
647 Ok(ConnState {
648 value: params.value + 1,
649 })
650 }
651
652 async fn on_connect(self: Arc<Self>, ctx: Ctx<Self>, conn: ConnCtx<Self>) -> Result<()> {
653 let conn_state = conn.state()?;
654 ctx.state_mut()
655 .log
656 .push(format!("on_connect:{}:{}", conn.id(), conn_state.value));
657 Ok(())
658 }
659
660 async fn on_disconnect(self: Arc<Self>, ctx: Ctx<Self>, conn: ConnCtx<Self>) {
661 ctx.state_mut()
662 .log
663 .push(format!("on_disconnect:{}", conn.id()));
664 }
665
666 async fn on_subscribe(
667 self: Arc<Self>,
668 ctx: Ctx<Self>,
669 conn: ConnCtx<Self>,
670 event_name: String,
671 ) -> Result<()> {
672 let conn_state = conn.state()?;
673 ctx.state_mut()
674 .log
675 .push(format!("on_subscribe:{event_name}:{}", conn_state.value));
676 if event_name == "denied" {
677 anyhow::bail!("subscribe denied");
678 }
679 Ok(())
680 }
681 }
682
683 #[derive(Debug, Default, PartialEq, Eq, Serialize, serde::Deserialize)]
684 struct LifecycleInput {
685 count: u32,
686 }
687
688 #[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)]
689 struct LifecycleState {
690 count: u32,
691 log: Vec<String>,
692 }
693
694 #[derive(Debug, Default, PartialEq, Eq, Serialize, serde::Deserialize)]
695 struct UnitInput;
696
697 #[derive(Debug, Default, PartialEq, Eq, Serialize, serde::Deserialize)]
698 struct ExampleState {
699 count: u32,
700 label: String,
701 }
702
703 #[test]
704 fn input_decode_round_trips_unit() {
705 let bytes = cbor(&());
706 let input = Input::<EmptyActor> {
707 bytes: Some(bytes.clone()),
708 _p: PhantomData,
709 };
710
711 assert!(input.is_present());
712 assert_eq!(input.raw(), Some(bytes.as_slice()));
713 assert_eq!(input.decode().expect("decode unit input"), ());
714 }
715
716 #[test]
717 fn input_decode_round_trips_unit_struct() {
718 let input = Input::<UnitActor> {
719 bytes: Some(cbor(&UnitInput)),
720 _p: PhantomData,
721 };
722
723 assert_eq!(input.decode().expect("decode unit struct input"), UnitInput);
724 }
725
726 #[test]
727 fn input_decode_or_default_uses_default_when_missing() {
728 let input = Input::<DefaultActor> {
729 bytes: None,
730 _p: PhantomData,
731 };
732
733 assert_eq!(
734 input.decode_or_default().expect("default input"),
735 DefaultInput { count: 7 }
736 );
737 }
738
739 #[test]
740 fn input_decode_treats_missing_unit_as_unit() {
741 let input = Input::<EmptyActor> {
742 bytes: None,
743 _p: PhantomData,
744 };
745
746 assert_eq!(input.decode().expect("missing unit input"), ());
747 }
748
749 #[test]
750 fn connection_params_decode_null_as_default() {
751 assert_eq!(
752 decode_conn_params::<LifecycleActor>(&[0xf6]).expect("decode null conn params"),
753 ConnParams::default()
754 );
755 assert_eq!(
756 decode_conn_params::<LifecycleActor>(&[]).expect("decode empty conn params"),
757 ConnParams::default()
758 );
759 }
760
761 #[test]
762 fn snapshot_decode_round_trips_map_struct() {
763 let snapshot = Snapshot {
764 is_new: false,
765 bytes: Some(cbor(&ExampleState {
766 count: 9,
767 label: "hi".into(),
768 })),
769 };
770
771 assert!(!snapshot.is_new());
772 assert_eq!(
773 snapshot.decode::<ExampleState>().expect("decode snapshot"),
774 Some(ExampleState {
775 count: 9,
776 label: "hi".into(),
777 })
778 );
779 }
780
781 #[test]
782 fn snapshot_decode_or_default_uses_default_when_missing() {
783 let snapshot = Snapshot {
784 is_new: true,
785 bytes: None,
786 };
787
788 assert!(snapshot.is_new());
789 assert_eq!(
790 snapshot
791 .decode_or_default::<ExampleState>()
792 .expect("default snapshot"),
793 ExampleState::default()
794 );
795 }
796
797 #[test]
798 fn empty_snapshot_decodes_as_missing_without_changing_newness() {
799 let snapshot = Snapshot {
800 is_new: false,
801 bytes: Some(Vec::new()),
802 };
803
804 assert!(!snapshot.is_new());
805 assert_eq!(
806 snapshot
807 .decode_or_default::<ExampleState>()
808 .expect("default empty snapshot"),
809 ExampleState::default()
810 );
811 }
812
813 #[test]
814 fn wrap_start_rehydrates_hibernated_connection_state() {
815 let (tx, rx) = unbounded_channel();
816 drop(tx);
817 let start = wrap_start::<ConnActor>(ActorStart {
818 ctx: rivetkit_core::testing::actor_context("actor-id", "test", Vec::new(), "local"),
819 input: None,
820 is_new: true,
821 snapshot: None,
822 hibernated: vec![(
823 rivetkit_core::ConnHandle::new(
824 "conn-id",
825 cbor(&()),
826 cbor(&ConnState { value: 1 }),
827 true,
828 ),
829 cbor(&ConnState { value: 5 }),
830 )],
831 events: rx.into(),
832 startup_ready: None,
833 })
834 .expect("wrap start");
835
836 assert_eq!(
837 start.hibernated[0]
838 .conn
839 .state()
840 .expect("decode hibernated conn state"),
841 ConnState { value: 5 }
842 );
843 }
844
845 #[test]
846 fn events_try_recv_wraps_core_events() {
847 let (tx, rx) = unbounded_channel();
848 tx.send(ActorEvent::ConnectionClosed {
849 conn: rivetkit_core::ConnHandle::new("conn-id", cbor(&()), cbor(&()), true),
850 })
851 .expect("queue event");
852
853 let mut events = Events::<EmptyActor> {
854 ctx: Ctx::new(rivetkit_core::testing::actor_context(
855 "actor-id",
856 "test",
857 Vec::new(),
858 "local",
859 )),
860 rx: rx.into(),
861 _p: PhantomData,
862 };
863
864 let Some(RuntimeEvent::ConnClosed(closed)) = events.try_recv() else {
865 panic!("expected typed connection-closed event");
866 };
867
868 assert_eq!(closed.conn.id(), "conn-id");
869 }
870
871 #[tokio::test]
872 async fn run_actor_creates_state_and_replies_with_snapshot() {
873 let (tx, rx) = unbounded_channel();
874 let start = lifecycle_start(Some(cbor(&LifecycleInput { count: 3 })), None, rx.into());
875 let actor = tokio::spawn(run_actor::<LifecycleActor>(start));
876
877 let deltas = request_serialize(&tx).await;
878 let state = decode_actor_state(deltas);
879 assert_eq!(state.count, 3);
880 assert_eq!(
881 state.log,
882 [
883 "create_state",
884 "create",
885 "on_create",
886 "on_start",
887 "on_state_change",
888 ]
889 );
890
891 request_sleep(&tx).await;
892 actor.await.expect("join run_actor").expect("run actor");
893 }
894
895 #[tokio::test]
896 async fn run_actor_rehydrates_snapshot_without_on_create() {
897 let snapshot = LifecycleState {
898 count: 8,
899 log: vec!["snapshot".into()],
900 };
901 let (tx, rx) = unbounded_channel();
902 let start = lifecycle_start(None, Some(cbor(&snapshot)), rx.into());
903 let actor = tokio::spawn(run_actor::<LifecycleActor>(start));
904
905 let state = decode_actor_state(request_serialize(&tx).await);
906 assert_eq!(state.count, 8);
907 assert_eq!(
908 state.log,
909 ["snapshot", "create", "on_start", "on_state_change"]
910 );
911
912 request_sleep(&tx).await;
913 actor.await.expect("join run_actor").expect("run actor");
914 }
915
916 #[tokio::test]
917 async fn run_actor_default_fetch_replies_404() {
918 let (tx, rx) = unbounded_channel();
919 let start = lifecycle_start(Some(cbor(&LifecycleInput { count: 1 })), None, rx.into());
920 let actor = tokio::spawn(run_actor::<LifecycleActor>(start));
921
922 let (reply_tx, reply_rx) = oneshot::channel();
923 tx.send(ActorEvent::HttpRequest {
924 request: rivetkit_core::Request::default(),
925 reply: reply_tx.into(),
926 })
927 .expect("send http event");
928
929 let response = reply_rx.await.expect("http reply").expect("http response");
930 assert_eq!(response.status().as_u16(), 404);
931
932 request_sleep(&tx).await;
933 actor.await.expect("join run_actor").expect("run actor");
934 }
935
936 #[tokio::test]
937 async fn run_actor_missing_unit_input_creates_state() {
938 let (tx, rx) = unbounded_channel();
939 let start = unit_creation_start(None, None, rx.into());
940 let actor = tokio::spawn(run_actor::<UnitCreationActor>(start));
941
942 let state = decode_unit_creation_state(request_serialize(&tx).await);
943 assert_eq!(state.created, 1);
944
945 request_sleep(&tx).await;
946 actor.await.expect("join run_actor").expect("run actor");
947 }
948
949 #[tokio::test]
950 async fn run_actor_default_websocket_rejects() {
951 let (tx, rx) = unbounded_channel();
952 let start = lifecycle_start(Some(cbor(&LifecycleInput { count: 1 })), None, rx.into());
953 let actor = tokio::spawn(run_actor::<LifecycleActor>(start));
954
955 let (reply_tx, reply_rx) = oneshot::channel();
956 tx.send(ActorEvent::WebSocketOpen {
957 conn: conn("ws-conn", (), ConnState::default()),
958 ws: WebSocket::new(),
959 request: Some(rivetkit_core::Request::default()),
960 reply: reply_tx.into(),
961 })
962 .expect("send websocket event");
963
964 let error = reply_rx
965 .await
966 .expect("websocket reply")
967 .expect_err("default websocket should reject");
968 assert!(error.to_string().contains("websockets not supported"));
969
970 request_sleep(&tx).await;
971 actor.await.expect("join run_actor").expect("run actor");
972 }
973
974 #[tokio::test]
975 async fn run_actor_connection_hooks_store_state_and_disconnect() {
976 let (tx, rx) = unbounded_channel();
977 let (start, ctx) =
978 lifecycle_start_with_ctx(Some(cbor(&LifecycleInput { count: 1 })), None, rx.into());
979 let actor = tokio::spawn(run_actor::<LifecycleActor>(start));
980 let params = ConnParams {
981 allow: true,
982 value: 41,
983 };
984 let conn = conn("conn-hooks", params.clone(), ConnState::default());
985
986 let (reply_tx, reply_rx) = oneshot::channel();
987 tx.send(ActorEvent::ConnectionPreflight {
988 conn: conn.clone(),
989 params: cbor(¶ms),
990 request: None,
991 reply: reply_tx.into(),
992 })
993 .expect("send connection preflight");
994 reply_rx
995 .await
996 .expect("connection preflight reply")
997 .expect("connection preflight result");
998 assert_eq!(
999 decode_cbor::<ConnState>(&conn.state(), "connection state").expect("conn state"),
1000 ConnState { value: 42 }
1001 );
1002
1003 tx.send(ActorEvent::ConnectionClosed { conn })
1004 .expect("send connection closed");
1005 request_sleep(&tx).await;
1006 actor.await.expect("join run_actor").expect("run actor");
1007
1008 let log = &ctx.state().log;
1009 assert!(log.contains(&"on_before_connect:41".to_owned()));
1010 assert!(log.contains(&"on_connect:conn-hooks:42".to_owned()));
1011 assert!(log.contains(&"on_disconnect:conn-hooks".to_owned()));
1012 }
1013
1014 #[tokio::test]
1015 async fn run_actor_subscribe_hook_allows_and_denies() {
1016 let (tx, rx) = unbounded_channel();
1017 let (start, ctx) =
1018 lifecycle_start_with_ctx(Some(cbor(&LifecycleInput { count: 1 })), None, rx.into());
1019 let actor = tokio::spawn(run_actor::<LifecycleActor>(start));
1020 let conn = conn(
1021 "subscribe-conn",
1022 ConnParams::default(),
1023 ConnState { value: 91 },
1024 );
1025
1026 request_subscribe(&tx, conn.clone(), "chat.message")
1027 .await
1028 .expect("subscribe should be allowed");
1029 let error = request_subscribe(&tx, conn, "denied")
1030 .await
1031 .expect_err("subscribe should be denied");
1032 assert!(error.to_string().contains("subscribe denied"));
1033
1034 request_sleep(&tx).await;
1035 actor.await.expect("join run_actor").expect("run actor");
1036
1037 let log = &ctx.state().log;
1038 assert!(log.contains(&"on_subscribe:chat.message:91".to_owned()));
1039 assert!(log.contains(&"on_subscribe:denied:91".to_owned()));
1040 }
1041
1042 #[tokio::test]
1043 async fn run_actor_connection_preflight_rejects_before_connect() {
1044 let (tx, rx) = unbounded_channel();
1045 let (start, ctx) =
1046 lifecycle_start_with_ctx(Some(cbor(&LifecycleInput { count: 1 })), None, rx.into());
1047 let actor = tokio::spawn(run_actor::<LifecycleActor>(start));
1048 let params = ConnParams {
1049 allow: false,
1050 value: 5,
1051 };
1052
1053 let (reply_tx, reply_rx) = oneshot::channel();
1054 tx.send(ActorEvent::ConnectionPreflight {
1055 conn: conn("conn-reject", params.clone(), ConnState::default()),
1056 params: cbor(¶ms),
1057 request: None,
1058 reply: reply_tx.into(),
1059 })
1060 .expect("send rejected connection preflight");
1061
1062 let error = reply_rx
1063 .await
1064 .expect("connection preflight reply")
1065 .expect_err("connection preflight should reject");
1066 assert!(error.to_string().contains("connection rejected"));
1067
1068 request_sleep(&tx).await;
1069 actor.await.expect("join run_actor").expect("run actor");
1070 assert!(
1071 !ctx.state()
1072 .log
1073 .iter()
1074 .any(|entry| entry.starts_with("on_connect:conn-reject"))
1075 );
1076 }
1077
1078 #[tokio::test]
1079 async fn run_actor_cancels_run_with_abort_signal() {
1080 let (tx, rx) = unbounded_channel();
1081 let (start, ctx) =
1082 lifecycle_start_with_ctx(Some(cbor(&LifecycleInput { count: 2 })), None, rx.into());
1083 let actor = tokio::spawn(run_actor::<LifecycleActor>(start));
1084
1085 request_sleep(&tx).await;
1086 actor.await.expect("join run_actor").expect("run actor");
1087
1088 assert!(ctx.state().log.iter().any(|entry| entry == "run_aborted"));
1089 }
1090
1091 #[tokio::test]
1092 async fn run_actor_destroy_cleanup_fires() {
1093 let (tx, rx) = unbounded_channel();
1094 let (start, ctx) =
1095 lifecycle_start_with_ctx(Some(cbor(&LifecycleInput { count: 2 })), None, rx.into());
1096 let actor = tokio::spawn(run_actor::<LifecycleActor>(start));
1097
1098 request_destroy(&tx).await;
1099 actor.await.expect("join run_actor").expect("run actor");
1100
1101 assert!(ctx.state().log.iter().any(|entry| entry == "on_destroy"));
1102 }
1103
1104 #[tokio::test]
1105 async fn run_actor_error_requests_errored_stop() {
1106 let (tx, rx) = unbounded_channel();
1107 let (start, ctx) = unit_start_with_ctx::<FailingRunActor>(rx.into());
1108 ctx.inner().set_started(true);
1111 let actor = tokio::spawn(run_actor::<FailingRunActor>(start));
1112
1113 for _ in 0..1000 {
1117 if ctx.inner().is_destroy_requested() {
1118 break;
1119 }
1120 tokio::task::yield_now().await;
1121 }
1122 assert!(
1123 ctx.inner().is_destroy_requested(),
1124 "run error should request an errored stop"
1125 );
1126
1127 drop(tx);
1128 let error = actor
1129 .await
1130 .expect("join run_actor")
1131 .expect_err("run_actor should propagate the run error");
1132 assert!(format!("{error:#}").contains("boom in run"));
1133 }
1134
1135 #[tokio::test]
1136 async fn run_actor_panic_requests_errored_stop() {
1137 let (tx, rx) = unbounded_channel();
1138 let (start, ctx) = unit_start_with_ctx::<PanickingRunActor>(rx.into());
1139 ctx.inner().set_started(true);
1140 let actor = tokio::spawn(run_actor::<PanickingRunActor>(start));
1141
1142 for _ in 0..1000 {
1145 if ctx.inner().is_destroy_requested() {
1146 break;
1147 }
1148 tokio::task::yield_now().await;
1149 }
1150 assert!(
1151 ctx.inner().is_destroy_requested(),
1152 "run panic should request an errored stop"
1153 );
1154
1155 drop(tx);
1156 let error = actor
1157 .await
1158 .expect("join run_actor")
1159 .expect_err("run_actor should propagate the panic as an error");
1160 assert!(format!("{error:#}").contains("panicked"));
1161 }
1162
1163 #[tokio::test]
1164 async fn run_actor_shutdown_error_is_not_reported_as_crash() {
1165 let (tx, rx) = unbounded_channel();
1166 let (start, ctx) = unit_start_with_ctx::<ShutdownErrRunActor>(rx.into());
1167 ctx.inner().set_started(true);
1168 let actor = tokio::spawn(run_actor::<ShutdownErrRunActor>(start));
1169
1170 drop(tx);
1173 actor
1174 .await
1175 .expect("join run_actor")
1176 .expect_err("run error should still propagate during shutdown");
1177 assert!(
1178 !ctx.inner().is_destroy_requested(),
1179 "shutdown-induced run errors must not be reported as crashes"
1180 );
1181 }
1182
1183 #[tokio::test]
1184 async fn run_actor_dispatches_typed_actions_by_arg_shape() {
1185 let (tx, rx) = unbounded_channel();
1186 let start = action_start(rx.into());
1187 let actor = tokio::spawn(run_actor::<ActionActor>(start));
1188
1189 assert_eq!(
1190 decode_cbor::<u32>(
1191 &request_action(
1192 &tx,
1193 "add",
1194 &encode_positional(&Add { left: 2, right: 3 }).expect("encode add"),
1195 None,
1196 )
1197 .await
1198 .expect("add result"),
1199 "add output",
1200 )
1201 .expect("decode add output"),
1202 5
1203 );
1204 assert_eq!(
1205 decode_cbor::<u32>(
1206 &request_action(
1207 &tx,
1208 "scale",
1209 &encode_positional(&Scale(4, 5)).expect("encode scale"),
1210 None,
1211 )
1212 .await
1213 .expect("scale result"),
1214 "scale output",
1215 )
1216 .expect("decode scale output"),
1217 20
1218 );
1219 assert_eq!(
1220 decode_cbor::<String>(
1221 &request_action(
1222 &tx,
1223 "echo",
1224 &encode_positional(&Echo("hi".to_owned())).expect("encode echo"),
1225 None,
1226 )
1227 .await
1228 .expect("echo result"),
1229 "echo output",
1230 )
1231 .expect("decode echo output"),
1232 "hi"
1233 );
1234 assert_eq!(
1235 decode_cbor::<String>(
1236 &request_action(
1237 &tx,
1238 "ping",
1239 &encode_positional(&Ping).expect("encode ping"),
1240 None,
1241 )
1242 .await
1243 .expect("ping result"),
1244 "ping output",
1245 )
1246 .expect("decode ping output"),
1247 "pong"
1248 );
1249
1250 request_sleep(&tx).await;
1251 actor.await.expect("join run_actor").expect("run actor");
1252 }
1253
1254 #[tokio::test]
1255 async fn run_actor_dispatches_actions_concurrently() {
1256 let (tx, rx) = unbounded_channel();
1257 let _ = ACTION_BARRIER.set(Arc::new(Barrier::new(2)));
1258 let start = action_start(rx.into());
1259 let actor = tokio::spawn(run_actor::<ActionActor>(start));
1260 let first_args = encode_positional(&WaitForPeer {
1261 label: "first".to_owned(),
1262 })
1263 .expect("encode first wait");
1264 let second_args = encode_positional(&WaitForPeer {
1265 label: "second".to_owned(),
1266 })
1267 .expect("encode second wait");
1268
1269 let first = request_action_rx(&tx, "wait", &first_args, None);
1270 let second = request_action_rx(&tx, "wait", &second_args, None);
1271
1272 let (first, second) = tokio::time::timeout(Duration::from_secs(1), async move {
1273 tokio::join!(first, second)
1274 })
1275 .await
1276 .expect("concurrent handlers should rendezvous");
1277 assert_eq!(
1278 decode_cbor::<String>(&first.expect("first result"), "first wait output")
1279 .expect("decode first wait"),
1280 "first"
1281 );
1282 assert_eq!(
1283 decode_cbor::<String>(&second.expect("second result"), "second wait output")
1284 .expect("decode second wait"),
1285 "second"
1286 );
1287
1288 request_sleep(&tx).await;
1289 actor.await.expect("join run_actor").expect("run actor");
1290 }
1291
1292 #[tokio::test]
1293 async fn run_actor_action_errors_and_unknown_action_are_structured() {
1294 let (tx, rx) = unbounded_channel();
1295 let start = action_start(rx.into());
1296 let actor = tokio::spawn(run_actor::<ActionActor>(start));
1297
1298 let error = request_action(
1299 &tx,
1300 "fail",
1301 &encode_positional(&Fail).expect("encode fail"),
1302 None,
1303 )
1304 .await
1305 .expect_err("fail action should error");
1306 assert!(error.to_string().contains("intentional action failure"));
1307
1308 let error = request_action(&tx, "missing", &[], None)
1309 .await
1310 .expect_err("missing action should error");
1311 let error = rivet_error::RivetError::extract(&error);
1312 assert_eq!(error.group(), "actor");
1313 assert_eq!(error.code(), "action_not_found");
1314
1315 request_sleep(&tx).await;
1316 actor.await.expect("join run_actor").expect("run actor");
1317 }
1318
1319 #[tokio::test]
1320 async fn run_actor_action_receives_per_call_connection_state() {
1321 let (tx, rx) = unbounded_channel();
1322 let start = action_start(rx.into());
1323 let actor = tokio::spawn(run_actor::<ActionActor>(start));
1324 let conn = conn(
1325 "action-conn",
1326 ConnParams::default(),
1327 ConnState { value: 77 },
1328 );
1329
1330 let output = request_action(
1331 &tx,
1332 "connValue",
1333 &encode_positional(&ConnValue).expect("encode conn value"),
1334 Some(conn),
1335 )
1336 .await
1337 .expect("conn value result");
1338
1339 assert_eq!(
1340 decode_cbor::<u32>(&output, "conn value output").expect("decode conn value"),
1341 77
1342 );
1343
1344 request_sleep(&tx).await;
1345 actor.await.expect("join run_actor").expect("run actor");
1346 }
1347
1348 #[tokio::test]
1349 async fn run_actor_dispatches_typed_queue_send() {
1350 let (tx, rx) = unbounded_channel();
1351 let start = action_start(rx.into());
1352 let actor = tokio::spawn(run_actor::<ActionActor>(start));
1353 let conn = conn("queue-conn", ConnParams::default(), ConnState::default());
1354
1355 let result = request_queue_send(
1356 &tx,
1357 "double",
1358 &cbor(&QueueDouble { value: 21 }),
1359 conn.clone(),
1360 )
1361 .await
1362 .expect("queue result");
1363 assert_eq!(result.status, QueueSendStatus::Completed);
1364 assert_eq!(
1365 decode_cbor::<u32>(
1366 result.response.as_deref().expect("queue response"),
1367 "queue response",
1368 )
1369 .expect("decode queue response"),
1370 42
1371 );
1372
1373 let error = request_queue_send(&tx, "missing", &[], conn)
1374 .await
1375 .expect_err("missing queue should error");
1376 let error = rivet_error::RivetError::extract(&error);
1377 assert_eq!(error.group(), "actor");
1378 assert_eq!(error.code(), "not_found");
1379
1380 request_sleep(&tx).await;
1381 actor.await.expect("join run_actor").expect("run actor");
1382 }
1383
1384 #[tokio::test]
1385 async fn run_actor_can_pull_typed_queue_backlog_until_abort() {
1386 let (done_tx, done_rx) = oneshot::channel();
1387 let drained = QUEUE_PULL_DRAINED.get_or_init(|| parking_lot::Mutex::new(None));
1388 assert!(
1389 drained.lock().replace(done_tx).is_none(),
1390 "queue pull notifier already installed"
1391 );
1392
1393 let (tx, rx) = unbounded_channel();
1394 let start = queue_pull_start(rx.into());
1395 start
1396 .ctx
1397 .queue()
1398 .send("double", &QueueDouble { value: 3 })
1399 .await
1400 .expect("send first queue message");
1401 start
1402 .ctx
1403 .queue()
1404 .send("double", &QueueDouble { value: 4 })
1405 .await
1406 .expect("send second queue message");
1407
1408 let actor = tokio::spawn(run_actor::<QueuePullActor>(start));
1409 let values = done_rx.await.expect("queue drain notification");
1410 assert_eq!(values, vec![3, 4]);
1411
1412 let state = request_serialize(&tx).await;
1413 let [StateDelta::ActorState(bytes)] = state.as_slice() else {
1414 panic!("expected actor state delta");
1415 };
1416 assert_eq!(
1417 decode_cbor::<Vec<u32>>(bytes, "queue pull state").expect("decode state"),
1418 vec![3, 4]
1419 );
1420
1421 request_sleep(&tx).await;
1422 actor.await.expect("join run_actor").expect("run actor");
1423 }
1424
1425 struct DefaultActor;
1426
1427 impl Actor for DefaultActor {
1428 type State = ();
1429 type Input = DefaultInput;
1430 type Actions = ();
1431 type Events = ();
1432 type Queue = ();
1433 type ConnParams = ();
1434 type ConnState = ();
1435 type Action = action::Raw;
1436 }
1437
1438 #[derive(Debug, PartialEq, Eq, Serialize, serde::Deserialize)]
1439 struct DefaultInput {
1440 count: u32,
1441 }
1442
1443 impl Default for DefaultInput {
1444 fn default() -> Self {
1445 Self { count: 7 }
1446 }
1447 }
1448
1449 struct ConnActor;
1450
1451 impl Actor for ConnActor {
1452 type State = ();
1453 type Input = ();
1454 type Actions = ();
1455 type Events = ();
1456 type Queue = ();
1457 type ConnParams = ();
1458 type ConnState = ConnState;
1459 type Action = action::Raw;
1460 }
1461
1462 #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, serde::Deserialize)]
1463 struct ConnParams {
1464 allow: bool,
1465 value: u32,
1466 }
1467
1468 #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, serde::Deserialize)]
1469 struct ConnState {
1470 value: u32,
1471 }
1472
1473 struct UnitCreationActor;
1474
1475 #[async_trait]
1476 impl Actor for UnitCreationActor {
1477 type State = UnitCreationState;
1478 type Input = ();
1479 type Actions = ();
1480 type Events = ();
1481 type Queue = ();
1482 type ConnParams = ();
1483 type ConnState = ();
1484 type Action = action::Raw;
1485
1486 async fn create_state(_ctx: &Ctx<Self>, (): Self::Input) -> Result<Self::State> {
1487 Ok(UnitCreationState { created: 1 })
1488 }
1489
1490 async fn create(_ctx: &Ctx<Self>) -> Result<Self> {
1491 Ok(Self)
1492 }
1493 }
1494
1495 #[derive(Debug, Clone, PartialEq, Eq, Serialize, serde::Deserialize)]
1496 struct UnitCreationState {
1497 created: u32,
1498 }
1499
1500 struct ActionActor;
1501
1502 #[async_trait]
1503 impl Actor for ActionActor {
1504 type State = ();
1505 type Input = ();
1506 type Actions = (Add, Scale, Echo, Ping, Fail, WaitForPeer, ConnValue);
1507 type Events = ();
1508 type Queue = (QueueDouble,);
1509 type ConnParams = ConnParams;
1510 type ConnState = ConnState;
1511 type Action = action::Raw;
1512
1513 async fn create_state(_ctx: &Ctx<Self>, (): Self::Input) -> Result<Self::State> {
1514 Ok(())
1515 }
1516
1517 async fn create(_ctx: &Ctx<Self>) -> Result<Self> {
1518 Ok(Self)
1519 }
1520 }
1521
1522 #[derive(Debug, Clone, Serialize, Deserialize)]
1523 struct Add {
1524 left: u32,
1525 right: u32,
1526 }
1527
1528 impl Action for Add {
1529 type Output = u32;
1530
1531 const NAME: &'static str = "add";
1532 }
1533
1534 impl Handles<Add> for ActionActor {
1535 type Future = BoxTestFuture<u32>;
1536
1537 fn handle(self: Arc<Self>, _ctx: Ctx<Self>, action: Add) -> Self::Future {
1538 Box::pin(async move { Ok(action.left + action.right) })
1539 }
1540 }
1541
1542 #[derive(Debug, Clone, Serialize, Deserialize)]
1543 struct Scale(u32, u32);
1544
1545 impl Action for Scale {
1546 type Output = u32;
1547
1548 const NAME: &'static str = "scale";
1549 }
1550
1551 impl Handles<Scale> for ActionActor {
1552 type Future = BoxTestFuture<u32>;
1553
1554 fn handle(self: Arc<Self>, _ctx: Ctx<Self>, action: Scale) -> Self::Future {
1555 Box::pin(async move { Ok(action.0 * action.1) })
1556 }
1557 }
1558
1559 #[derive(Debug, Clone, Serialize, Deserialize)]
1560 struct Echo(String);
1561
1562 impl Action for Echo {
1563 type Output = String;
1564
1565 const NAME: &'static str = "echo";
1566 }
1567
1568 impl Handles<Echo> for ActionActor {
1569 type Future = BoxTestFuture<String>;
1570
1571 fn handle(self: Arc<Self>, _ctx: Ctx<Self>, action: Echo) -> Self::Future {
1572 Box::pin(async move { Ok(action.0) })
1573 }
1574 }
1575
1576 #[derive(Debug, Clone, Serialize, Deserialize)]
1577 struct Ping;
1578
1579 impl Action for Ping {
1580 type Output = String;
1581
1582 const NAME: &'static str = "ping";
1583 }
1584
1585 impl Handles<Ping> for ActionActor {
1586 type Future = BoxTestFuture<String>;
1587
1588 fn handle(self: Arc<Self>, _ctx: Ctx<Self>, _action: Ping) -> Self::Future {
1589 Box::pin(async move { Ok("pong".to_owned()) })
1590 }
1591 }
1592
1593 #[derive(Debug, Clone, Serialize, Deserialize)]
1594 struct Fail;
1595
1596 impl Action for Fail {
1597 type Output = ();
1598
1599 const NAME: &'static str = "fail";
1600 }
1601
1602 impl Handles<Fail> for ActionActor {
1603 type Future = BoxTestFuture<()>;
1604
1605 fn handle(self: Arc<Self>, _ctx: Ctx<Self>, _action: Fail) -> Self::Future {
1606 Box::pin(async move { anyhow::bail!("intentional action failure") })
1607 }
1608 }
1609
1610 #[derive(Debug, Clone, Serialize, Deserialize)]
1611 struct WaitForPeer {
1612 label: String,
1613 }
1614
1615 impl Action for WaitForPeer {
1616 type Output = String;
1617
1618 const NAME: &'static str = "wait";
1619 }
1620
1621 impl Handles<WaitForPeer> for ActionActor {
1622 type Future = BoxTestFuture<String>;
1623
1624 fn handle(self: Arc<Self>, _ctx: Ctx<Self>, action: WaitForPeer) -> Self::Future {
1625 Box::pin(async move {
1626 ACTION_BARRIER
1627 .get()
1628 .expect("action barrier should be installed")
1629 .wait()
1630 .await;
1631 Ok(action.label)
1632 })
1633 }
1634 }
1635
1636 #[derive(Debug, Clone, Serialize, Deserialize)]
1637 struct ConnValue;
1638
1639 impl Action for ConnValue {
1640 type Output = u32;
1641
1642 const NAME: &'static str = "connValue";
1643 }
1644
1645 impl Handles<ConnValue> for ActionActor {
1646 type Future = BoxTestFuture<u32>;
1647
1648 fn handle(self: Arc<Self>, ctx: Ctx<Self>, _action: ConnValue) -> Self::Future {
1649 Box::pin(async move {
1650 let conn = ctx.conn().context("missing action connection")?;
1651 Ok(conn.state()?.value)
1652 })
1653 }
1654 }
1655
1656 #[derive(Debug, Clone, Serialize, Deserialize)]
1657 struct QueueDouble {
1658 value: u32,
1659 }
1660
1661 impl QueueMessage for QueueDouble {
1662 type Reply = u32;
1663
1664 const NAME: &'static str = "double";
1665 }
1666
1667 impl HandlesQueue<QueueDouble> for ActionActor {
1668 type Future = BoxTestFuture<u32>;
1669
1670 fn handle_queue(self: Arc<Self>, _ctx: Ctx<Self>, message: QueueDouble) -> Self::Future {
1671 Box::pin(async move { Ok(message.value * 2) })
1672 }
1673 }
1674
1675 struct QueuePullActor;
1676
1677 #[async_trait]
1678 impl Actor for QueuePullActor {
1679 type State = Vec<u32>;
1680 type Input = ();
1681 type Actions = ();
1682 type Events = ();
1683 type Queue = ();
1684 type ConnParams = ();
1685 type ConnState = ();
1686 type Action = action::Raw;
1687
1688 async fn create_state(_ctx: &Ctx<Self>, (): Self::Input) -> Result<Self::State> {
1689 Ok(Vec::new())
1690 }
1691
1692 async fn create(_ctx: &Ctx<Self>) -> Result<Self> {
1693 Ok(Self)
1694 }
1695
1696 async fn run(self: Arc<Self>, ctx: Ctx<Self>) -> Result<()> {
1697 let mut values = Vec::new();
1698
1699 for _ in 0..2 {
1700 let message = ctx
1701 .queue()
1702 .next_typed::<QueueDouble>(QueueNextOpts {
1703 completable: true,
1704 ..Default::default()
1705 })
1706 .await?
1707 .context("expected queued message")?;
1708 let value = message.body().value;
1709 message.complete(value * 2).await?;
1710 values.push(value);
1711 }
1712
1713 ctx.state_mut().extend(values.iter().copied());
1714 if let Some(done) = QUEUE_PULL_DRAINED
1715 .get_or_init(|| parking_lot::Mutex::new(None))
1716 .lock()
1717 .take()
1718 {
1719 let _ = done.send(values);
1720 }
1721
1722 ctx.abort_signal().cancelled().await;
1723 Ok(())
1724 }
1725 }
1726
1727 struct FailingRunActor;
1728
1729 #[async_trait]
1730 impl Actor for FailingRunActor {
1731 type State = ();
1732 type Input = ();
1733 type Actions = ();
1734 type Events = ();
1735 type Queue = ();
1736 type ConnParams = ();
1737 type ConnState = ();
1738 type Action = action::Raw;
1739
1740 async fn create_state(_ctx: &Ctx<Self>, (): Self::Input) -> Result<Self::State> {
1741 Ok(())
1742 }
1743
1744 async fn create(_ctx: &Ctx<Self>) -> Result<Self> {
1745 Ok(Self)
1746 }
1747
1748 async fn run(self: Arc<Self>, _ctx: Ctx<Self>) -> Result<()> {
1749 anyhow::bail!("boom in run")
1750 }
1751 }
1752
1753 struct PanickingRunActor;
1754
1755 #[async_trait]
1756 impl Actor for PanickingRunActor {
1757 type State = ();
1758 type Input = ();
1759 type Actions = ();
1760 type Events = ();
1761 type Queue = ();
1762 type ConnParams = ();
1763 type ConnState = ();
1764 type Action = action::Raw;
1765
1766 async fn create_state(_ctx: &Ctx<Self>, (): Self::Input) -> Result<Self::State> {
1767 Ok(())
1768 }
1769
1770 async fn create(_ctx: &Ctx<Self>) -> Result<Self> {
1771 Ok(Self)
1772 }
1773
1774 async fn run(self: Arc<Self>, _ctx: Ctx<Self>) -> Result<()> {
1775 panic!("run task panic under test");
1776 }
1777 }
1778
1779 struct ShutdownErrRunActor;
1780
1781 #[async_trait]
1782 impl Actor for ShutdownErrRunActor {
1783 type State = ();
1784 type Input = ();
1785 type Actions = ();
1786 type Events = ();
1787 type Queue = ();
1788 type ConnParams = ();
1789 type ConnState = ();
1790 type Action = action::Raw;
1791
1792 async fn create_state(_ctx: &Ctx<Self>, (): Self::Input) -> Result<Self::State> {
1793 Ok(())
1794 }
1795
1796 async fn create(_ctx: &Ctx<Self>) -> Result<Self> {
1797 Ok(Self)
1798 }
1799
1800 async fn run(self: Arc<Self>, ctx: Ctx<Self>) -> Result<()> {
1801 ctx.abort_signal().cancelled().await;
1802 anyhow::bail!("wait cancelled by shutdown")
1803 }
1804 }
1805
1806 fn unit_start_with_ctx<A: Actor>(rx: ActorEvents) -> (Start<A>, Ctx<A>) {
1807 let ctx = Ctx::new(rivetkit_core::testing::actor_context(
1808 "actor-id",
1809 "test",
1810 Vec::new(),
1811 "local",
1812 ));
1813
1814 let start = Start {
1815 ctx: ctx.clone(),
1816 input: Input {
1817 bytes: None,
1818 _p: PhantomData,
1819 },
1820 is_new: true,
1821 snapshot: Snapshot {
1822 is_new: true,
1823 bytes: None,
1824 },
1825 hibernated: Vec::new(),
1826 events: Events {
1827 ctx: ctx.clone(),
1828 rx,
1829 _p: PhantomData,
1830 },
1831 startup_ready: None,
1832 };
1833
1834 (start, ctx)
1835 }
1836
1837 fn cbor<T: Serialize>(value: &T) -> Vec<u8> {
1838 let mut encoded = Vec::new();
1839 ciborium::into_writer(value, &mut encoded).expect("encode test value as cbor");
1840 encoded
1841 }
1842
1843 fn lifecycle_start(
1844 input: Option<Vec<u8>>,
1845 snapshot: Option<Vec<u8>>,
1846 rx: ActorEvents,
1847 ) -> Start<LifecycleActor> {
1848 lifecycle_start_with_ctx(input, snapshot, rx).0
1849 }
1850
1851 fn lifecycle_start_with_ctx(
1852 input: Option<Vec<u8>>,
1853 snapshot: Option<Vec<u8>>,
1854 rx: ActorEvents,
1855 ) -> (Start<LifecycleActor>, Ctx<LifecycleActor>) {
1856 let ctx = Ctx::new(rivetkit_core::testing::actor_context(
1857 "actor-id",
1858 "test",
1859 Vec::new(),
1860 "local",
1861 ));
1862
1863 let is_new = snapshot.is_none();
1864 let start = Start {
1865 ctx: ctx.clone(),
1866 input: Input {
1867 bytes: input,
1868 _p: PhantomData,
1869 },
1870 is_new,
1871 snapshot: Snapshot {
1872 is_new,
1873 bytes: snapshot,
1874 },
1875 hibernated: Vec::new(),
1876 events: Events {
1877 ctx: ctx.clone(),
1878 rx,
1879 _p: PhantomData,
1880 },
1881 startup_ready: None,
1882 };
1883
1884 (start, ctx)
1885 }
1886
1887 fn unit_creation_start(
1888 input: Option<Vec<u8>>,
1889 snapshot: Option<Vec<u8>>,
1890 rx: ActorEvents,
1891 ) -> Start<UnitCreationActor> {
1892 let ctx = Ctx::new(rivetkit_core::testing::actor_context(
1893 "actor-id",
1894 "unit-creation",
1895 Vec::new(),
1896 "local",
1897 ));
1898
1899 let is_new = snapshot.is_none();
1900 Start {
1901 ctx: ctx.clone(),
1902 input: Input {
1903 bytes: input,
1904 _p: PhantomData,
1905 },
1906 is_new,
1907 snapshot: Snapshot {
1908 is_new,
1909 bytes: snapshot,
1910 },
1911 hibernated: Vec::new(),
1912 events: Events {
1913 ctx,
1914 rx,
1915 _p: PhantomData,
1916 },
1917 startup_ready: None,
1918 }
1919 }
1920
1921 fn action_start(rx: ActorEvents) -> Start<ActionActor> {
1922 let ctx = Ctx::new(rivetkit_core::testing::actor_context(
1923 "actor-id",
1924 "action-test",
1925 Vec::new(),
1926 "local",
1927 ));
1928
1929 Start {
1930 ctx: ctx.clone(),
1931 input: Input {
1932 bytes: None,
1933 _p: PhantomData,
1934 },
1935 is_new: true,
1936 snapshot: Snapshot {
1937 is_new: true,
1938 bytes: None,
1939 },
1940 hibernated: Vec::new(),
1941 events: Events {
1942 ctx,
1943 rx,
1944 _p: PhantomData,
1945 },
1946 startup_ready: None,
1947 }
1948 }
1949
1950 fn queue_pull_start(rx: ActorEvents) -> Start<QueuePullActor> {
1951 let ctx = Ctx::new(rivetkit_core::testing::actor_context(
1952 "actor-id",
1953 "queue-pull-test",
1954 Vec::new(),
1955 "local",
1956 ));
1957
1958 Start {
1959 ctx: ctx.clone(),
1960 input: Input {
1961 bytes: None,
1962 _p: PhantomData,
1963 },
1964 is_new: true,
1965 snapshot: Snapshot {
1966 is_new: true,
1967 bytes: None,
1968 },
1969 hibernated: Vec::new(),
1970 events: Events {
1971 ctx,
1972 rx,
1973 _p: PhantomData,
1974 },
1975 startup_ready: None,
1976 }
1977 }
1978
1979 async fn request_serialize(
1980 tx: &tokio::sync::mpsc::UnboundedSender<ActorEvent>,
1981 ) -> Vec<StateDelta> {
1982 let (reply_tx, reply_rx) = oneshot::channel();
1983 tx.send(ActorEvent::SerializeState {
1984 reason: rivetkit_core::SerializeStateReason::Save,
1985 reply: reply_tx.into(),
1986 })
1987 .expect("send serialize event");
1988 reply_rx
1989 .await
1990 .expect("serialize reply")
1991 .expect("serialize result")
1992 }
1993
1994 async fn request_sleep(tx: &tokio::sync::mpsc::UnboundedSender<ActorEvent>) {
1995 let (reply_tx, reply_rx) = oneshot::channel();
1996 tx.send(ActorEvent::RunGracefulCleanup {
1997 reason: ShutdownKind::Sleep,
1998 reply: reply_tx.into(),
1999 })
2000 .expect("send sleep cleanup");
2001 reply_rx.await.expect("sleep reply").expect("sleep result");
2002 }
2003
2004 async fn request_action(
2005 tx: &tokio::sync::mpsc::UnboundedSender<ActorEvent>,
2006 name: &str,
2007 args: &[u8],
2008 conn: Option<ConnHandle>,
2009 ) -> Result<Vec<u8>> {
2010 request_action_rx(tx, name, args, conn).await
2011 }
2012
2013 async fn request_action_rx(
2014 tx: &tokio::sync::mpsc::UnboundedSender<ActorEvent>,
2015 name: &str,
2016 args: &[u8],
2017 conn: Option<ConnHandle>,
2018 ) -> Result<Vec<u8>> {
2019 let (reply_tx, reply_rx) = oneshot::channel();
2020 tx.send(ActorEvent::Action {
2021 name: name.to_owned(),
2022 args: args.to_vec(),
2023 conn,
2024 scheduled_fire: None,
2025 reply: reply_tx.into(),
2026 })
2027 .expect("send action event");
2028 reply_rx.await.expect("action reply")
2029 }
2030
2031 async fn request_queue_send(
2032 tx: &tokio::sync::mpsc::UnboundedSender<ActorEvent>,
2033 name: &str,
2034 body: &[u8],
2035 conn: ConnHandle,
2036 ) -> Result<QueueSendResult> {
2037 let (reply_tx, reply_rx) = oneshot::channel();
2038 tx.send(ActorEvent::QueueSend {
2039 name: name.to_owned(),
2040 body: body.to_vec(),
2041 conn,
2042 request: rivetkit_core::Request::default(),
2043 wait: true,
2044 timeout_ms: None,
2045 reply: reply_tx.into(),
2046 })
2047 .expect("send queue event");
2048 reply_rx.await.expect("queue reply")
2049 }
2050
2051 async fn request_subscribe(
2052 tx: &tokio::sync::mpsc::UnboundedSender<ActorEvent>,
2053 conn: ConnHandle,
2054 event_name: &str,
2055 ) -> Result<()> {
2056 let (reply_tx, reply_rx) = oneshot::channel();
2057 tx.send(ActorEvent::SubscribeRequest {
2058 conn,
2059 event_name: event_name.to_owned(),
2060 reply: reply_tx.into(),
2061 })
2062 .expect("send subscribe event");
2063 reply_rx.await.expect("subscribe reply")
2064 }
2065
2066 async fn request_destroy(tx: &tokio::sync::mpsc::UnboundedSender<ActorEvent>) {
2067 let (reply_tx, reply_rx) = oneshot::channel();
2068 tx.send(ActorEvent::RunGracefulCleanup {
2069 reason: ShutdownKind::Destroy,
2070 reply: reply_tx.into(),
2071 })
2072 .expect("send destroy cleanup");
2073 reply_rx
2074 .await
2075 .expect("destroy reply")
2076 .expect("destroy result");
2077 }
2078
2079 fn decode_actor_state(deltas: Vec<StateDelta>) -> LifecycleState {
2080 let [StateDelta::ActorState(bytes)] = deltas.as_slice() else {
2081 panic!("expected one actor state delta");
2082 };
2083 decode_cbor(bytes, "actor state").expect("decode actor state")
2084 }
2085
2086 fn decode_unit_creation_state(deltas: Vec<StateDelta>) -> UnitCreationState {
2087 let [StateDelta::ActorState(bytes)] = deltas.as_slice() else {
2088 panic!("expected one actor state delta");
2089 };
2090 decode_cbor(bytes, "actor state").expect("decode actor state")
2091 }
2092
2093 fn conn<P: Serialize, S: Serialize>(id: &str, params: P, state: S) -> ConnHandle {
2094 ConnHandle::new(id, cbor(¶ms), cbor(&state), true)
2095 }
2096}