1use std::time::Duration;
13
14use tokio::time::Instant;
15
16use crate::behavior::{
17 Actions, Address, Behavior, Births, Create, Delivery, Recipient, SendAlgebra, SendProduct,
18 ServiceSends, User, UserEvent,
19};
20use crate::protocol::{
21 ChildEvent, ChildStopped, ObserveChild, PeerEvent, PeerStopped, ReportWorkerStopped,
22 ShutdownEvent, ShutdownRequested, TimeEvent, TimeReached, WorkerEvent, WorkerStopped,
23};
24use crate::verdict::{Never, Step};
25use crate::{Become, Crash, Exit, RestartDenial, SupervisionFailureReason};
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum Strategy {
29 OneForOne,
30 OneForAll,
31 RestForOne,
32}
33
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum RestartPolicy {
36 Permanent,
37 Transient,
38 Temporary,
39}
40
41#[must_use]
42pub const fn restart_one() -> Strategy {
43 Strategy::OneForOne
44}
45
46#[must_use]
47pub const fn restart_all() -> Strategy {
48 Strategy::OneForAll
49}
50
51#[must_use]
52pub const fn restart_rest() -> Strategy {
53 Strategy::RestForOne
54}
55
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63pub struct SupervisionFailure<A: Address> {
64 pub child: A::Nonce,
66 pub outcome: Result<Exit<A>, Crash>,
68 pub reason: SupervisionFailureReason,
70}
71
72impl<A: Address> SupervisionFailure<A> {
73 #[must_use]
76 pub const fn into_exit(self) -> Exit<A> {
77 Exit::SupervisionFailed(self.reason)
78 }
79}
80
81pub type SupervisionFailureReaction<B> =
88 fn(
89 &mut B,
90 &SupervisionFailure<<B as Behavior>::Addr>,
91 ) -> Result<Become<<B as Behavior>::Addr>, <B as Behavior>::Error>;
92
93pub fn retire_on_supervision_failure<B: Behavior>(
102 _behavior: &mut B,
103 _failure: &SupervisionFailure<B::Addr>,
104) -> Result<Become<B::Addr>, B::Error> {
105 Ok(Step::Continue)
106}
107
108pub fn stop_on_supervision_failure<B: Behavior>(
118 _behavior: &mut B,
119 failure: &SupervisionFailure<B::Addr>,
120) -> Result<Become<B::Addr>, B::Error> {
121 Ok(Step::Stop(failure.into_exit()))
122}
123
124#[derive(Clone, PartialEq, Eq)]
125pub enum SupervisionEvent<E, A: Address> {
126 Inner(E),
127 ChildStopped(ChildStopped<A>),
128 WorkerStopped(WorkerStopped<A>),
129}
130
131impl<E, A: Address> ChildEvent<A> for SupervisionEvent<E, A> {
132 fn child_stopped(event: ChildStopped<A>) -> Option<Self> {
133 Some(Self::ChildStopped(event))
134 }
135}
136
137impl<E, A: Address> WorkerEvent<A> for SupervisionEvent<E, A> {
138 fn worker_stopped(event: WorkerStopped<A>) -> Option<Self> {
139 Some(Self::WorkerStopped(event))
140 }
141}
142
143impl<E: UserEvent, A: Address> UserEvent for SupervisionEvent<E, A> {
144 type Addr = E::Addr;
145 type Message = E::Message;
146
147 fn user(from: Self::Addr, message: Self::Message) -> Self {
148 Self::Inner(E::user(from, message))
149 }
150
151 fn into_user(self) -> Result<User<Self::Addr, Self::Message>, Self> {
152 match self {
153 Self::Inner(event) => event.into_user().map_err(Self::Inner),
154 stopped @ (Self::ChildStopped(_) | Self::WorkerStopped(_)) => Err(stopped),
155 }
156 }
157}
158
159impl<E: TimeEvent, A: Address> TimeEvent for SupervisionEvent<E, A> {
160 fn time_reached(event: TimeReached) -> Option<Self> {
161 E::time_reached(event).map(Self::Inner)
162 }
163}
164
165impl<E: PeerEvent<A>, A: Address> PeerEvent<A> for SupervisionEvent<E, A> {
166 fn peer_stopped(event: PeerStopped<A>) -> Option<Self> {
167 E::peer_stopped(event).map(Self::Inner)
168 }
169}
170
171impl<E: ShutdownEvent, A: Address> ShutdownEvent for SupervisionEvent<E, A> {
172 fn shutdown_requested(event: ShutdownRequested) -> Option<Self> {
173 E::shutdown_requested(event).map(Self::Inner)
174 }
175}
176
177#[derive(Debug)]
183pub enum ProxyCommand<C: Behavior> {
184 Forward(C::Msg),
186 Replace(C),
188}
189
190pub type SupervisorSends<A, Sends, C> = SendProduct<
191 Sends,
192 SendProduct<ServiceSends<ObserveChild<A>>, Vec<Delivery<A, ProxyCommand<C>>>>,
193>;
194
195pub type SupervisorActions<B, C> = Actions<
196 <B as Behavior>::Addr,
197 <B as Behavior>::Ph,
198 SupervisorSends<<B as Behavior>::Addr, <B as Behavior>::Sends, C>,
199 Births<Proxy<C>>,
200>;
201
202pub struct Proxy<C: Behavior<Ph = Never>> {
209 worker: Option<C>,
210 generation: u64,
211 worker_alive: bool,
212 pending: Option<C>,
213}
214
215impl<C: Behavior<Ph = Never>> Proxy<C> {
216 #[must_use]
217 pub fn new(worker: C) -> Self {
218 Self {
219 worker: Some(worker),
220 generation: 0,
221 worker_alive: false,
222 pending: None,
223 }
224 }
225}
226
227impl<C> Behavior for Proxy<C>
228where
229 C: Behavior<Ph = Never> + Send,
230 C::Addr: Send,
231 <C::Addr as Address>::Nonce: From<u64> + Send,
232 C::Msg: Send,
233 C: Send,
234{
235 type Addr = C::Addr;
236 type Msg = ProxyCommand<C>;
237 type Event = SupervisionEvent<User<C::Addr, ProxyCommand<C>>, C::Addr>;
238 type Sends = SendProduct<
239 Vec<Delivery<C::Addr, C::Msg>>,
240 SendProduct<
241 ServiceSends<ObserveChild<C::Addr>>,
242 ServiceSends<ReportWorkerStopped<C::Addr>>,
243 >,
244 >;
245 type Ph = Never;
246 type Error = Never;
247 type Birth = Births<C>;
248
249 async fn init(&mut self) -> Result<Actions<C::Addr, Never, Self::Sends, Births<C>>, Never> {
250 let child = self.worker.take().expect("a proxy initializes once");
251 self.worker_alive = true;
252 Ok(Actions {
253 sends: SendProduct {
254 inner: Vec::new(),
255 own: SendProduct {
256 inner: ServiceSends::one(ObserveChild {
257 nonce: <C::Addr as Address>::Nonce::from(self.generation),
258 }),
259 own: ServiceSends::empty(),
260 },
261 },
262 creates: vec![Create::birth(
263 <C::Addr as Address>::Nonce::from(self.generation),
264 child,
265 )],
266 become_: Step::Continue,
267 })
268 }
269
270 async fn step(
271 &mut self,
272 event: Self::Event,
273 ) -> Result<Actions<C::Addr, Never, Self::Sends, Births<C>>, Never> {
274 let SupervisionEvent::Inner(event) = event else {
275 return match event {
276 SupervisionEvent::ChildStopped(event)
277 if event.nonce == <C::Addr as Address>::Nonce::from(self.generation) =>
278 {
279 self.worker_alive = false;
280 let report = ReportWorkerStopped {
281 outcome: event.outcome,
282 at: event.at,
283 };
284 let creates = self.pending.take().map_or_else(Vec::new, |child| {
285 self.generation = self
286 .generation
287 .checked_add(1)
288 .expect("proxy generation exhausted");
289 self.worker_alive = true;
290 vec![Create::replacement_incarnation(
291 <C::Addr as Address>::Nonce::from(self.generation),
292 child,
293 )]
294 });
295 let observes = creates
296 .iter()
297 .map(|create| ObserveChild {
298 nonce: create.nonce,
299 })
300 .collect();
301 Ok(Actions {
302 sends: SendProduct {
303 inner: Vec::new(),
304 own: SendProduct {
305 inner: ServiceSends::new(observes),
306 own: ServiceSends::one(report),
307 },
308 },
309 creates,
310 become_: Step::Continue,
311 })
312 }
313 SupervisionEvent::ChildStopped(_) | SupervisionEvent::WorkerStopped(_) => {
314 Ok(Actions::cont())
315 }
316 SupervisionEvent::Inner(_) => unreachable!(),
317 };
318 };
319 match event.message {
320 ProxyCommand::Forward(message) => Ok(Actions {
321 sends: SendProduct {
322 inner: self
323 .worker_alive
324 .then(|| {
325 Delivery::new(
326 Recipient::child(<C::Addr as Address>::Nonce::from(
327 self.generation,
328 )),
329 message,
330 )
331 })
332 .into_iter()
333 .collect(),
334 own: SendProduct {
335 inner: ServiceSends::empty(),
336 own: ServiceSends::empty(),
337 },
338 },
339 creates: Vec::new(),
340 become_: Step::Continue,
341 }),
342 ProxyCommand::Replace(child) => {
343 if self.worker_alive {
344 self.pending = Some(child);
345 return Ok(Actions::cont());
346 }
347 self.generation = self
348 .generation
349 .checked_add(1)
350 .expect("proxy generation exhausted");
351 self.worker_alive = true;
352 let nonce = <C::Addr as Address>::Nonce::from(self.generation);
353 Ok(Actions {
354 sends: SendProduct {
355 inner: Vec::new(),
356 own: SendProduct {
357 inner: ServiceSends::one(ObserveChild { nonce }),
358 own: ServiceSends::empty(),
359 },
360 },
361 creates: vec![Create::replacement_incarnation(nonce, child)],
362 become_: Step::Continue,
363 })
364 }
365 }
366 }
367}
368
369struct Slot {
370 alive: bool,
371 sequence: u64,
372}
373
374enum ReplacementDecision<A: Address, C: Behavior<Addr = A>> {
375 Retire,
376 Replace(Vec<Delivery<A, ProxyCommand<C>>>),
377 Failed(SupervisionFailure<A>),
378}
379
380pub struct Supervising<B: Behavior, C: Behavior<Ph = Never, Addr = B::Addr>> {
381 inner: B,
382 slots: Vec<(<B::Addr as Address>::Nonce, Slot)>,
383 configured_count: usize,
384 next_sequence: u64,
385 build: fn(usize) -> C,
386 strategy: Strategy,
387 policy: RestartPolicy,
388 max_restarts: u32,
389 window: Duration,
390 restarts: Vec<Instant>,
391 on_failure: SupervisionFailureReaction<B>,
392}
393
394impl<B, C> Supervising<B, C>
395where
396 B: Behavior<Birth = Births<C>>,
397 C: Behavior<Ph = Never, Addr = B::Addr>,
398{
399 #[allow(clippy::too_many_arguments, reason = "hidden by Spec")]
400 #[must_use]
405 pub fn new(
406 inner: B,
407 nonces: fn(usize) -> <B::Addr as Address>::Nonce,
408 count: usize,
409 build: fn(usize) -> C,
410 strategy: Strategy,
411 policy: RestartPolicy,
412 max_restarts: u32,
413 window: Duration,
414 ) -> Self {
415 let slots = (0..count)
416 .map(|index| {
417 (
418 nonces(index),
419 Slot {
420 alive: true,
421 sequence: u64::try_from(index).expect("fleet index fits u64"),
422 },
423 )
424 })
425 .collect();
426 Self {
427 inner,
428 slots,
429 configured_count: count,
430 next_sequence: u64::try_from(count).expect("fleet size fits u64"),
431 build,
432 strategy,
433 policy,
434 max_restarts,
435 window,
436 restarts: Vec::new(),
437 on_failure: retire_on_supervision_failure::<B>,
438 }
439 }
440
441 #[must_use]
442 pub fn with_strategy(mut self, strategy: Strategy) -> Self {
443 self.strategy = strategy;
444 self
445 }
446
447 #[must_use]
448 pub fn with_policy(mut self, policy: RestartPolicy) -> Self {
449 self.policy = policy;
450 self
451 }
452
453 #[must_use]
454 pub fn with_budget(mut self, max: u32, window: Duration) -> Self {
455 self.max_restarts = max;
456 self.window = window;
457 self
458 }
459
460 #[must_use]
461 pub fn with_failure_reaction(mut self, reaction: SupervisionFailureReaction<B>) -> Self {
463 self.on_failure = reaction;
464 self
465 }
466
467 fn position(&self, nonce: <B::Addr as Address>::Nonce) -> Option<usize> {
468 self.slots.iter().position(|(known, _)| *known == nonce)
469 }
470
471 #[must_use]
472 pub fn is_alive(&self, nonce: <B::Addr as Address>::Nonce) -> bool {
477 self.slots[self.position(nonce).expect("unknown supervised nonce")]
478 .1
479 .alive
480 }
481
482 #[must_use]
483 pub fn child_count(&self) -> usize {
484 self.slots.len()
485 }
486
487 #[must_use]
488 pub fn restarts_in_window(&self) -> usize {
489 self.restarts.len()
490 }
491
492 fn replacement_decision(
493 &mut self,
494 event: &WorkerStopped<B::Addr>,
495 ) -> ReplacementDecision<B::Addr, C> {
496 let dead = self
497 .position(event.proxy)
498 .expect("unknown supervised nonce");
499 let eligible = match self.policy {
500 RestartPolicy::Permanent => true,
501 RestartPolicy::Transient => {
502 !matches!(&event.outcome, Ok(Exit::Normal | Exit::Collected))
503 }
504 RestartPolicy::Temporary => false,
505 };
506 if !eligible {
507 self.slots[dead].1.alive = false;
508 return ReplacementDecision::Retire;
509 }
510 if self.window != Duration::MAX {
511 self.restarts.retain(|stamp| {
512 event
513 .at
514 .checked_duration_since(*stamp)
515 .is_none_or(|age| age <= self.window)
516 });
517 }
518 let sequence = self.slots[dead].1.sequence;
519 let candidates: Vec<usize> = match self.strategy {
520 Strategy::OneForOne => vec![dead],
521 Strategy::OneForAll => self
522 .slots
523 .iter()
524 .enumerate()
525 .filter_map(|(index, (_, slot))| slot.alive.then_some(index))
526 .collect(),
527 Strategy::RestForOne => self
528 .slots
529 .iter()
530 .enumerate()
531 .filter_map(|(index, (_, slot))| {
532 (slot.alive && slot.sequence >= sequence).then_some(index)
533 })
534 .collect(),
535 };
536 if self.restarts.len() + candidates.len() > self.max_restarts as usize {
537 self.slots[dead].1.alive = false;
538 return ReplacementDecision::Failed(SupervisionFailure {
539 child: event.proxy,
540 outcome: event.outcome,
541 reason: SupervisionFailureReason::RestartDenied(RestartDenial::BudgetExceeded {
542 restarts_in_window: self.restarts.len(),
543 replacements_requested: candidates.len(),
544 maximum_restarts: self.max_restarts,
545 }),
546 });
547 }
548 self.restarts
549 .resize(self.restarts.len() + candidates.len(), event.at);
550 ReplacementDecision::Replace(
551 candidates
552 .into_iter()
553 .map(|index| {
554 self.slots[index].1.alive = true;
555 Delivery::new(
556 Recipient::child(self.slots[index].0),
557 ProxyCommand::Replace((self.build)(index)),
558 )
559 })
560 .collect(),
561 )
562 }
563
564 fn react_to_failure(
565 &mut self,
566 failure: &SupervisionFailure<B::Addr>,
567 ) -> Result<Become<B::Addr, B::Ph>, B::Error> {
568 Ok(match (self.on_failure)(&mut self.inner, failure)? {
569 Step::Continue => Step::Continue,
570 Step::Goto(never) => match never {},
571 Step::Stop(exit) => Step::Stop(exit),
572 })
573 }
574
575 fn wrap(
576 &mut self,
577 actions: Actions<B::Addr, B::Ph, B::Sends, Births<C>>,
578 ) -> SupervisorActions<B, C> {
579 let born: Vec<_> = actions.creates.iter().map(|create| create.nonce).collect();
580 for create in &actions.creates {
581 assert!(
582 self.position(create.nonce).is_none(),
583 "a child birth nonce must be fresh"
584 );
585 self.slots.push((
586 create.nonce,
587 Slot {
588 alive: true,
589 sequence: self.next_sequence,
590 },
591 ));
592 self.next_sequence = self
593 .next_sequence
594 .checked_add(1)
595 .expect("birth sequence exhausted");
596 }
597 Actions {
598 sends: SendProduct {
599 inner: actions.sends,
600 own: SendProduct {
601 inner: ServiceSends::new(
602 born.into_iter()
603 .map(|nonce| ObserveChild { nonce })
604 .collect(),
605 ),
606 own: Vec::new(),
607 },
608 },
609 creates: actions
610 .creates
611 .into_iter()
612 .map(|create| Create {
613 nonce: create.nonce,
614 child: Proxy::new(create.child),
615 kind: create.kind,
616 })
617 .collect(),
618 become_: actions.become_,
619 }
620 }
621}
622
623impl<B, C, A, Ph, Sends> Behavior for Supervising<B, C>
624where
625 A: Address + Send,
626 Sends: SendAlgebra,
627 B: Behavior<Addr = A, Ph = Ph, Sends = Sends, Birth = Births<C>> + Send,
628 B::Event: ChildEvent<B::Addr> + Send,
629 A::Nonce: From<u64> + Send,
630 B::Msg: Send,
631 C: Behavior<Ph = Never, Addr = B::Addr> + Send,
632{
633 type Addr = A;
634 type Msg = B::Msg;
635 type Event = SupervisionEvent<B::Event, B::Addr>;
636 type Sends = SupervisorSends<A, Sends, C>;
637 type Ph = Ph;
638 type Error = B::Error;
639 type Birth = Births<Proxy<C>>;
640
641 async fn init(&mut self) -> Result<SupervisorActions<B, C>, B::Error> {
642 let actions = self.inner.init().await?;
643 let mut actions = self.wrap(actions);
644 actions.creates.extend(
645 self.slots[..self.configured_count]
646 .iter()
647 .enumerate()
648 .map(|(index, (nonce, _))| Create::birth(*nonce, Proxy::new((self.build)(index)))),
649 );
650 actions.sends.own.inner.extend(
651 self.slots[..self.configured_count]
652 .iter()
653 .map(|(nonce, _)| ObserveChild { nonce: *nonce }),
654 );
655 Ok(actions)
656 }
657
658 async fn step(&mut self, event: Self::Event) -> Result<SupervisorActions<B, C>, B::Error> {
659 match event {
660 SupervisionEvent::WorkerStopped(event) => {
661 let decision = self.replacement_decision(&event);
662 match decision {
663 ReplacementDecision::Retire => Ok(Actions::cont()),
664 ReplacementDecision::Replace(replacements) => Ok(Actions {
665 sends: SendProduct {
666 inner: B::Sends::empty(),
667 own: SendProduct {
668 inner: ServiceSends::empty(),
669 own: replacements,
670 },
671 },
672 creates: Vec::new(),
673 become_: Step::Continue,
674 }),
675 ReplacementDecision::Failed(failure) => {
676 Ok(Actions::just(self.react_to_failure(&failure)?))
677 }
678 }
679 }
680 SupervisionEvent::ChildStopped(event) => {
681 let dead = self
682 .position(event.nonce)
683 .expect("unknown supervised nonce");
684 self.slots[dead].1.alive = false;
685 let failure = SupervisionFailure {
686 child: event.nonce,
687 outcome: event.outcome,
688 reason: SupervisionFailureReason::StableChildStopped,
689 };
690 Ok(Actions::just(self.react_to_failure(&failure)?))
691 }
692 SupervisionEvent::Inner(event) => {
693 let actions = self.inner.step(event).await?;
694 Ok(self.wrap(actions))
695 }
696 }
697 }
698}