Skip to main content

rightkit_process/
restart.rs

1use std::{collections::VecDeque, error::Error, fmt, time::Duration};
2
3#[derive(Clone, Copy, Debug, Eq, PartialEq)]
4pub enum RestartMode {
5    Automatic,
6    OnDemand,
7    Disabled,
8}
9
10#[derive(Clone, Copy, Debug, Eq, PartialEq)]
11pub struct RestartPolicy {
12    mode: RestartMode,
13    max_restarts: u32,
14    window: Duration,
15    base_backoff: Duration,
16    max_backoff: Duration,
17}
18
19#[derive(Clone, Copy, Debug, Eq, PartialEq)]
20pub struct InvalidRestartPolicy;
21
22impl fmt::Display for InvalidRestartPolicy {
23    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24        f.write_str("restart window must be non-zero and max backoff must cover base backoff")
25    }
26}
27
28impl Error for InvalidRestartPolicy {}
29
30impl RestartPolicy {
31    pub fn new(
32        mode: RestartMode,
33        max_restarts: u32,
34        window: Duration,
35        base_backoff: Duration,
36        max_backoff: Duration,
37    ) -> Result<Self, InvalidRestartPolicy> {
38        if window.is_zero() || max_backoff < base_backoff {
39            return Err(InvalidRestartPolicy);
40        }
41        Ok(Self {
42            mode,
43            max_restarts,
44            window,
45            base_backoff,
46            max_backoff,
47        })
48    }
49}
50
51#[derive(Clone, Copy, Debug, Eq, PartialEq)]
52pub struct Generation(u64);
53
54#[derive(Clone, Copy, Debug, Eq, PartialEq)]
55pub struct RestartPermit {
56    epoch: u64,
57    delay: Duration,
58    not_before: Duration,
59}
60
61impl RestartPermit {
62    pub fn delay(&self) -> Duration {
63        self.delay
64    }
65}
66
67#[derive(Clone, Copy, Debug, Eq, PartialEq)]
68pub enum RestartOutcome {
69    Restart(RestartPermit),
70    AwaitingDemand,
71    Exhausted,
72    IgnoredStale,
73    Stopped,
74}
75
76#[derive(Clone, Debug)]
77pub struct RestartTracker {
78    policy: RestartPolicy,
79    active: Option<Generation>,
80    next_generation: u64,
81    restart_times: VecDeque<Duration>,
82    permit_epoch: u64,
83    awaiting_demand: bool,
84    shutdown: bool,
85}
86
87impl RestartTracker {
88    pub fn new(policy: RestartPolicy) -> Self {
89        Self {
90            policy,
91            active: None,
92            next_generation: 0,
93            restart_times: VecDeque::new(),
94            permit_epoch: 0,
95            awaiting_demand: false,
96            shutdown: false,
97        }
98    }
99
100    pub fn started(&mut self, _at: Duration) -> Generation {
101        self.next_generation = self.next_generation.saturating_add(1);
102        self.permit_epoch = self.permit_epoch.saturating_add(1);
103        self.awaiting_demand = false;
104        let generation = Generation(self.next_generation);
105        self.active = Some(generation);
106        generation
107    }
108
109    pub fn exited(&mut self, generation: Generation, at: Duration) -> RestartOutcome {
110        if self.shutdown {
111            return RestartOutcome::Stopped;
112        }
113        if self.active != Some(generation) {
114            return RestartOutcome::IgnoredStale;
115        }
116        self.active = None;
117        match self.policy.mode {
118            RestartMode::Automatic => self.schedule(at),
119            RestartMode::OnDemand => {
120                self.awaiting_demand = true;
121                RestartOutcome::AwaitingDemand
122            }
123            RestartMode::Disabled => RestartOutcome::Stopped,
124        }
125    }
126
127    pub fn request_restart(&mut self, at: Duration) -> RestartOutcome {
128        if self.shutdown || self.policy.mode == RestartMode::Disabled {
129            return RestartOutcome::Stopped;
130        }
131        if self.policy.mode == RestartMode::OnDemand && !self.awaiting_demand {
132            return RestartOutcome::AwaitingDemand;
133        }
134        self.awaiting_demand = false;
135        self.schedule(at)
136    }
137
138    pub fn permit_valid(&self, permit: &RestartPermit, at: Duration) -> bool {
139        !self.shutdown && permit.epoch == self.permit_epoch && at >= permit.not_before
140    }
141
142    pub fn shutdown(&mut self) {
143        self.shutdown = true;
144        self.active = None;
145        self.awaiting_demand = false;
146        self.permit_epoch = self.permit_epoch.saturating_add(1);
147    }
148
149    fn schedule(&mut self, at: Duration) -> RestartOutcome {
150        let window_start = at.saturating_sub(self.policy.window);
151        while self
152            .restart_times
153            .front()
154            .is_some_and(|restart| *restart < window_start)
155        {
156            self.restart_times.pop_front();
157        }
158        if self.restart_times.len() >= self.policy.max_restarts as usize {
159            return RestartOutcome::Exhausted;
160        }
161
162        let exponent = u32::try_from(self.restart_times.len()).unwrap_or(u32::MAX);
163        let multiplier = 1_u32.checked_shl(exponent.min(31)).unwrap_or(u32::MAX);
164        let delay = self
165            .policy
166            .base_backoff
167            .checked_mul(multiplier)
168            .unwrap_or(self.policy.max_backoff)
169            .min(self.policy.max_backoff);
170        self.restart_times.push_back(at);
171        self.permit_epoch = self.permit_epoch.saturating_add(1);
172        let permit = RestartPermit {
173            epoch: self.permit_epoch,
174            delay,
175            not_before: at.checked_add(delay).unwrap_or(Duration::MAX),
176        };
177        RestartOutcome::Restart(permit)
178    }
179}