Skip to main content

behavior/supervision/
supervisor.rs

1//! Fleet coordination for supervised stable proxy actors.
2
3use std::time::Duration;
4
5use tokio::time::Instant;
6
7use super::policy::{
8    RestartPolicy, Strategy, SupervisionFailure, SupervisionFailureReaction,
9    retire_on_supervision_failure,
10};
11use super::protocol::{ProxyCommand, SupervisionEvent};
12use super::proxy::Proxy;
13use crate::behavior::{
14    Actions, Address, Behavior, Births, Create, Delivery, Recipient, SendAlgebra, SendProduct,
15    ServiceSends,
16};
17use crate::protocol::{ChildEvent, ObserveChild, WorkerStopped};
18use crate::verdict::{Never, Step};
19use crate::{Become, Exit, RestartDenial, SupervisionFailureReason};
20
21pub type SupervisorSends<A, Sends, C> = SendProduct<
22    Sends,
23    SendProduct<ServiceSends<ObserveChild<A>>, Vec<Delivery<A, ProxyCommand<C>>>>,
24>;
25
26pub type SupervisorActions<B, C> = Actions<
27    <B as Behavior>::Addr,
28    <B as Behavior>::Ph,
29    SupervisorSends<<B as Behavior>::Addr, <B as Behavior>::Sends, C>,
30    Births<Proxy<C>>,
31>;
32
33struct Slot {
34    alive: bool,
35    sequence: u64,
36}
37
38enum ReplacementDecision<A: Address, C: Behavior<Addr = A>> {
39    Retire,
40    Replace(Vec<Delivery<A, ProxyCommand<C>>>),
41    Failed(SupervisionFailure<A>),
42}
43
44pub struct Supervising<B: Behavior, C: Behavior<Ph = Never, Addr = B::Addr>> {
45    inner: B,
46    slots: Vec<(<B::Addr as Address>::Nonce, Slot)>,
47    configured_count: usize,
48    next_sequence: u64,
49    build: fn(usize) -> C,
50    strategy: Strategy,
51    policy: RestartPolicy,
52    max_restarts: u32,
53    window: Duration,
54    restarts: Vec<Instant>,
55    on_failure: SupervisionFailureReaction<B>,
56}
57
58impl<B, C> Supervising<B, C>
59where
60    B: Behavior<Birth = Births<C>>,
61    C: Behavior<Ph = Never, Addr = B::Addr>,
62{
63    #[allow(clippy::too_many_arguments, reason = "hidden by Spec")]
64    /// Construct the concrete supervisor behavior hidden by `Spec`.
65    ///
66    /// # Panics
67    /// Panics only if a fleet index cannot be represented by `u64`.
68    #[must_use]
69    pub fn new(
70        inner: B,
71        nonces: fn(usize) -> <B::Addr as Address>::Nonce,
72        count: usize,
73        build: fn(usize) -> C,
74        strategy: Strategy,
75        policy: RestartPolicy,
76        max_restarts: u32,
77        window: Duration,
78    ) -> Self {
79        let slots = (0..count)
80            .map(|index| {
81                (
82                    nonces(index),
83                    Slot {
84                        alive: true,
85                        sequence: u64::try_from(index).expect("fleet index fits u64"),
86                    },
87                )
88            })
89            .collect();
90        Self {
91            inner,
92            slots,
93            configured_count: count,
94            next_sequence: u64::try_from(count).expect("fleet size fits u64"),
95            build,
96            strategy,
97            policy,
98            max_restarts,
99            window,
100            restarts: Vec::new(),
101            on_failure: retire_on_supervision_failure::<B>,
102        }
103    }
104
105    #[must_use]
106    pub fn with_strategy(mut self, strategy: Strategy) -> Self {
107        self.strategy = strategy;
108        self
109    }
110
111    #[must_use]
112    pub fn with_policy(mut self, policy: RestartPolicy) -> Self {
113        self.policy = policy;
114        self
115    }
116
117    #[must_use]
118    pub fn with_budget(mut self, max: u32, window: Duration) -> Self {
119        self.max_restarts = max;
120        self.window = window;
121        self
122    }
123
124    #[must_use]
125    /// Replace the pure reaction used for typed supervision failures.
126    pub fn with_failure_reaction(mut self, reaction: SupervisionFailureReaction<B>) -> Self {
127        self.on_failure = reaction;
128        self
129    }
130
131    fn position(&self, nonce: <B::Addr as Address>::Nonce) -> Option<usize> {
132        self.slots.iter().position(|(known, _)| *known == nonce)
133    }
134
135    #[must_use]
136    /// Report whether a known supervised proxy is alive.
137    ///
138    /// # Panics
139    /// Panics when `nonce` is not part of this supervisor topology.
140    pub fn is_alive(&self, nonce: <B::Addr as Address>::Nonce) -> bool {
141        self.slots[self.position(nonce).expect("unknown supervised nonce")]
142            .1
143            .alive
144    }
145
146    #[must_use]
147    pub fn child_count(&self) -> usize {
148        self.slots.len()
149    }
150
151    #[must_use]
152    pub fn restarts_in_window(&self) -> usize {
153        self.restarts.len()
154    }
155
156    fn replacement_decision(
157        &mut self,
158        event: &WorkerStopped<B::Addr>,
159    ) -> ReplacementDecision<B::Addr, C> {
160        let dead = self
161            .position(event.proxy)
162            .expect("unknown supervised nonce");
163        let eligible = match self.policy {
164            RestartPolicy::Permanent => true,
165            RestartPolicy::Transient => {
166                !matches!(&event.outcome, Ok(Exit::Normal | Exit::Collected))
167            }
168            RestartPolicy::Temporary => false,
169        };
170        if !eligible {
171            self.slots[dead].1.alive = false;
172            return ReplacementDecision::Retire;
173        }
174        if self.window != Duration::MAX {
175            self.restarts.retain(|stamp| {
176                event
177                    .at
178                    .checked_duration_since(*stamp)
179                    .is_none_or(|age| age <= self.window)
180            });
181        }
182        let sequence = self.slots[dead].1.sequence;
183        let candidates: Vec<usize> = match self.strategy {
184            Strategy::OneForOne => vec![dead],
185            Strategy::OneForAll => self
186                .slots
187                .iter()
188                .enumerate()
189                .filter_map(|(index, (_, slot))| slot.alive.then_some(index))
190                .collect(),
191            Strategy::RestForOne => self
192                .slots
193                .iter()
194                .enumerate()
195                .filter_map(|(index, (_, slot))| {
196                    (slot.alive && slot.sequence >= sequence).then_some(index)
197                })
198                .collect(),
199        };
200        if self.restarts.len() + candidates.len() > self.max_restarts as usize {
201            self.slots[dead].1.alive = false;
202            return ReplacementDecision::Failed(SupervisionFailure {
203                child: event.proxy,
204                outcome: event.outcome,
205                reason: SupervisionFailureReason::RestartDenied(RestartDenial::BudgetExceeded {
206                    restarts_in_window: self.restarts.len(),
207                    replacements_requested: candidates.len(),
208                    maximum_restarts: self.max_restarts,
209                }),
210            });
211        }
212        self.restarts
213            .resize(self.restarts.len() + candidates.len(), event.at);
214        ReplacementDecision::Replace(
215            candidates
216                .into_iter()
217                .map(|index| {
218                    self.slots[index].1.alive = true;
219                    Delivery::new(
220                        Recipient::child(self.slots[index].0),
221                        ProxyCommand::Replace((self.build)(index)),
222                    )
223                })
224                .collect(),
225        )
226    }
227
228    fn react_to_failure(
229        &mut self,
230        failure: &SupervisionFailure<B::Addr>,
231    ) -> Result<Become<B::Addr, B::Ph>, B::Error> {
232        Ok(match (self.on_failure)(&mut self.inner, failure)? {
233            Step::Continue => Step::Continue,
234            Step::Goto(never) => match never {},
235            Step::Stop(exit) => Step::Stop(exit),
236        })
237    }
238
239    fn wrap(
240        &mut self,
241        actions: Actions<B::Addr, B::Ph, B::Sends, Births<C>>,
242    ) -> SupervisorActions<B, C> {
243        let born: Vec<_> = actions.creates.iter().map(|create| create.nonce).collect();
244        for create in &actions.creates {
245            assert!(
246                self.position(create.nonce).is_none(),
247                "a child birth nonce must be fresh"
248            );
249            self.slots.push((
250                create.nonce,
251                Slot {
252                    alive: true,
253                    sequence: self.next_sequence,
254                },
255            ));
256            self.next_sequence = self
257                .next_sequence
258                .checked_add(1)
259                .expect("birth sequence exhausted");
260        }
261        Actions {
262            sends: SendProduct {
263                inner: actions.sends,
264                own: SendProduct {
265                    inner: ServiceSends::new(
266                        born.into_iter()
267                            .map(|nonce| ObserveChild { nonce })
268                            .collect(),
269                    ),
270                    own: Vec::new(),
271                },
272            },
273            creates: actions
274                .creates
275                .into_iter()
276                .map(|create| Create {
277                    nonce: create.nonce,
278                    child: Proxy::new(create.child),
279                    kind: create.kind,
280                })
281                .collect(),
282            become_: actions.become_,
283        }
284    }
285}
286
287impl<B, C, A, Ph, Sends> Behavior for Supervising<B, C>
288where
289    A: Address + Send,
290    Sends: SendAlgebra,
291    B: Behavior<Addr = A, Ph = Ph, Sends = Sends, Birth = Births<C>> + Send,
292    B::Event: ChildEvent<B::Addr> + Send,
293    A::Nonce: From<u64> + Send,
294    B::Msg: Send,
295    C: Behavior<Ph = Never, Addr = B::Addr> + Send,
296{
297    type Addr = A;
298    type Msg = B::Msg;
299    type Event = SupervisionEvent<B::Event, B::Addr>;
300    type Sends = SupervisorSends<A, Sends, C>;
301    type Ph = Ph;
302    type Error = B::Error;
303    type Birth = Births<Proxy<C>>;
304
305    async fn init(&mut self) -> Result<SupervisorActions<B, C>, B::Error> {
306        let actions = self.inner.init().await?;
307        let mut actions = self.wrap(actions);
308        actions.creates.extend(
309            self.slots[..self.configured_count]
310                .iter()
311                .enumerate()
312                .map(|(index, (nonce, _))| Create::birth(*nonce, Proxy::new((self.build)(index)))),
313        );
314        actions.sends.own.inner.extend(
315            self.slots[..self.configured_count]
316                .iter()
317                .map(|(nonce, _)| ObserveChild { nonce: *nonce }),
318        );
319        Ok(actions)
320    }
321
322    async fn step(&mut self, event: Self::Event) -> Result<SupervisorActions<B, C>, B::Error> {
323        match event {
324            SupervisionEvent::WorkerStopped(event) => {
325                let decision = self.replacement_decision(&event);
326                match decision {
327                    ReplacementDecision::Retire => Ok(Actions::cont()),
328                    ReplacementDecision::Replace(replacements) => Ok(Actions {
329                        sends: SendProduct {
330                            inner: B::Sends::empty(),
331                            own: SendProduct {
332                                inner: ServiceSends::empty(),
333                                own: replacements,
334                            },
335                        },
336                        creates: Vec::new(),
337                        become_: Step::Continue,
338                    }),
339                    ReplacementDecision::Failed(failure) => {
340                        Ok(Actions::just(self.react_to_failure(&failure)?))
341                    }
342                }
343            }
344            SupervisionEvent::ChildStopped(event) => {
345                let dead = self
346                    .position(event.nonce)
347                    .expect("unknown supervised nonce");
348                self.slots[dead].1.alive = false;
349                let failure = SupervisionFailure {
350                    child: event.nonce,
351                    outcome: event.outcome,
352                    reason: SupervisionFailureReason::StableChildStopped,
353                };
354                Ok(Actions::just(self.react_to_failure(&failure)?))
355            }
356            SupervisionEvent::Inner(event) => {
357                let actions = self.inner.step(event).await?;
358                Ok(self.wrap(actions))
359            }
360        }
361    }
362}