behavior/
pool.rs

1//! A bounded, FIFO worker pool expressed entirely as a pure behavior.
2//!
3//! Pool scheduling is a derived Bombay construction, not an actor-model
4//! primitive. Runtime installation, delivery, and observation remain effects
5//! for an interpreter; this module owns only their typed protocol and fold.
6
7use core::convert::Infallible;
8use core::marker::PhantomData;
9use std::collections::{BTreeMap, VecDeque};
10use std::time::Duration;
11
12use crate::{
13    Actions, Address, Behavior, Births, Crash, CreationRejection, Delivery, Exit, Never, Own,
14    Proxy, ProxyCommand, Recipient, RestartPolicy, SendAlgebra, SendInput, Strategy,
15    SupervisionEvent, Supervisor, SupervisorSends, User, WorkerCreationResolved, WorkerStopped,
16    delegate_transition,
17};
18
19/// Caller-chosen identity used to correlate pool responses.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
21pub struct JobId(pub u64);
22
23/// Pool-owned correlation token for one exact dispatch attempt.
24///
25/// This is not an actor identity or evidence that delivery occurred.
26#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
27pub struct AssignmentId(pub u64);
28
29/// One assignment accepted by a worker behavior.
30#[derive(Debug, Clone, PartialEq, Eq)]
31pub struct PoolAssignment<J> {
32    pub assignment: AssignmentId,
33    pub job: JobId,
34    pub payload: J,
35}
36
37/// Messages accepted by a pool coordinator.
38#[derive(Clone, PartialEq, Eq)]
39pub enum PoolMessage<A: Address, J, R> {
40    Submit {
41        job: JobId,
42        payload: J,
43        reply_to: Recipient<A, PoolResponse<J, R, A>>,
44    },
45    Completed {
46        worker: A::Nonce,
47        assignment: AssignmentId,
48        result: R,
49    },
50}
51
52/// Messages accepted by a key-persistent pool coordinator.
53///
54/// `Rebalance` is the only input that can change an established key binding.
55/// It affects later submissions; jobs already accepted retain their selected
56/// stable worker slot.
57#[derive(Clone, PartialEq, Eq)]
58pub enum KeyedPoolMessage<A: Address, K, J, R> {
59    Submit {
60        key: K,
61        job: JobId,
62        payload: J,
63        reply_to: Recipient<A, PoolResponse<J, R, A>>,
64    },
65    Completed {
66        worker: A::Nonce,
67        assignment: AssignmentId,
68        result: R,
69    },
70    Rebalance {
71        key: K,
72        worker: A::Nonce,
73    },
74}
75
76/// Why a submitted job was not accepted by the pool.
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78pub enum PoolRejection {
79    BacklogFull,
80    /// The key's selected stable slot is unknown or permanently retired.
81    AffinityUnavailable,
82}
83
84/// Why an accepted assignment ended without a worker completion.
85#[derive(Clone, PartialEq, Eq)]
86pub enum PoolInterruption<A: Address> {
87    WorkerStopped {
88        worker: A::Nonce,
89        outcome: Result<Exit<A>, Crash>,
90    },
91    NoRecoverableWorkers,
92    /// The job's selected stable slot retired while the job was queued.
93    AffinityRetired {
94        worker: A::Nonce,
95        reason: WorkerRetirement,
96    },
97}
98
99/// Complete response protocol for one submitted job.
100#[derive(Clone, PartialEq, Eq)]
101pub enum PoolResponse<J, R, A: Address> {
102    Accepted {
103        job: JobId,
104    },
105    Rejected {
106        job: JobId,
107        payload: J,
108        reason: PoolRejection,
109    },
110    Completed {
111        job: JobId,
112        result: R,
113    },
114    Interrupted {
115        job: JobId,
116        payload: J,
117        reason: PoolInterruption<A>,
118    },
119}
120
121/// Bombay policy for an assigned job whose worker incarnation stops.
122#[derive(Debug, Clone, Copy, PartialEq, Eq)]
123pub enum InterruptionPolicy {
124    /// End pool ownership and report the still-owned job to its submitter.
125    Fail,
126    /// Put the job at the front of the backlog for at-least-once assignment.
127    Retry,
128}
129
130/// Public, payload-free view of one stable worker slot.
131#[derive(Debug, Clone, Copy, PartialEq, Eq)]
132pub enum WorkerPhase {
133    Installing,
134    Idle,
135    Assigned {
136        assignment: AssignmentId,
137        job: JobId,
138    },
139    Retired {
140        reason: WorkerRetirement,
141    },
142}
143
144/// Why a stable worker slot is no longer eligible for dispatch.
145#[derive(Debug, Clone, Copy, PartialEq, Eq)]
146pub enum WorkerRetirement {
147    CreationRejected(CreationRejection),
148    ReplacementUnavailable,
149}
150
151/// Invalid static pool topology.
152#[derive(Debug, Clone, Copy, PartialEq, Eq)]
153pub enum PoolConfigError<N> {
154    /// No stable worker slot exists, so accepted ownership could never end.
155    NoWorkers,
156    /// Two configured positions selected the same stable worker nonce.
157    DuplicateWorker(N),
158}
159
160/// Pure, statically dispatched policy for a previously unseen affinity key.
161pub trait AffinitySelector<K, N> {
162    /// Select the stable worker nonce for a key that has no binding yet.
163    fn select(&self, key: &K) -> N;
164}
165
166impl<K, N, F> AffinitySelector<K, N> for F
167where
168    F: Fn(&K) -> N,
169{
170    fn select(&self, key: &K) -> N {
171        self(key)
172    }
173}
174
175/// Typed rejection of an event that cannot apply to the current pool state.
176#[derive(Debug, Clone, Copy, PartialEq, Eq)]
177pub enum PoolError<N> {
178    UnknownWorker(N),
179    CompletionForUnavailableWorker {
180        worker: N,
181        phase: WorkerPhase,
182    },
183    StaleCompletion {
184        worker: N,
185        expected: AssignmentId,
186        received: AssignmentId,
187    },
188    WorkerStoppedWhileUnavailable {
189        worker: N,
190        phase: WorkerPhase,
191    },
192    CreationResolvedWhileUnavailable {
193        worker: N,
194        phase: WorkerPhase,
195    },
196    RebalanceToRetiredWorker {
197        worker: N,
198        reason: WorkerRetirement,
199    },
200}
201
202struct AcceptedJob<A: Address, J, R> {
203    id: JobId,
204    payload: J,
205    reply_to: Recipient<A, PoolResponse<J, R, A>>,
206    interruption: Option<PoolInterruption<A>>,
207    target: Option<A::Nonce>,
208}
209
210struct QueuedJob<A: Address, J, R> {
211    accepted: AcceptedJob<A, J, R>,
212    dispatch_payload: J,
213}
214
215enum SlotState<A: Address, J, R> {
216    Installing,
217    Idle,
218    Assigned {
219        assignment: AssignmentId,
220        job: AcceptedJob<A, J, R>,
221    },
222    Retired {
223        reason: WorkerRetirement,
224    },
225}
226
227struct Slot<A: Address, J, R> {
228    nonce: A::Nonce,
229    state: SlotState<A, J, R>,
230}
231
232struct PlannedDispatch {
233    slot_position: usize,
234    job_position: usize,
235}
236
237enum Admission {
238    Accepted,
239    Rejected,
240}
241
242/// The pool's concrete event sum, including existing supervision facts.
243pub type PoolEvent<A, J, R> = SupervisionEvent<User<A, PoolMessage<A, J, R>>>;
244
245/// Concrete event sum for a [`KeyedWorkerPool`].
246pub type KeyedPoolEvent<A, K, J, R> = SupervisionEvent<User<A, KeyedPoolMessage<A, K, J, R>>>;
247
248/// Named pool-owned delivery lanes.
249pub struct PoolBehaviorSends<A: Address, J, R, C: Behavior<Addr = A>> {
250    /// Admission and terminal responses addressed to submitters.
251    pub responses: Vec<Delivery<A, PoolResponse<J, R, A>>>,
252    /// Assignments addressed to the selected stable worker proxies.
253    pub assignments: Vec<Delivery<A, ProxyCommand<C>>>,
254}
255
256impl<A: Address, J, R, C: Behavior<Addr = A>> SendAlgebra for PoolBehaviorSends<A, J, R, C> {
257    fn empty() -> Self {
258        Self {
259            responses: Vec::new(),
260            assignments: Vec::new(),
261        }
262    }
263
264    fn append(&mut self, mut other: Self) {
265        self.responses.append(&mut other.responses);
266        self.assignments.append(&mut other.assignments);
267    }
268}
269
270impl<A: Address, J, R, C: Behavior<Addr = A>> SendInput<Delivery<A, PoolResponse<J, R, A>>, Own>
271    for PoolBehaviorSends<A, J, R, C>
272{
273    fn emit(&mut self, input: Delivery<A, PoolResponse<J, R, A>>) {
274        self.responses.push(input);
275    }
276}
277
278impl<A: Address, J, R, C: Behavior<Addr = A>> SendInput<Delivery<A, ProxyCommand<C>>, Own>
279    for PoolBehaviorSends<A, J, R, C>
280{
281    fn emit(&mut self, input: Delivery<A, ProxyCommand<C>>) {
282        self.assignments.push(input);
283    }
284}
285
286type KernelSends<A, J, R, C> = PoolBehaviorSends<A, J, R, C>;
287
288/// Pool effects keep responses and assignments in named, independently
289/// appendable lanes within the supervised behavior send product.
290pub type PoolSends<A, J, R, C> = SupervisorSends<A, KernelSends<A, J, R, C>, C>;
291
292/// Complete action type returned by a [`WorkerPool`] transition.
293pub type PoolActions<A, J, R, C> = Actions<A, Never, PoolSends<A, J, R, C>, Births<Proxy<C>>>;
294
295struct PoolKernel<A: Address, J, R, C>(PhantomData<fn(A, J, R, C)>);
296
297impl<A: Address, J, R, C> PoolKernel<A, J, R, C> {
298    const fn new() -> Self {
299        Self(PhantomData)
300    }
301}
302
303impl<A, J, R, C> Behavior for PoolKernel<A, J, R, C>
304where
305    A: Address,
306    C: Behavior<Addr = A, Msg = PoolAssignment<J>, Ph = Never>,
307{
308    type Addr = A;
309    type Msg = PoolMessage<A, J, R>;
310    type Event = User<A, PoolMessage<A, J, R>>;
311    type Sends = KernelSends<A, J, R, C>;
312    type Ph = Never;
313    type Error = Infallible;
314    type Birth = Births<C>;
315
316    fn init(&mut self) -> crate::BehaviorActed<Self> {
317        Ok(Actions::cont())
318    }
319
320    fn transition(&mut self, _event: Self::Event) -> crate::BehaviorActed<Self> {
321        Ok(Actions::cont())
322    }
323}
324
325type PoolSupervisor<A, J, R, C> = Supervisor<PoolKernel<A, J, R, C>, C>;
326
327/// A fixed, homogeneous, bounded FIFO worker pool.
328///
329/// Each configured nonce names one stable supervised proxy. Jobs are assigned
330/// only after a successful worker-creation result makes that slot idle. The
331/// retained state records an assignment before the corresponding delivery is
332/// returned, and a completion must carry the exact assignment token.
333///
334/// # Panics
335///
336/// Admission or retry propagates a panic from the application payload's
337/// `Clone` implementation before changing pool state. Dispatch panics at the
338/// physical assignment-counter boundary before committing its dispatch plan;
339/// the executor's poison-before-step contract makes that actor turn terminal
340/// rather than exposing partial successor state. The final counter value is
341/// deliberately reserved so every successful batch has a representable
342/// successor counter. This is a Bombay implementation boundary, not an actor
343/// model law.
344///
345/// A worker with any other message protocol cannot form a pool:
346///
347/// ```compile_fail
348/// use behavior::{Behavior, MailAddr, Never, NoBirths, PoolAssignment, User, WorkerPool};
349///
350/// struct WrongWorker;
351/// impl Behavior for WrongWorker {
352///     type Addr = MailAddr;
353///     type Msg = u8;
354///     type Event = User<MailAddr, u8>;
355///     type Sends = Vec<behavior::Delivery<MailAddr, Never>>;
356///     type Ph = Never;
357///     type Error = Never;
358///     type Birth = NoBirths;
359///     fn init(&mut self) -> behavior::BehaviorActed<Self> { unimplemented!() }
360///     fn transition(&mut self, _: Self::Event) -> behavior::BehaviorActed<Self> { unimplemented!() }
361/// }
362///
363/// let _: Option<WorkerPool<MailAddr, String, (), WrongWorker>> = None;
364/// ```
365pub struct WorkerPool<A: Address, J, R, C>
366where
367    C: Behavior<Addr = A, Msg = PoolAssignment<J>, Ph = Never>,
368{
369    supervisor: PoolSupervisor<A, J, R, C>,
370    slots: Vec<Slot<A, J, R>>,
371    backlog: VecDeque<QueuedJob<A, J, R>>,
372    backlog_capacity: usize,
373    next_assignment: u64,
374    interruption: InterruptionPolicy,
375}
376
377impl<A, J, R, C> WorkerPool<A, J, R, C>
378where
379    A: Address,
380    C: Behavior<Addr = A, Msg = PoolAssignment<J>, Ph = Never>,
381{
382    /// Construct a pool after proving that every configured child route is
383    /// unique.
384    ///
385    /// # Errors
386    ///
387    /// Returns [`PoolConfigError::NoWorkers`] for an empty topology or
388    /// [`PoolConfigError::DuplicateWorker`] for the first repeated
389    /// creator-local nonce. No behavior or creation request is produced.
390    #[allow(
391        clippy::too_many_arguments,
392        reason = "the arguments expose the complete pool policy"
393    )]
394    pub fn new(
395        nonces: fn(usize) -> A::Nonce,
396        count: usize,
397        build: fn(usize) -> C,
398        backlog_capacity: usize,
399        interruption: InterruptionPolicy,
400        restart_policy: RestartPolicy,
401        max_restarts: u32,
402        restart_window: Duration,
403    ) -> Result<Self, PoolConfigError<A::Nonce>> {
404        if count == 0 {
405            return Err(PoolConfigError::NoWorkers);
406        }
407        let mut slots = Vec::with_capacity(count);
408        for index in 0..count {
409            let nonce = nonces(index);
410            if slots.iter().any(|slot: &Slot<A, J, R>| slot.nonce == nonce) {
411                return Err(PoolConfigError::DuplicateWorker(nonce));
412            }
413            slots.push(Slot {
414                nonce,
415                state: SlotState::Installing,
416            });
417        }
418        Ok(Self {
419            supervisor: Supervisor::new(
420                PoolKernel::new(),
421                nonces,
422                count,
423                build,
424                Strategy::OneForOne,
425                restart_policy,
426                max_restarts,
427                restart_window,
428            ),
429            slots,
430            backlog: VecDeque::new(),
431            backlog_capacity,
432            next_assignment: 0,
433            interruption,
434        })
435    }
436
437    #[must_use]
438    pub fn backlog_len(&self) -> usize {
439        self.backlog.len()
440    }
441
442    #[must_use]
443    pub fn worker_phase(&self, worker: A::Nonce) -> Option<WorkerPhase> {
444        self.slots
445            .iter()
446            .find(|slot| slot.nonce == worker)
447            .map(|slot| match &slot.state {
448                SlotState::Installing => WorkerPhase::Installing,
449                SlotState::Idle => WorkerPhase::Idle,
450                SlotState::Assigned { assignment, job } => WorkerPhase::Assigned {
451                    assignment: *assignment,
452                    job: job.id,
453                },
454                SlotState::Retired { reason } => WorkerPhase::Retired { reason: *reason },
455            })
456    }
457
458    fn slot_position(&self, worker: A::Nonce) -> Result<usize, PoolError<A::Nonce>> {
459        self.slots
460            .iter()
461            .position(|slot| slot.nonce == worker)
462            .ok_or(PoolError::UnknownWorker(worker))
463    }
464}
465
466impl<A, J, R, C> WorkerPool<A, J, R, C>
467where
468    A: Address,
469    A::Nonce: From<u64>,
470    J: Clone,
471    C: Behavior<Addr = A, Msg = PoolAssignment<J>, Ph = Never>,
472{
473    fn supervisor_transition(&mut self, event: PoolEvent<A, J, R>) -> PoolActions<A, J, R, C> {
474        match delegate_transition(&mut self.supervisor, event) {
475            Ok(actions) => actions,
476            Err(never) => match never {},
477        }
478    }
479
480    fn submit(
481        &mut self,
482        job: JobId,
483        payload: J,
484        reply_to: Recipient<A, PoolResponse<J, R, A>>,
485        actions: &mut PoolActions<A, J, R, C>,
486    ) {
487        let can_dispatch = self
488            .slots
489            .iter()
490            .any(|slot| matches!(slot.state, SlotState::Idle));
491        if !can_dispatch && self.backlog.len() == self.backlog_capacity {
492            actions.sends.behavior.send::<_, Own>(Delivery::new(
493                reply_to,
494                PoolResponse::Rejected {
495                    job,
496                    payload,
497                    reason: PoolRejection::BacklogFull,
498                },
499            ));
500            return;
501        }
502        let dispatch_payload = payload.clone();
503        self.backlog.push_back(QueuedJob {
504            accepted: AcceptedJob {
505                id: job,
506                payload,
507                reply_to,
508                interruption: None,
509                target: None,
510            },
511            dispatch_payload,
512        });
513        actions
514            .sends
515            .behavior
516            .send::<_, Own>(Delivery::new(reply_to, PoolResponse::Accepted { job }));
517    }
518
519    fn submit_to(
520        &mut self,
521        target: A::Nonce,
522        job: JobId,
523        payload: J,
524        reply_to: Recipient<A, PoolResponse<J, R, A>>,
525        actions: &mut PoolActions<A, J, R, C>,
526    ) -> Admission {
527        let Some(slot) = self.slots.iter().find(|slot| slot.nonce == target) else {
528            actions.sends.behavior.send::<_, Own>(Delivery::new(
529                reply_to,
530                PoolResponse::Rejected {
531                    job,
532                    payload,
533                    reason: PoolRejection::AffinityUnavailable,
534                },
535            ));
536            return Admission::Rejected;
537        };
538        if matches!(slot.state, SlotState::Retired { .. }) {
539            actions.sends.behavior.send::<_, Own>(Delivery::new(
540                reply_to,
541                PoolResponse::Rejected {
542                    job,
543                    payload,
544                    reason: PoolRejection::AffinityUnavailable,
545                },
546            ));
547            return Admission::Rejected;
548        }
549        let can_dispatch = matches!(slot.state, SlotState::Idle);
550        if !can_dispatch && self.backlog.len() == self.backlog_capacity {
551            actions.sends.behavior.send::<_, Own>(Delivery::new(
552                reply_to,
553                PoolResponse::Rejected {
554                    job,
555                    payload,
556                    reason: PoolRejection::BacklogFull,
557                },
558            ));
559            return Admission::Rejected;
560        }
561        let dispatch_payload = payload.clone();
562        self.backlog.push_back(QueuedJob {
563            accepted: AcceptedJob {
564                id: job,
565                payload,
566                reply_to,
567                interruption: None,
568                target: Some(target),
569            },
570            dispatch_payload,
571        });
572        actions
573            .sends
574            .behavior
575            .send::<_, Own>(Delivery::new(reply_to, PoolResponse::Accepted { job }));
576        Admission::Accepted
577    }
578
579    fn complete(
580        &mut self,
581        worker: A::Nonce,
582        assignment: AssignmentId,
583        result: R,
584        actions: &mut PoolActions<A, J, R, C>,
585    ) -> Result<(), PoolError<A::Nonce>> {
586        let position = self.slot_position(worker)?;
587        let phase = self
588            .worker_phase(worker)
589            .expect("position proves the slot exists");
590        let SlotState::Assigned {
591            assignment: expected,
592            ..
593        } = &self.slots[position].state
594        else {
595            return Err(PoolError::CompletionForUnavailableWorker { worker, phase });
596        };
597        if *expected != assignment {
598            return Err(PoolError::StaleCompletion {
599                worker,
600                expected: *expected,
601                received: assignment,
602            });
603        }
604        let SlotState::Assigned { job, .. } =
605            core::mem::replace(&mut self.slots[position].state, SlotState::Idle)
606        else {
607            unreachable!("the state was proven assigned")
608        };
609        actions.sends.behavior.send::<_, Own>(Delivery::new(
610            job.reply_to,
611            PoolResponse::Completed {
612                job: job.id,
613                result,
614            },
615        ));
616        Ok(())
617    }
618
619    fn worker_stopped(
620        &mut self,
621        stopped: &WorkerStopped<A>,
622        responses: &mut Vec<Delivery<A, PoolResponse<J, R, A>>>,
623    ) -> Result<(), PoolError<A::Nonce>> {
624        let position = self.slot_position(stopped.proxy)?;
625        let phase = self
626            .worker_phase(stopped.proxy)
627            .expect("position proves the slot exists");
628        if matches!(phase, WorkerPhase::Installing | WorkerPhase::Retired { .. }) {
629            return Err(PoolError::WorkerStoppedWhileUnavailable {
630                worker: stopped.proxy,
631                phase,
632            });
633        }
634        if self.interruption == InterruptionPolicy::Retry {
635            if let SlotState::Assigned { job, .. } = &self.slots[position].state {
636                let dispatch_payload = job.payload.clone();
637                let SlotState::Assigned { mut job, .. } =
638                    core::mem::replace(&mut self.slots[position].state, SlotState::Installing)
639                else {
640                    unreachable!("the assigned state was matched before committing retry")
641                };
642                job.interruption = Some(PoolInterruption::WorkerStopped {
643                    worker: stopped.proxy,
644                    outcome: stopped.outcome,
645                });
646                self.backlog.push_front(QueuedJob {
647                    accepted: job,
648                    dispatch_payload,
649                });
650            } else {
651                self.slots[position].state = SlotState::Installing;
652            }
653            return Ok(());
654        }
655
656        let previous = core::mem::replace(&mut self.slots[position].state, SlotState::Installing);
657        if let SlotState::Assigned { job, .. } = previous {
658            responses.push(Delivery::new(
659                job.reply_to,
660                PoolResponse::Interrupted {
661                    job: job.id,
662                    payload: job.payload,
663                    reason: PoolInterruption::WorkerStopped {
664                        worker: stopped.proxy,
665                        outcome: stopped.outcome,
666                    },
667                },
668            ));
669        }
670        Ok(())
671    }
672
673    fn fail_backlog_if_irrecoverable(&mut self, actions: &mut PoolActions<A, J, R, C>) {
674        if self
675            .slots
676            .iter()
677            .any(|slot| !matches!(slot.state, SlotState::Retired { .. }))
678        {
679            return;
680        }
681        for queued in self.backlog.drain(..) {
682            let job = queued.accepted;
683            actions.sends.behavior.send::<_, Own>(Delivery::new(
684                job.reply_to,
685                PoolResponse::Interrupted {
686                    job: job.id,
687                    payload: job.payload,
688                    reason: job
689                        .interruption
690                        .unwrap_or(PoolInterruption::NoRecoverableWorkers),
691                },
692            ));
693        }
694    }
695
696    fn fail_jobs_for_retired_slot(
697        &mut self,
698        worker: A::Nonce,
699        reason: WorkerRetirement,
700        actions: &mut PoolActions<A, J, R, C>,
701    ) {
702        let mut retained = VecDeque::with_capacity(self.backlog.len());
703        while let Some(queued) = self.backlog.pop_front() {
704            if queued.accepted.target == Some(worker) {
705                let job = queued.accepted;
706                actions.sends.behavior.send::<_, Own>(Delivery::new(
707                    job.reply_to,
708                    PoolResponse::Interrupted {
709                        job: job.id,
710                        payload: job.payload,
711                        reason: job
712                            .interruption
713                            .unwrap_or(PoolInterruption::AffinityRetired { worker, reason }),
714                    },
715                ));
716            } else {
717                retained.push_back(queued);
718            }
719        }
720        self.backlog = retained;
721    }
722
723    fn creation_resolved(
724        &mut self,
725        resolved: &WorkerCreationResolved<A::Nonce>,
726    ) -> Result<(), PoolError<A::Nonce>> {
727        let position = self.slot_position(resolved.proxy)?;
728        let phase = self
729            .worker_phase(resolved.proxy)
730            .expect("position proves the slot exists");
731        if !matches!(phase, WorkerPhase::Installing) {
732            return Err(PoolError::CreationResolvedWhileUnavailable {
733                worker: resolved.proxy,
734                phase,
735            });
736        }
737        self.slots[position].state = match resolved.result {
738            Ok(()) => SlotState::Idle,
739            Err(rejection) => SlotState::Retired {
740                reason: WorkerRetirement::CreationRejected(rejection),
741            },
742        };
743        Ok(())
744    }
745
746    fn dispatch(&mut self, actions: &mut PoolActions<A, J, R, C>) {
747        let mut selected_jobs = Vec::new();
748        let mut plan = Vec::new();
749        for (slot_position, slot) in self.slots.iter().enumerate() {
750            if !matches!(slot.state, SlotState::Idle) {
751                continue;
752            }
753            let Some(job_position) = self.backlog.iter().enumerate().find_map(|(position, job)| {
754                (!selected_jobs.contains(&position)
755                    && job
756                        .accepted
757                        .target
758                        .is_none_or(|target| target == slot.nonce))
759                .then_some(position)
760            }) else {
761                continue;
762            };
763            selected_jobs.push(job_position);
764            plan.push(PlannedDispatch {
765                slot_position,
766                job_position,
767            });
768        }
769        let count =
770            u64::try_from(plan.len()).expect("a pool cannot contain more than u64::MAX slots");
771        let next_assignment = self
772            .next_assignment
773            .checked_add(count)
774            .expect("pool assignment identifiers exhausted");
775
776        let mut selected_by_position = BTreeMap::new();
777        for planned in plan {
778            selected_by_position.insert(planned.job_position, planned.slot_position);
779        }
780        let mut selected_by_slot: Vec<Option<QueuedJob<A, J, R>>> = std::iter::repeat_with(|| None)
781            .take(self.slots.len())
782            .collect();
783        let mut remaining = VecDeque::new();
784        for (position, queued) in self.backlog.drain(..).enumerate() {
785            if let Some(slot_position) = selected_by_position.remove(&position) {
786                selected_by_slot[slot_position] = Some(queued);
787            } else {
788                remaining.push_back(queued);
789            }
790        }
791        self.backlog = remaining;
792
793        for (offset, (slot_position, queued)) in selected_by_slot
794            .into_iter()
795            .enumerate()
796            .filter_map(|(slot_position, queued)| queued.map(|queued| (slot_position, queued)))
797            .enumerate()
798        {
799            let payload = queued.dispatch_payload;
800            let job = queued.accepted;
801            let assignment = AssignmentId(
802                self.next_assignment
803                    + u64::try_from(offset).expect("offset is bounded by the checked plan length"),
804            );
805            let nonce = self.slots[slot_position].nonce;
806            let job_id = job.id;
807            self.slots[slot_position].state = SlotState::Assigned { assignment, job };
808            actions.sends.behavior.send::<_, Own>(Delivery::new(
809                Recipient::child(nonce),
810                ProxyCommand::Forward(PoolAssignment {
811                    assignment,
812                    job: job_id,
813                    payload,
814                }),
815            ));
816        }
817        self.next_assignment = next_assignment;
818    }
819}
820
821impl<A, J, R, C> Behavior for WorkerPool<A, J, R, C>
822where
823    A: Address,
824    A::Nonce: From<u64>,
825    J: Clone,
826    C: Behavior<Addr = A, Msg = PoolAssignment<J>, Ph = Never>,
827{
828    type Addr = A;
829    type Msg = PoolMessage<A, J, R>;
830    type Event = PoolEvent<A, J, R>;
831    type Sends = PoolSends<A, J, R, C>;
832    type Ph = Never;
833    type Error = PoolError<A::Nonce>;
834    type Birth = Births<Proxy<C>>;
835
836    fn init(&mut self) -> crate::BehaviorActed<Self> {
837        match self.supervisor.init() {
838            Ok(actions) => Ok(actions),
839            Err(never) => match never {},
840        }
841    }
842
843    fn transition(&mut self, event: Self::Event) -> crate::BehaviorActed<Self> {
844        match event {
845            SupervisionEvent::Inner(User {
846                message:
847                    PoolMessage::Submit {
848                        job,
849                        payload,
850                        reply_to,
851                    },
852                ..
853            }) => {
854                let mut actions = Actions::cont();
855                self.submit(job, payload, reply_to, &mut actions);
856                self.dispatch(&mut actions);
857                Ok(actions)
858            }
859            SupervisionEvent::Inner(User {
860                message:
861                    PoolMessage::Completed {
862                        worker,
863                        assignment,
864                        result,
865                    },
866                ..
867            }) => {
868                let mut actions = Actions::cont();
869                self.complete(worker, assignment, result, &mut actions)?;
870                self.dispatch(&mut actions);
871                Ok(actions)
872            }
873            SupervisionEvent::WorkerStopped(stopped) => {
874                let proxy = stopped.proxy;
875                let mut responses = Vec::new();
876                self.worker_stopped(&stopped, &mut responses)?;
877                let mut actions =
878                    self.supervisor_transition(SupervisionEvent::WorkerStopped(stopped));
879                actions.sends.behavior.responses.extend(responses);
880                let replacement_requested = actions
881                    .sends
882                    .replacement_commands
883                    .iter()
884                    .any(|delivery| delivery.to.route() == crate::Route::Child(proxy));
885                if !replacement_requested {
886                    let position = self.slot_position(proxy)?;
887                    let reason = WorkerRetirement::ReplacementUnavailable;
888                    self.slots[position].state = SlotState::Retired { reason };
889                    self.fail_jobs_for_retired_slot(proxy, reason, &mut actions);
890                }
891                self.dispatch(&mut actions);
892                self.fail_backlog_if_irrecoverable(&mut actions);
893                Ok(actions)
894            }
895            SupervisionEvent::WorkerCreationResolved(resolved) => {
896                let proxy = resolved.proxy;
897                self.creation_resolved(&resolved)?;
898                let mut actions =
899                    self.supervisor_transition(SupervisionEvent::WorkerCreationResolved(resolved));
900                if let Some(WorkerPhase::Retired { reason }) = self.worker_phase(proxy) {
901                    self.fail_jobs_for_retired_slot(proxy, reason, &mut actions);
902                }
903                self.dispatch(&mut actions);
904                self.fail_backlog_if_irrecoverable(&mut actions);
905                Ok(actions)
906            }
907            SupervisionEvent::ChildStopped(stopped) => {
908                Ok(self.supervisor_transition(SupervisionEvent::ChildStopped(stopped)))
909            }
910            SupervisionEvent::CreationResolved(resolved) => {
911                Ok(self.supervisor_transition(SupervisionEvent::CreationResolved(resolved)))
912            }
913        }
914    }
915}
916
917#[cfg(test)]
918mod tests {
919    use super::*;
920    use crate::{MailAddr, NoBirths, Route};
921
922    #[derive(Clone, Copy)]
923    struct TestWorker;
924
925    impl Behavior for TestWorker {
926        type Addr = MailAddr;
927        type Msg = PoolAssignment<u8>;
928        type Event = User<MailAddr, PoolAssignment<u8>>;
929        type Sends = Vec<Delivery<MailAddr, Never>>;
930        type Ph = Never;
931        type Error = Never;
932        type Birth = NoBirths;
933
934        fn init(&mut self) -> crate::BehaviorActed<Self> {
935            Ok(Actions::cont())
936        }
937
938        fn transition(&mut self, _: Self::Event) -> crate::BehaviorActed<Self> {
939            Ok(Actions::cont())
940        }
941    }
942
943    fn test_worker(_: usize) -> TestWorker {
944        TestWorker
945    }
946
947    #[test]
948    fn one_dispatch_batch_preserves_fifo_jobs_across_index_removal() {
949        let mut pool = WorkerPool::new(
950            |index| u64::try_from(index).unwrap(),
951            2,
952            test_worker,
953            3,
954            InterruptionPolicy::Fail,
955            RestartPolicy::Permanent,
956            1,
957            Duration::from_secs(1),
958        )
959        .unwrap();
960        pool.init().unwrap();
961
962        for job in 1..=3 {
963            pool.transition(SupervisionEvent::Inner(User::new(
964                MailAddr(90),
965                PoolMessage::Submit {
966                    job: JobId(job),
967                    payload: u8::try_from(job).unwrap(),
968                    reply_to: Recipient::global(MailAddr(91)),
969                },
970            )))
971            .unwrap();
972        }
973        pool.slots[0].state = SlotState::Idle;
974        pool.slots[1].state = SlotState::Idle;
975
976        let mut actions: PoolActions<MailAddr, u8, (), TestWorker> = Actions::cont();
977        pool.dispatch(&mut actions);
978
979        let assignments = &actions.sends.behavior.assignments;
980        assert_eq!(assignments.len(), 2);
981        for (index, expected_job) in [JobId(1), JobId(2)].into_iter().enumerate() {
982            assert_eq!(
983                assignments[index].to.route(),
984                Route::Child(u64::try_from(index).unwrap())
985            );
986            let ProxyCommand::Forward(assignment) = &assignments[index].message else {
987                panic!("pool dispatches with Forward");
988            };
989            assert_eq!(
990                assignment.assignment,
991                AssignmentId(u64::try_from(index).unwrap())
992            );
993            assert_eq!(assignment.job, expected_job);
994        }
995        assert_eq!(pool.backlog.len(), 1);
996        assert_eq!(pool.backlog[0].accepted.id, JobId(3));
997    }
998}
999
1000/// A worker pool whose admitted keys remain bound to stable worker slots.
1001///
1002/// The selector chooses a stable proxy nonce only when a key is first
1003/// admitted. Replacement incarnations remain behind that proxy, so they do
1004/// not alter affinity. [`KeyedPoolMessage::Rebalance`] is the sole transition
1005/// that changes an established binding, and jobs accepted before it retain
1006/// their original target.
1007///
1008/// Keys must have a concrete equality relation; a key type without `Eq` cannot
1009/// form an affinity table:
1010///
1011/// ```compile_fail
1012/// use behavior::{Actions, Delivery, KeyedWorkerPool, MailAddr, Never, NoBirths};
1013/// struct NonKey(f64);
1014/// struct Worker;
1015/// #[behavior::behavior(
1016///     addr = MailAddr,
1017///     message = behavior::PoolAssignment<u8>,
1018///     sends = Vec<Delivery<MailAddr, Never>>,
1019///     births = NoBirths,
1020///     error = Never,
1021/// )]
1022/// impl Worker {
1023///     fn init(&mut self) -> behavior::Acted<MailAddr, Never, Vec<Delivery<MailAddr, Never>>, NoBirths, Never> {
1024///         Ok(Actions::cont())
1025///     }
1026///     fn receive(&mut self, _: MailAddr, _: behavior::PoolAssignment<u8>) -> behavior::Acted<MailAddr, Never, Vec<Delivery<MailAddr, Never>>, NoBirths, Never> {
1027///         Ok(Actions::cont())
1028///     }
1029/// }
1030/// let _: Option<KeyedWorkerPool<MailAddr, NonKey, u8, (), Worker, fn(&NonKey) -> u64>> = None;
1031/// ```
1032pub struct KeyedWorkerPool<A: Address, K, J, R, C, S>
1033where
1034    K: Eq,
1035    C: Behavior<Addr = A, Msg = PoolAssignment<J>, Ph = Never>,
1036    S: AffinitySelector<K, A::Nonce>,
1037{
1038    pool: WorkerPool<A, J, R, C>,
1039    bindings: Vec<(K, A::Nonce)>,
1040    selector: S,
1041}
1042
1043impl<A, K, J, R, C, S> KeyedWorkerPool<A, K, J, R, C, S>
1044where
1045    A: Address,
1046    K: Eq,
1047    C: Behavior<Addr = A, Msg = PoolAssignment<J>, Ph = Never>,
1048    S: AffinitySelector<K, A::Nonce>,
1049{
1050    /// Construct a key-persistent pool over the same fixed supervised slots as
1051    /// [`WorkerPool`]. The selector is pure and is consulted once per
1052    /// previously unseen key. It chooses behavior policy; runtime route
1053    /// resolution remains outside this type.
1054    ///
1055    /// # Errors
1056    ///
1057    /// Returns [`PoolConfigError::NoWorkers`] for an empty topology or
1058    /// [`PoolConfigError::DuplicateWorker`] for a repeated stable nonce.
1059    #[allow(
1060        clippy::too_many_arguments,
1061        reason = "the arguments expose the complete pool and affinity policy"
1062    )]
1063    pub fn new(
1064        nonces: fn(usize) -> A::Nonce,
1065        count: usize,
1066        build: fn(usize) -> C,
1067        backlog_capacity: usize,
1068        interruption: InterruptionPolicy,
1069        restart_policy: RestartPolicy,
1070        max_restarts: u32,
1071        restart_window: Duration,
1072        selector: S,
1073    ) -> Result<Self, PoolConfigError<A::Nonce>> {
1074        Ok(Self {
1075            pool: WorkerPool::new(
1076                nonces,
1077                count,
1078                build,
1079                backlog_capacity,
1080                interruption,
1081                restart_policy,
1082                max_restarts,
1083                restart_window,
1084            )?,
1085            bindings: Vec::new(),
1086            selector,
1087        })
1088    }
1089
1090    /// Return the stable slot currently bound to `key`.
1091    #[must_use]
1092    pub fn affinity(&self, key: &K) -> Option<A::Nonce> {
1093        self.bindings
1094            .iter()
1095            .find_map(|(bound, worker)| (bound == key).then_some(*worker))
1096    }
1097
1098    #[must_use]
1099    pub fn backlog_len(&self) -> usize {
1100        self.pool.backlog_len()
1101    }
1102
1103    #[must_use]
1104    pub fn worker_phase(&self, worker: A::Nonce) -> Option<WorkerPhase> {
1105        self.pool.worker_phase(worker)
1106    }
1107
1108    fn rebalance(&mut self, key: K, worker: A::Nonce) -> Result<(), PoolError<A::Nonce>> {
1109        let position = self.pool.slot_position(worker)?;
1110        if let SlotState::Retired { reason } = self.pool.slots[position].state {
1111            return Err(PoolError::RebalanceToRetiredWorker { worker, reason });
1112        }
1113        if let Some((_, bound)) = self.bindings.iter_mut().find(|(bound, _)| *bound == key) {
1114            *bound = worker;
1115        } else {
1116            self.bindings.push((key, worker));
1117        }
1118        Ok(())
1119    }
1120}
1121
1122impl<A, K, J, R, C, S> Behavior for KeyedWorkerPool<A, K, J, R, C, S>
1123where
1124    A: Address,
1125    A::Nonce: From<u64>,
1126    K: Eq,
1127    J: Clone,
1128    C: Behavior<Addr = A, Msg = PoolAssignment<J>, Ph = Never>,
1129    S: AffinitySelector<K, A::Nonce>,
1130{
1131    type Addr = A;
1132    type Msg = KeyedPoolMessage<A, K, J, R>;
1133    type Event = KeyedPoolEvent<A, K, J, R>;
1134    type Sends = PoolSends<A, J, R, C>;
1135    type Ph = Never;
1136    type Error = PoolError<A::Nonce>;
1137    type Birth = Births<Proxy<C>>;
1138
1139    fn init(&mut self) -> crate::BehaviorActed<Self> {
1140        self.pool.init()
1141    }
1142
1143    fn transition(&mut self, event: Self::Event) -> crate::BehaviorActed<Self> {
1144        match event {
1145            SupervisionEvent::Inner(User {
1146                message:
1147                    KeyedPoolMessage::Submit {
1148                        key,
1149                        job,
1150                        payload,
1151                        reply_to,
1152                    },
1153                ..
1154            }) => {
1155                let existing = self.affinity(&key);
1156                let target = existing.unwrap_or_else(|| self.selector.select(&key));
1157                let mut actions = Actions::cont();
1158                let admission = self
1159                    .pool
1160                    .submit_to(target, job, payload, reply_to, &mut actions);
1161                match (admission, existing) {
1162                    (Admission::Accepted, None) => self.bindings.push((key, target)),
1163                    (Admission::Accepted | Admission::Rejected, Some(_))
1164                    | (Admission::Rejected, None) => {}
1165                }
1166                self.pool.dispatch(&mut actions);
1167                Ok(actions)
1168            }
1169            SupervisionEvent::Inner(User {
1170                message:
1171                    KeyedPoolMessage::Completed {
1172                        worker,
1173                        assignment,
1174                        result,
1175                    },
1176                ..
1177            }) => {
1178                let mut actions = Actions::cont();
1179                self.pool
1180                    .complete(worker, assignment, result, &mut actions)?;
1181                self.pool.dispatch(&mut actions);
1182                Ok(actions)
1183            }
1184            SupervisionEvent::Inner(User {
1185                message: KeyedPoolMessage::Rebalance { key, worker },
1186                ..
1187            }) => {
1188                self.rebalance(key, worker)?;
1189                Ok(Actions::cont())
1190            }
1191            SupervisionEvent::WorkerStopped(stopped) => self
1192                .pool
1193                .transition(SupervisionEvent::WorkerStopped(stopped)),
1194            SupervisionEvent::WorkerCreationResolved(resolved) => self
1195                .pool
1196                .transition(SupervisionEvent::WorkerCreationResolved(resolved)),
1197            SupervisionEvent::ChildStopped(stopped) => self
1198                .pool
1199                .transition(SupervisionEvent::ChildStopped(stopped)),
1200            SupervisionEvent::CreationResolved(resolved) => self
1201                .pool
1202                .transition(SupervisionEvent::CreationResolved(resolved)),
1203        }
1204    }
1205}
behavior/pool.rs

behavior/
pool.rs
