1use 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 #[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 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 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}