1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
//! Graceful shutdown and single-node activity drain coordination.
use std::process::ExitCode;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tokio::sync::Notify;
use tracing::{error, info, warn};
use crate::ServerState;
use crate::error::ServerError;
use crate::worker::LostWorkerReport;
/// Process exit selected by the shutdown coordinator.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ShutdownOutcome {
/// Drain completed before the configured timeout.
Clean,
/// In-flight activities outlived the drain timeout and were parked for
/// restart recovery (#207): nothing recorded, nothing delivered — the
/// recoverable-by-design state, so a fully-parked drain is a SUCCESS. A
/// long-running activity (an agent round runs hours) outliving any sane
/// drain window is the expected case, and a non-zero exit on every routine
/// deploy would train operators to ignore failures.
Parked,
/// The drain timed out AND the park itself failed (lock poison, sink
/// error): in-flight state could not be handed to restart recovery.
TimedOut,
/// A second termination signal requested immediate process exit.
Forced,
}
impl ShutdownOutcome {
/// Convert the outcome to the process exit code required by the operations contract.
#[must_use]
pub fn exit_code(self) -> ExitCode {
match self {
Self::Clean | Self::Parked => ExitCode::SUCCESS,
Self::TimedOut => ExitCode::FAILURE,
Self::Forced => ExitCode::from(130),
}
}
}
/// Cloneable gate shared by transports, dispatchers, worker streams, and the
/// shutdown coordinator.
#[derive(Clone, Debug, Default)]
pub struct DrainState {
inner: Arc<DrainStateInner>,
}
#[derive(Debug, Default)]
struct DrainStateInner {
draining: AtomicBool,
empty: Notify,
}
impl DrainState {
/// Return whether drain has begun and new workflow/activity starts must be rejected.
#[must_use]
pub fn is_draining(&self) -> bool {
self.inner.draining.load(Ordering::Acquire)
}
/// Mark the server draining. Returns true for the first caller that changed the state.
#[must_use]
pub fn begin(&self) -> bool {
!self.inner.draining.swap(true, Ordering::AcqRel)
}
/// Reject a new unit of work if drain has already begun.
///
/// # Errors
///
/// Returns [`ServerError::WorkerDispatch`] with a stable drain message when work is closed.
pub fn ensure_accepting(
&self,
namespace: &str,
activity_type: &str,
) -> Result<(), ServerError> {
if self.is_draining() {
Err(ServerError::worker_dispatch(
namespace.to_owned(),
activity_type.to_owned(),
"server is draining and not accepting new activity tasks",
))
} else {
Ok(())
}
}
/// Wake waiters after in-flight accounting may have reached zero.
pub fn notify_activity_drained(&self) {
self.inner.empty.notify_waiters();
}
async fn wait_for_empty(&self, state: &ServerState) -> Result<(), ServerError> {
loop {
let in_flight = state.heartbeat_tracker().in_flight_count()?;
if in_flight == 0 {
return Ok(());
}
let notified = self.inner.empty.notified();
if state.heartbeat_tracker().in_flight_count()? == 0 {
return Ok(());
}
notified.await;
}
}
}
/// Run the graceful drain after the first termination signal.
///
/// The caller is responsible for stopping transports as soon as drain begins.
///
/// # Errors
///
/// Returns [`ServerError`] if worker-drain broadcast, in-flight accounting, timeout failure
/// surfacing, or engine shutdown fails.
pub async fn drain_after_first_signal(
state: ServerState,
second_signal: impl std::future::Future<Output = ()>,
) -> Result<ShutdownOutcome, ServerError> {
let drain = state.drain_state().clone();
let first = drain.begin();
if first {
info!("shutdown signal received; beginning graceful drain");
}
let delivered_workers = state.worker_registry().broadcast_drain()?;
info!(delivered_workers, "sent drain request to connected workers");
let timeout = state.runtime_config().drain_timeout;
tokio::pin!(second_signal);
let outcome = tokio::select! {
() = &mut second_signal => {
warn!("second shutdown signal received; forcing immediate exit");
ShutdownOutcome::Forced
}
result = wait_for_drain_or_timeout(&state, &drain, timeout) => result?,
};
if matches!(outcome, ShutdownOutcome::Forced) {
return Ok(outcome);
}
state.shutdown()?;
Ok(outcome)
}
async fn wait_for_drain_or_timeout(
state: &ServerState,
drain: &DrainState,
timeout: Duration,
) -> Result<ShutdownOutcome, ServerError> {
match tokio::time::timeout(timeout, drain.wait_for_empty(state)).await {
Ok(result) => {
result?;
info!("activity drain completed cleanly");
Ok(ShutdownOutcome::Clean)
}
Err(_elapsed) => {
// #207 drain-timeout backstop: PARK the remaining in-flight
// dispatches for restart recovery instead of synthesizing
// retryable lost-worker failures. Nothing is recorded, so the
// durable log converges on the kill -9 shape and post-restart
// replay re-dispatches every parked ordinal. A park that itself
// fails leaves in-flight state unhanded — the one remaining
// FAILURE-worthy drain outcome.
match state
.heartbeat_tracker()
.park_all_in_flight_workers(state.worker_registry(), state.pending_activities())
{
Ok(reports) => {
log_parked_workers(&reports);
Ok(ShutdownOutcome::Parked)
}
Err(park_error) => {
error!(
%park_error,
"activity drain timed out and parking the remaining in-flight \
activities failed; exiting with the failure drain outcome"
);
Ok(ShutdownOutcome::TimedOut)
}
}
}
}
}
fn log_parked_workers(reports: &[LostWorkerReport]) {
let parked_tasks: usize = reports.iter().map(|report| report.tasks.len()).sum();
if parked_tasks == 0 {
info!("activity drain timed out with no tracked in-flight activities to park");
} else {
info!(
parked_workers = reports.len(),
parked_tasks,
"activity drain timed out; remaining activities parked for restart recovery"
);
}
}
#[cfg(test)]
mod tests {
use std::process::ExitCode;
use super::{DrainState, ShutdownOutcome};
#[test]
fn begin_is_idempotent_and_sets_draining() {
let drain = DrainState::default();
assert!(!drain.is_draining());
assert!(drain.begin());
assert!(drain.is_draining());
assert!(!drain.begin());
}
/// #207 exit contract: a fully-parked drain is a SUCCESS (parked state is
/// recoverable by design); FAILURE is reserved for a park that itself
/// failed; a forced exit keeps 130. `ExitCode` carries no `PartialEq`, so
/// the mapping is asserted through its debug representation.
#[test]
fn exit_codes_map_parked_to_success_and_timed_out_to_failure() {
let debug = |code: ExitCode| format!("{code:?}");
assert_eq!(
debug(ShutdownOutcome::Clean.exit_code()),
debug(ExitCode::SUCCESS)
);
assert_eq!(
debug(ShutdownOutcome::Parked.exit_code()),
debug(ExitCode::SUCCESS)
);
assert_eq!(
debug(ShutdownOutcome::TimedOut.exit_code()),
debug(ExitCode::FAILURE)
);
assert_eq!(
debug(ShutdownOutcome::Forced.exit_code()),
debug(ExitCode::from(130))
);
}
}