Line data Source code
1 : // SPDX-License-Identifier: EUPL-1.2
2 : //! Graceful shutdown drain.
3 : //!
4 : //! The ORDER is the whole feature, and it exists because of one asymmetry: a
5 : //! load balancer learns this instance is going away only by polling
6 : //! `/q/health`, but the orchestrator kills it on its own schedule. So the
7 : //! health endpoint must go unhealthy while the socket still works, and only
8 : //! then may the socket close.
9 : //!
10 : //! 1. flip `draining` → `/q/health` answers 503 (see `antares_api::health`)
11 : //! 2. keep accepting for `ANTARES_DRAIN_DELAY_MS` — the LB's notice window;
12 : //! this is the step people skip, and skipping it is what turns a rolling
13 : //! update into a burst of connection-refused
14 : //! 3. stop accepting
15 : //! 4. wait for in-flight connections, bounded by `ANTARES_DRAIN_DEADLINE_SECS`
16 : //! 5. wait for the outbox to empty — same deadline, whatever step 4 left of it
17 : //! (see the note in `drain`)
18 : //! 6. close the pools
19 : //!
20 : //! The two numbers are for different jobs and are easy to confuse: the delay
21 : //! is the LB's notice window in MILLIseconds (default 2000), the deadline is
22 : //! the ceiling on in-flight work in SECONDS (default 20). Steps 4 and 5 share
23 : //! that one deadline.
24 : //!
25 : //! Operational contract: the container `stopGracePeriod` (compose
26 : //! `stop_grace_period`, K8s `terminationGracePeriodSeconds`) MUST exceed
27 : //! delay + deadline, or the orchestrator turns a drain into a kill. The
28 : //! defaults below (2 s + 20 s) overrun Docker's 10 s default as soon as the
29 : //! in-flight work is not short; the reference manifests set both explicitly.
30 :
31 : use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
32 : use std::sync::Arc;
33 : use std::time::{Duration, Instant};
34 :
35 : /// The LB's notice window (default 2 s): how long to keep serving AFTER going
36 : /// unhealthy, sized so a load balancer's health poll actually observes the 503
37 : /// before the socket goes. It is NOT the in-flight ceiling — that is
38 : /// `drain_deadline`, and it is 10× longer.
39 75 : pub fn drain_delay() -> Result<Duration, String> {
40 75 : env_num("ANTARES_DRAIN_DELAY_MS", 2000).map(Duration::from_millis)
41 75 : }
42 :
43 : /// The real shutdown deadline (default 20 s): the ceiling on waiting for
44 : /// in-flight work — connections, then the outbox — once the listener is
45 : /// closed. Both share it; it does not extend the notice window above.
46 65 : pub fn drain_deadline() -> Result<Duration, String> {
47 65 : env_num("ANTARES_DRAIN_DEADLINE_SECS", 20).map(Duration::from_secs)
48 65 : }
49 :
50 : /// Absent = the documented default; present-but-unparsable is fatal. A
51 : /// misread drain window silently running at the default is the same class of
52 : /// misconfiguration as an unknown key, and is refused the same way. Zero is a
53 : /// real choice on both knobs (no notice window / close at once).
54 140 : fn env_num(key: &str, default: u64) -> Result<u64, String> {
55 140 : match std::env::var(key) {
56 102 : Err(std::env::VarError::NotPresent) => Ok(default),
57 0 : Err(e) => Err(format!("{key} is unreadable: {e}")),
58 38 : Ok(v) => v
59 38 : .parse::<u64>()
60 38 : .map_err(|e| format!("{key} must be a non-negative integer, got {v:?} ({e})")),
61 : }
62 140 : }
63 :
64 : /// Resolves on SIGTERM or SIGINT. SIGTERM is the one that matters —
65 : /// it is what every orchestrator sends — and listening only for ctrl_c (the
66 : /// v0 behaviour) meant a `docker stop` or a pod eviction dropped every
67 : /// in-flight request on the floor.
68 8 : pub async fn signal() {
69 : #[cfg(unix)]
70 : {
71 : use tokio::signal::unix::{signal, SignalKind};
72 8 : let mut term = match signal(SignalKind::terminate()) {
73 8 : Ok(s) => s,
74 0 : Err(e) => {
75 0 : tracing::warn!("cannot listen for SIGTERM ({e}); ctrl-c only");
76 0 : let _ = tokio::signal::ctrl_c().await;
77 0 : return;
78 : }
79 : };
80 8 : tokio::select! {
81 8 : _ = term.recv() => tracing::info!("SIGTERM received"),
82 8 : _ = tokio::signal::ctrl_c() => tracing::info!("SIGINT received"),
83 : }
84 : }
85 : #[cfg(not(unix))]
86 : {
87 : let _ = tokio::signal::ctrl_c().await;
88 : }
89 8 : }
90 :
91 : /// Steps 4–6. Steps 1–3 belong to the accept loop, which owns the listener.
92 : /// The deadline is passed in, not read here: the composition root parses every
93 : /// config value once, at startup, so a garbage window fails before serving.
94 16 : pub async fn drain(
95 16 : inflight: &Arc<AtomicUsize>,
96 16 : pending_changes: &AtomicUsize,
97 16 : store: &dyn antares_store::CurrentStateDriver,
98 16 : temporal: &dyn antares_store::TemporalDriver,
99 16 : deadline: Duration,
100 16 : flush_outbox: bool,
101 16 : ) {
102 16 : let started = Instant::now();
103 : // A request is not over when its response is: the remote leg of a
104 : // distributed subscription, an initial Context Source notification and a
105 : // forwarded notification all run as tasks after the 2xx, and a stop that
106 : // dropped them left the subscription chain half-built on every roll.
107 : // The matcher queue is part of a request too: a change accepted before
108 : // the listener closed still owes its notifications (5.8.6).
109 66 : while inflight.load(Ordering::Relaxed) > 0
110 38 : || antares_api::background_tasks() > 0
111 14 : || pending_changes.load(Ordering::SeqCst) > 0
112 : {
113 52 : if started.elapsed() >= deadline {
114 2 : tracing::warn!(
115 : "drain deadline {deadline:?} hit with {} connection(s), {} task(s) and {} change batch(es) still open — closing anyway",
116 0 : inflight.load(Ordering::Relaxed),
117 0 : antares_api::background_tasks(),
118 0 : pending_changes.load(Ordering::SeqCst)
119 : );
120 2 : break;
121 50 : }
122 50 : tokio::time::sleep(Duration::from_millis(25)).await;
123 : }
124 : // The outbox is drained by a background task on the api pods. Stopping
125 : // here — after the last request has committed its row, before the pool
126 : // closes — gives that task the chance to publish what is still pending,
127 : // so a rolling update does not leave events sitting in the table until
128 : // another pod's fallback poll finds them. Stores without an outbox
129 : // (memory, file) answer an empty page and fall straight through.
130 : //
131 : // Two residuals, stated rather than papered over. The table is shared, so
132 : // this waits for rows OTHER pods are still producing too, and under
133 : // sustained write load it therefore runs to the deadline; and the deadline
134 : // is the same one step 4 just spent, so a slow in-flight wait can leave
135 : // the flush no time at all. Either way the rows stay committed — the
136 : // fallback poll on a surviving pod publishes them — so the ceiling costs
137 : // latency, never an event.
138 16 : if flush_outbox {
139 : loop {
140 730 : match store.outbox_peek(1).await {
141 730 : Ok(rows) if rows.is_empty() => break,
142 721 : Ok(_) => {}
143 0 : Err(e) => {
144 0 : tracing::warn!("outbox flush gave up: {e}");
145 0 : break;
146 : }
147 : }
148 721 : if started.elapsed() >= deadline {
149 1 : tracing::warn!("drain deadline {deadline:?} hit with outbox rows still pending");
150 1 : break;
151 720 : }
152 720 : tokio::time::sleep(Duration::from_millis(25)).await;
153 : }
154 6 : }
155 16 : store.close().await;
156 : // Both seams, because the temporal half may be a store of its own
157 : // (`ANTARES_TEMPORAL` naming a second backend) with its own pool. When
158 : // one instance serves both, the second call lands on an already-closed
159 : // pool and does nothing.
160 16 : temporal.close().await;
161 16 : tracing::info!("drain complete in {:?}", started.elapsed());
162 16 : }
163 :
164 : /// Step 1, so the flip and the log line stay in one place.
165 10 : pub fn begin(draining: &Arc<AtomicBool>, delay: Duration) {
166 10 : draining.store(true, Ordering::Relaxed);
167 : // Immediate, not sampler-paced — a roll must be visible on a
168 : // dashboard for its whole (short) duration.
169 10 : metrics::gauge!("antares_draining").set(1.0);
170 10 : tracing::info!("draining: /q/health now 503 for {delay:?} before the listener closes");
171 10 : }
172 :
173 : #[cfg(test)]
174 : mod tests {
175 : use super::*;
176 :
177 : /// The two drain variables are process-global: every test that sets or
178 : /// clears them holds this, or a parallel test reads its value.
179 : static DRAIN_ENV: std::sync::Mutex<()> = std::sync::Mutex::new(());
180 :
181 4 : fn drain_env() -> std::sync::MutexGuard<'static, ()> {
182 4 : DRAIN_ENV
183 4 : .lock()
184 4 : .unwrap_or_else(std::sync::PoisonError::into_inner)
185 4 : }
186 :
187 : /// A temporal driver that only records whether it was closed; every
188 : /// operation answers "no temporal store", the same shape as `NoTemporal`.
189 : struct CountsCloses(Arc<AtomicUsize>);
190 :
191 : impl CountsCloses {
192 0 : fn off<T>() -> Result<T, antares_model::NgsiError> {
193 0 : Err(antares_model::NgsiError::OperationNotSupported(
194 0 : "test driver".into(),
195 0 : ))
196 0 : }
197 : }
198 :
199 : #[async_trait::async_trait]
200 : impl antares_store::TemporalDriver for CountsCloses {
201 2 : async fn close(&self) {
202 : self.0.fetch_add(1, Ordering::SeqCst);
203 2 : }
204 : async fn temporal_append(
205 : &self,
206 : _t: &antares_model::TenantId,
207 : _id: &str,
208 : _shell: &serde_json::Value,
209 : _add: &serde_json::Value,
210 0 : ) -> Result<(), antares_model::NgsiError> {
211 : Self::off()
212 0 : }
213 : async fn query_temporal(
214 : &self,
215 : _t: &antares_model::TenantId,
216 : _f: &antares_store::filter::TemporalFilter<'_>,
217 0 : ) -> Result<antares_store::filter::TemporalOutcome, antares_model::NgsiError> {
218 : Self::off()
219 0 : }
220 : async fn get_temporal(
221 : &self,
222 : _t: &antares_model::TenantId,
223 : _id: &str,
224 : _f: &antares_store::filter::TemporalFilter<'_>,
225 0 : ) -> Result<Option<serde_json::Value>, antares_model::NgsiError> {
226 : Self::off()
227 0 : }
228 : async fn get(
229 : &self,
230 : _t: &antares_model::TenantId,
231 : _id: &str,
232 0 : ) -> Result<Option<serde_json::Value>, antares_model::NgsiError> {
233 : Self::off()
234 0 : }
235 : async fn create(
236 : &self,
237 : _t: &antares_model::TenantId,
238 : _id: &str,
239 : _d: serde_json::Value,
240 0 : ) -> Result<bool, antares_model::NgsiError> {
241 : Self::off()
242 0 : }
243 : async fn upsert(
244 : &self,
245 : _t: &antares_model::TenantId,
246 : _id: &str,
247 : _d: serde_json::Value,
248 0 : ) -> Result<bool, antares_model::NgsiError> {
249 : Self::off()
250 0 : }
251 : async fn delete(
252 : &self,
253 : _t: &antares_model::TenantId,
254 : _id: &str,
255 0 : ) -> Result<bool, antares_model::NgsiError> {
256 : Self::off()
257 0 : }
258 : async fn list(
259 : &self,
260 : _t: &antares_model::TenantId,
261 0 : ) -> Result<Vec<serde_json::Value>, antares_model::NgsiError> {
262 : Self::off()
263 0 : }
264 : async fn mutate_boxed<'a>(
265 : &self,
266 : _t: &antares_model::TenantId,
267 : _id: &str,
268 : _f: antares_store::MutateFn<'a>,
269 0 : ) -> Result<Option<Result<(), ()>>, antares_model::NgsiError> {
270 : Self::off()
271 0 : }
272 : }
273 :
274 : /// The drain closes BOTH driver seams. With `ANTARES_TEMPORAL` naming a
275 : /// backend of its own the temporal half is a second store holding its own
276 : /// connection pool; closing only the current-state store left that pool
277 : /// open for process teardown to sever, with whatever it still owed
278 : /// in flight.
279 : #[tokio::test(flavor = "multi_thread")]
280 2 : async fn drain_closes_the_temporal_driver_too() {
281 2 : let closes = Arc::new(AtomicUsize::new(0));
282 2 : let temporal = CountsCloses(Arc::clone(&closes));
283 2 : let store = antares_sql::store::any::AnyStore::Mem(antares_sql::store::Store::default());
284 2 : drain(
285 2 : &Arc::new(AtomicUsize::new(0)),
286 2 : &AtomicUsize::new(0),
287 2 : &store,
288 2 : &temporal,
289 2 : Duration::from_millis(50),
290 2 : false,
291 2 : )
292 2 : .await;
293 2 : assert_eq!(
294 2 : closes.load(Ordering::SeqCst),
295 2 : 1,
296 2 : "the temporal driver must be closed by the drain, not by process exit"
297 2 : );
298 2 : }
299 :
300 : /// The book states these two defaults, and their SUM is an operator
301 : /// contract: the container stop grace period has to exceed it or the
302 : /// orchestrator turns a drain into a kill. Nothing tied the stated
303 : /// numbers to the ones the binary uses, and the delay drifted to a value
304 : /// 1.5 s shorter than the truth in three chapters at once — a grace
305 : /// period sized from the book would then have been under the real drain.
306 : /// `dev/check-env-docs.sh` proves each variable is documented; this
307 : /// proves the documented numbers are the ones that run.
308 : #[test]
309 2 : fn the_documented_drain_defaults_are_the_ones_the_binary_uses() {
310 2 : let _env = drain_env();
311 2 : let book = std::fs::read_to_string(concat!(
312 : env!("CARGO_MANIFEST_DIR"),
313 : "/../../docs/src/configuration.md"
314 : ))
315 2 : .expect("the configuration chapter");
316 4 : let stated = |var: &str| -> String {
317 4 : let row = book
318 4 : .lines()
319 398 : .find(|l| l.starts_with(&format!("| `{var}` |")))
320 4 : .unwrap_or_else(|| panic!("{var} has no row in the configuration table"));
321 4 : row.split('|')
322 4 : .nth(2)
323 4 : .expect("the default column")
324 4 : .trim()
325 4 : .trim_matches('`')
326 4 : .to_owned()
327 4 : };
328 2 : std::env::remove_var("ANTARES_DRAIN_DELAY_MS");
329 2 : std::env::remove_var("ANTARES_DRAIN_DEADLINE_SECS");
330 2 : assert_eq!(
331 2 : stated("ANTARES_DRAIN_DELAY_MS"),
332 2 : drain_delay().expect("absent").as_millis().to_string(),
333 : "the book's notice window is not the compiled one"
334 : );
335 2 : assert_eq!(
336 2 : stated("ANTARES_DRAIN_DEADLINE_SECS"),
337 2 : drain_deadline().expect("absent").as_secs().to_string(),
338 : "the book's in-flight ceiling is not the compiled one"
339 : );
340 :
341 : // The shipped compose files justify their stop_grace_period against
342 : // the same two defaults, in a comment an operator copies the number
343 : // out of. Both said 0.5 s long after the delay became 2 s.
344 2 : let secs = drain_delay().expect("absent").as_secs_f64();
345 2 : let want = format!(
346 : "drain delay ({} s)",
347 2 : if secs.fract() == 0.0 {
348 2 : format!("{secs:.0}")
349 : } else {
350 0 : format!("{secs}")
351 : }
352 : );
353 4 : for name in ["docker-compose-ha.yml", "docker-compose-roles.yml"] {
354 4 : let path = format!(
355 : concat!(env!("CARGO_MANIFEST_DIR"), "/../../compose-files/{}"),
356 : name
357 : );
358 4 : let text = std::fs::read_to_string(&path).expect("the compose file");
359 4 : assert!(
360 4 : text.contains(&want),
361 : "{name} does not justify stop_grace_period against \"{want}\""
362 : );
363 : }
364 2 : }
365 :
366 : /// Both drain knobs in ONE test: the environment is process-global, so
367 : /// parsing them from parallel test threads would race.
368 : ///
369 : /// Contract: absent = the documented default; present-but-unparsable is
370 : /// FATAL, never a silent default — a misread timeout is exactly the class
371 : /// of misconfiguration the unknown-key policy exists to catch. The two
372 : /// defaults are different numbers for different jobs: 2000 MILLIseconds of
373 : /// LB notice, 20 SECONDS of in-flight ceiling.
374 : #[test]
375 2 : fn drain_knobs_default_when_absent_and_refuse_garbage() {
376 2 : let _env = drain_env();
377 2 : std::env::remove_var("ANTARES_DRAIN_DELAY_MS");
378 2 : std::env::remove_var("ANTARES_DRAIN_DEADLINE_SECS");
379 2 : assert_eq!(drain_delay().expect("absent"), Duration::from_millis(2000));
380 2 : assert_eq!(drain_deadline().expect("absent"), Duration::from_secs(20));
381 2 : assert_ne!(
382 2 : drain_delay().expect("absent"),
383 2 : drain_deadline().expect("absent"),
384 : "the notice window and the in-flight ceiling are not the same number"
385 : );
386 :
387 2 : std::env::set_var("ANTARES_DRAIN_DELAY_MS", "750");
388 2 : std::env::set_var("ANTARES_DRAIN_DEADLINE_SECS", "10");
389 2 : assert_eq!(drain_delay().expect("set"), Duration::from_millis(750));
390 2 : assert_eq!(drain_deadline().expect("set"), Duration::from_secs(10));
391 :
392 : // Zero is a real choice on both knobs (no notice window / close at
393 : // once), so it must NOT be rejected with the garbage.
394 2 : std::env::set_var("ANTARES_DRAIN_DELAY_MS", "0");
395 2 : std::env::set_var("ANTARES_DRAIN_DEADLINE_SECS", "0");
396 2 : assert_eq!(drain_delay().expect("zero"), Duration::ZERO);
397 2 : assert_eq!(drain_deadline().expect("zero"), Duration::ZERO);
398 :
399 14 : for bad in [
400 2 : "soon",
401 2 : "",
402 2 : "-1",
403 2 : "2.5",
404 2 : "500ms",
405 2 : "99999999999999999999999",
406 2 : " 5",
407 2 : ] {
408 14 : std::env::set_var("ANTARES_DRAIN_DELAY_MS", bad);
409 14 : let err =
410 14 : drain_delay().expect_err(&format!("ANTARES_DRAIN_DELAY_MS={bad:?} must be fatal"));
411 14 : assert!(
412 14 : err.contains("ANTARES_DRAIN_DELAY_MS"),
413 : "the error must name the key: {err}"
414 : );
415 : }
416 2 : std::env::set_var("ANTARES_DRAIN_DELAY_MS", "2000");
417 8 : for bad in ["soon", "", "-1", "20.0"] {
418 8 : std::env::set_var("ANTARES_DRAIN_DEADLINE_SECS", bad);
419 8 : let err = drain_deadline().expect_err(&format!(
420 8 : "ANTARES_DRAIN_DEADLINE_SECS={bad:?} must be fatal"
421 8 : ));
422 8 : assert!(err.contains("ANTARES_DRAIN_DEADLINE_SECS"), "{err}");
423 : }
424 2 : std::env::remove_var("ANTARES_DRAIN_DELAY_MS");
425 2 : std::env::remove_var("ANTARES_DRAIN_DEADLINE_SECS");
426 2 : }
427 :
428 6 : fn rt() -> tokio::runtime::Runtime {
429 6 : tokio::runtime::Builder::new_current_thread()
430 6 : .enable_all()
431 6 : .build()
432 6 : .expect("runtime")
433 6 : }
434 :
435 6 : fn mem_store() -> antares_sql::store::any::AnyStore {
436 6 : antares_sql::store::any::AnyStore::Mem(antares_sql::store::Store::default())
437 6 : }
438 :
439 : /// A request's follow-up work (the remote leg of a distributed
440 : /// subscription, a forwarded notification) runs after its response; the
441 : /// drain waits for it like it waits for the request itself.
442 : #[test]
443 2 : fn drain_waits_for_request_born_tasks() {
444 2 : rt().block_on(async {
445 2 : let inflight = Arc::new(AtomicUsize::new(0));
446 2 : let done = Arc::new(AtomicBool::new(false));
447 2 : let d = done.clone();
448 2 : antares_api::spawn(async move {
449 2 : tokio::time::sleep(Duration::from_millis(300)).await;
450 2 : d.store(true, Ordering::SeqCst);
451 2 : });
452 2 : let store = mem_store();
453 2 : drain(
454 2 : &inflight,
455 2 : &AtomicUsize::new(0),
456 2 : &store,
457 2 : &antares_store::NoTemporal,
458 2 : Duration::from_secs(5),
459 2 : false,
460 2 : )
461 2 : .await;
462 2 : assert!(
463 2 : done.load(Ordering::SeqCst),
464 : "drain returned before the task finished"
465 : );
466 2 : assert_eq!(antares_api::background_tasks(), 0);
467 2 : });
468 2 : }
469 :
470 : /// Nothing in flight = nothing to wait for: the drain must not sit out
471 : /// its deadline, and the outbox flush must not hang a store that has no
472 : /// outbox (memory/file).
473 : #[test]
474 2 : fn drain_returns_at_once_when_nothing_is_in_flight() {
475 2 : let inflight = Arc::new(AtomicUsize::new(0));
476 2 : let store = mem_store();
477 2 : let started = Instant::now();
478 2 : rt().block_on(drain(
479 2 : &inflight,
480 2 : &AtomicUsize::new(0),
481 2 : &store,
482 2 : &antares_store::NoTemporal,
483 2 : Duration::from_secs(20),
484 : true,
485 : ));
486 2 : assert!(
487 2 : started.elapsed() < Duration::from_secs(1),
488 : "an idle drain waited {:?} — it must not burn the deadline",
489 0 : started.elapsed()
490 : );
491 2 : }
492 :
493 : /// The deadline is a CEILING, not a promise: a connection that never
494 : /// finishes must not hold the process open forever.
495 : #[test]
496 2 : fn drain_gives_up_at_the_deadline_with_a_stuck_connection() {
497 2 : let inflight = Arc::new(AtomicUsize::new(1)); // never released
498 2 : let store = mem_store();
499 2 : let started = Instant::now();
500 2 : rt().block_on(drain(
501 2 : &inflight,
502 2 : &AtomicUsize::new(0),
503 2 : &store,
504 2 : &antares_store::NoTemporal,
505 2 : Duration::from_millis(300),
506 : false,
507 : ));
508 2 : let waited = started.elapsed();
509 2 : assert!(
510 2 : waited >= Duration::from_millis(300),
511 : "the drain must actually wait for in-flight work: {waited:?}"
512 : );
513 2 : assert!(
514 2 : waited < Duration::from_secs(3),
515 : "the drain must give up AT the deadline, not later: {waited:?}"
516 : );
517 2 : }
518 :
519 : /// The outbox flush must WAIT while rows are still pending, and only when
520 : /// it is asked to: a pod whose own drain is off publishes nothing, so
521 : /// waiting there would only burn the deadline. Needs a live database —
522 : /// the outbox table exists on the Pg arm alone, so the memory store can
523 : /// never exercise the wait (it answers an empty page and falls through).
524 : #[tokio::test(flavor = "multi_thread")]
525 : #[ignore = "needs a live database (ANTARES_TEST_DATABASE_URL)"]
526 0 : async fn outbox_flush_waits_for_pending_rows_and_only_when_asked() {
527 : use antares_sql::store::any::{AnyStore, PgBackend};
528 : use antares_sql::store::Kind;
529 0 : let url = std::env::var("ANTARES_TEST_DATABASE_URL")
530 0 : .expect("ANTARES_TEST_DATABASE_URL: this test is asked for by name where a DB exists");
531 : // a nested fn, not a closure: connecting is awaited, and an async
532 : // closure is not a stable language feature
533 0 : async fn connect(url: &str) -> AnyStore {
534 0 : AnyStore::Pg(PgBackend::new(
535 0 : antares_sql::store::pg::connect(url, 5)
536 0 : .await
537 0 : .expect("connect+migrate"),
538 : ))
539 0 : }
540 0 : let run = std::time::SystemTime::now()
541 0 : .duration_since(std::time::UNIX_EPOCH)
542 0 : .expect("clock")
543 0 : .as_millis();
544 0 : let tenant = antares_model::TenantId::new(&format!("drain{run}")).expect("tenant");
545 0 : let id = format!("urn:ngsi-ld:DrainProbe:{run}");
546 0 : let inflight = Arc::new(AtomicUsize::new(0));
547 :
548 : // One committed-but-unpublished row: outbox on, then a write. Nothing
549 : // publishes it here — a unit test wires no bus drain task.
550 0 : let store = connect(&url).await;
551 0 : store.set_outbox(true);
552 0 : store
553 0 : .create(
554 0 : &tenant,
555 0 : Kind::Entity,
556 0 : &id,
557 0 : serde_json::json!({"id": id.as_str(), "type": "DrainProbe"}),
558 0 : )
559 0 : .await
560 0 : .expect("write");
561 0 : let mine: Vec<i64> = store
562 0 : .outbox_peek(500)
563 0 : .await
564 0 : .expect("peek")
565 0 : .into_iter()
566 0 : .filter(|(_, t, _)| t == tenant.as_str())
567 0 : .map(|(seq, ..)| seq)
568 0 : .collect();
569 0 : assert!(
570 0 : !mine.is_empty(),
571 : "the write enqueued no outbox row — the rest of this test would prove nothing"
572 : );
573 :
574 : // Not asked to flush: the pending row may not delay the close at all.
575 0 : let t0 = Instant::now();
576 0 : drain(
577 0 : &inflight,
578 0 : &AtomicUsize::new(0),
579 0 : &store,
580 0 : &antares_store::NoTemporal,
581 0 : Duration::from_millis(400),
582 0 : false,
583 0 : )
584 0 : .await;
585 0 : let closed = t0.elapsed();
586 0 : assert!(
587 0 : closed < Duration::from_millis(250),
588 : "flush_outbox=false waited {closed:?} on a row it never intended to publish"
589 : );
590 :
591 : // Asked to flush: the rows stay pending, so the flush must hold the
592 : // process to its deadline rather than exit on top of them.
593 0 : let store = connect(&url).await;
594 0 : let t0 = Instant::now();
595 0 : drain(
596 0 : &inflight,
597 0 : &AtomicUsize::new(0),
598 0 : &store,
599 0 : &antares_store::NoTemporal,
600 0 : Duration::from_millis(400),
601 0 : true,
602 0 : )
603 0 : .await;
604 0 : let waited = t0.elapsed();
605 0 : assert!(
606 0 : waited >= Duration::from_millis(400),
607 : "the flush returned after {waited:?} with rows still pending"
608 : );
609 0 : assert!(
610 0 : waited < Duration::from_secs(5),
611 : "the flush must give up AT the deadline, not later: {waited:?}"
612 : );
613 :
614 : // Leave the shared table as it was found: ack only our own seqs (a
615 : // blanket ack would delete another test's pending rows) and drop the
616 : // probe entity with the outbox off, so the delete enqueues nothing.
617 0 : let store = connect(&url).await;
618 0 : store
619 0 : .delete(&tenant, Kind::Entity, &id)
620 0 : .await
621 0 : .expect("probe cleanup");
622 0 : store.outbox_ack(&mine).await.expect("outbox cleanup");
623 0 : assert!(
624 0 : store
625 0 : .outbox_peek(500)
626 0 : .await
627 0 : .expect("peek")
628 0 : .into_iter()
629 0 : .all(|(_, t, _)| t != tenant.as_str()),
630 0 : "the test left its own rows in the shared outbox"
631 0 : );
632 0 : }
633 :
634 : /// Step 1 is the flag the health endpoint reads; nothing else may flip it.
635 : #[test]
636 2 : fn begin_flips_the_health_flag() {
637 2 : let draining = Arc::new(AtomicBool::new(false));
638 2 : assert!(!draining.load(Ordering::Relaxed));
639 2 : begin(&draining, Duration::from_millis(500));
640 2 : assert!(draining.load(Ordering::Relaxed), "/q/health must go 503");
641 2 : }
642 : }
|