1use std::{
2 collections::{HashMap, VecDeque},
3 sync::LazyLock,
4};
5
6use dora_message::{
7 config::{DEFAULT_QUEUE_SIZE, QueuePolicy},
8 daemon_to_node::NodeEvent,
9 id::DataId,
10 metadata::{GOAL_ID, GOAL_STATUS, REQUEST_ID, get_string_param},
11};
12
13use super::thread::EventItem;
14
15fn is_correlated(event: &EventItem) -> bool {
23 let params = match event {
24 EventItem::NodeEvent {
25 event: NodeEvent::Input { metadata, .. },
26 ..
27 } => &metadata.parameters,
28 EventItem::ZenohInput { metadata, .. } => &metadata.parameters,
29 _ => return false,
30 };
31 params.contains_key(REQUEST_ID)
32 || params.contains_key(GOAL_ID)
33 || params.contains_key(GOAL_STATUS)
34}
35
36enum Eviction {
38 RemoveAt(usize),
40 DropIncoming,
43 DropFrontLoud,
46}
47
48fn select_eviction(queue: &VecDeque<EventItem>, incoming: &EventItem) -> Eviction {
54 if let Some(idx) = queue.iter().position(|e| !is_correlated(e)) {
55 return Eviction::RemoveAt(idx);
56 }
57 if !is_correlated(incoming) {
58 return Eviction::DropIncoming;
59 }
60 Eviction::DropFrontLoud
61}
62
63fn log_correlation_drop(event_id: &DataId, dropped: &EventItem) {
67 let params = match dropped {
68 EventItem::NodeEvent {
69 event: NodeEvent::Input { metadata, .. },
70 ..
71 } => &metadata.parameters,
72 EventItem::ZenohInput { metadata, .. } => &metadata.parameters,
73 _ => return,
74 };
75 let request_id = get_string_param(params, REQUEST_ID);
76 let goal_id = get_string_param(params, GOAL_ID);
77 let goal_status = get_string_param(params, GOAL_STATUS);
78 tracing::error!(
79 input = %event_id,
80 ?request_id,
81 ?goal_id,
82 ?goal_status,
83 "queue full of correlated messages; dropping oldest correlation. \
84 This breaks the service/action request-response contract. \
85 Consider increasing queue_size or switching this input to \
86 `queue_policy: backpressure`."
87 );
88}
89pub(crate) const NON_INPUT_EVENT: &str = "dora.non_input_event";
90
91static NON_INPUT_EVENT_ID: LazyLock<DataId> =
94 LazyLock::new(|| DataId::from(NON_INPUT_EVENT.to_string()));
95
96#[derive(Debug)]
168pub struct Scheduler {
169 last_used: VecDeque<DataId>,
171 event_queues: HashMap<DataId, (usize, VecDeque<EventItem>)>,
173 queue_policies: HashMap<DataId, QueuePolicy>,
175 dropped: HashMap<DataId, u64>,
177}
178
179impl Scheduler {
180 pub(crate) fn with_policies(
181 event_queues: HashMap<DataId, (usize, VecDeque<EventItem>)>,
182 queue_policies: HashMap<DataId, QueuePolicy>,
183 ) -> Self {
184 let topic = VecDeque::from_iter(
185 event_queues
186 .keys()
187 .filter(|t| **t != *NON_INPUT_EVENT_ID)
188 .cloned(),
189 );
190 Self {
191 last_used: topic,
192 event_queues,
193 queue_policies,
194 dropped: HashMap::new(),
195 }
196 }
197
198 pub fn drain_drop_counts(&mut self) -> HashMap<DataId, u64> {
200 std::mem::take(&mut self.dropped)
201 }
202
203 pub(crate) fn add_event(&mut self, event: EventItem) {
204 let (event_id, should_flush) = match &event {
205 EventItem::NodeEvent {
206 event: NodeEvent::Input { id, metadata, .. },
207 ..
208 } => {
209 let flush = dora_message::metadata::get_bool_param(
210 &metadata.parameters,
211 dora_message::metadata::FLUSH,
212 ) == Some(true);
213 (id, flush)
214 }
215 EventItem::ZenohInput { id, metadata, .. } => {
216 let flush = dora_message::metadata::get_bool_param(
217 &metadata.parameters,
218 dora_message::metadata::FLUSH,
219 ) == Some(true);
220 (id, flush)
221 }
222 _ => (&*NON_INPUT_EVENT_ID, false),
223 };
224
225 if should_flush && let Some((_size, queue)) = self.event_queues.get_mut(event_id) {
234 let before = queue.len();
235 queue.retain(is_correlated);
236 let drained = before - queue.len();
237 if drained > 0 {
238 tracing::debug!(
239 "Flushed {drained} queued event(s) for input `{event_id}` (flush signal)"
240 );
241 }
242 if !queue.is_empty() {
243 tracing::debug!(
244 input = %event_id,
245 preserved = queue.len(),
246 "flush signal retained correlated (request_id/goal_id) events"
247 );
248 }
249 }
250
251 if !self.event_queues.contains_key(event_id) {
258 tracing::warn!(
259 "no queue config for input `{event_id}`, using default size {DEFAULT_QUEUE_SIZE}"
260 );
261 self.last_used.push_back(event_id.clone());
262 self.event_queues
263 .insert(event_id.clone(), (DEFAULT_QUEUE_SIZE, Default::default()));
264 }
265 let Some((size, queue)) = self.event_queues.get_mut(event_id) else {
266 return;
268 };
269
270 let policy = self
271 .queue_policies
272 .get(event_id)
273 .copied()
274 .unwrap_or_default();
275
276 let cap = policy.effective_cap(*size);
277 if queue.len() >= cap {
278 if policy == QueuePolicy::Backpressure {
279 tracing::error!(
280 "Backpressure input `{event_id}` hit hard cap ({cap}), \
281 dropping oldest to prevent OOM"
282 );
283 } else {
284 tracing::warn!("Discarding event for input `{event_id}` due to queue size limit");
285 }
286 *self.dropped.entry(event_id.clone()).or_insert(0) += 1;
287 match select_eviction(queue, &event) {
288 Eviction::RemoveAt(idx) => {
289 queue.remove(idx);
290 }
291 Eviction::DropIncoming => {
292 return;
295 }
296 Eviction::DropFrontLoud => {
297 if let Some(front) = queue.pop_front() {
298 log_correlation_drop(event_id, &front);
299 }
300 }
301 }
302 }
303 queue.push_back(event);
304 }
305
306 pub(crate) fn next(&mut self) -> Option<EventItem> {
307 if let Some((_size, queue)) = self.event_queues.get_mut(&*NON_INPUT_EVENT_ID)
309 && let Some(event) = queue.pop_front()
310 {
311 return Some(event);
312 }
313
314 for index in 0..self.last_used.len() {
316 let id = &self.last_used[index];
317 if let Some((_size, queue)) = self.event_queues.get_mut(id)
318 && let Some(event) = queue.pop_front()
319 {
320 if let Some(id) = self.last_used.remove(index) {
322 self.last_used.push_back(id);
323 }
324 return Some(event);
325 }
326 }
327
328 None
329 }
330
331 pub(crate) fn is_empty(&self) -> bool {
332 self.event_queues
333 .iter()
334 .all(|(_id, (_size, queue))| queue.is_empty())
335 }
336}
337
338#[cfg(test)]
339mod tests {
340 use super::*;
341 use crate::uhlc;
342 use dora_message::{
343 daemon_to_node::NodeEvent,
344 metadata::{FLUSH, Metadata, MetadataParameters, Parameter},
345 };
346
347 fn make_input(id: &str, params: MetadataParameters) -> EventItem {
348 let ts = uhlc::HLC::default().new_timestamp();
349 let metadata = Metadata::from_parameters(ts, params);
350 EventItem::NodeEvent {
351 event: NodeEvent::Input {
352 id: DataId::from(id.to_string()),
353 metadata: std::sync::Arc::new(metadata),
354 data: None,
355 },
356 }
357 }
358
359 fn make_scheduler(audio_capacity: usize) -> (Scheduler, DataId) {
360 let id = DataId::from("audio".to_string());
361 let mut queues = HashMap::new();
362 queues.insert(id.clone(), (audio_capacity, VecDeque::new()));
363 queues.insert(
364 DataId::from(NON_INPUT_EVENT.to_string()),
365 (10, VecDeque::new()),
366 );
367 (Scheduler::with_policies(queues, HashMap::new()), id)
368 }
369
370 #[test]
371 fn flush_clears_older_queued_events() {
372 let (mut sched, id) = make_scheduler(10);
373
374 sched.add_event(make_input("audio", MetadataParameters::new()));
375 sched.add_event(make_input("audio", MetadataParameters::new()));
376 sched.add_event(make_input("audio", MetadataParameters::new()));
377 assert_eq!(sched.event_queues[&id].1.len(), 3);
378
379 let mut flush_params = MetadataParameters::new();
381 flush_params.insert(FLUSH.into(), Parameter::Bool(true));
382 sched.add_event(make_input("audio", flush_params));
383
384 assert_eq!(sched.event_queues[&id].1.len(), 1);
385 }
386
387 #[test]
388 fn non_flush_does_not_clear_queue() {
389 let (mut sched, id) = make_scheduler(10);
390
391 sched.add_event(make_input("audio", MetadataParameters::new()));
392 sched.add_event(make_input("audio", MetadataParameters::new()));
393 sched.add_event(make_input("audio", MetadataParameters::new()));
394 assert_eq!(sched.event_queues[&id].1.len(), 3);
395 }
396
397 #[test]
398 fn flush_false_does_not_clear_queue() {
399 let (mut sched, id) = make_scheduler(10);
400
401 sched.add_event(make_input("audio", MetadataParameters::new()));
402 sched.add_event(make_input("audio", MetadataParameters::new()));
403
404 let mut params = MetadataParameters::new();
405 params.insert(FLUSH.into(), Parameter::Bool(false));
406 sched.add_event(make_input("audio", params));
407
408 assert_eq!(sched.event_queues[&id].1.len(), 3);
409 }
410
411 #[test]
412 fn flush_with_queue_size_one_retains_flush_message() {
413 let (mut sched, id) = make_scheduler(1);
414
415 sched.add_event(make_input("audio", MetadataParameters::new()));
416 assert_eq!(sched.event_queues[&id].1.len(), 1);
417
418 let mut flush_params = MetadataParameters::new();
420 flush_params.insert(FLUSH.into(), Parameter::Bool(true));
421 sched.add_event(make_input("audio", flush_params));
422
423 assert_eq!(sched.event_queues[&id].1.len(), 1);
425 }
426
427 #[test]
428 fn drop_oldest_tracks_drop_count() {
429 let (mut sched, id) = make_scheduler(2);
430
431 sched.add_event(make_input("audio", MetadataParameters::new()));
433 sched.add_event(make_input("audio", MetadataParameters::new()));
434 assert_eq!(sched.event_queues[&id].1.len(), 2);
435
436 sched.add_event(make_input("audio", MetadataParameters::new()));
438 sched.add_event(make_input("audio", MetadataParameters::new()));
439 sched.add_event(make_input("audio", MetadataParameters::new()));
440
441 assert_eq!(sched.event_queues[&id].1.len(), 2);
443
444 let counts = sched.drain_drop_counts();
446 assert_eq!(counts.get(&id), Some(&3));
447
448 let counts = sched.drain_drop_counts();
450 assert!(counts.is_empty());
451 }
452
453 #[test]
456 fn flush_retains_correlated_events() {
457 let (mut sched, id) = make_scheduler(10);
460
461 sched.add_event(make_input("audio", with_request_id("req-1")));
462 sched.add_event(make_input("audio", MetadataParameters::new()));
463 sched.add_event(make_input("audio", MetadataParameters::new()));
464 assert_eq!(sched.event_queues[&id].1.len(), 3);
465
466 let mut flush_params = MetadataParameters::new();
467 flush_params.insert(FLUSH.into(), Parameter::Bool(true));
468 sched.add_event(make_input("audio", flush_params));
469
470 let queue = &sched.event_queues[&id].1;
471 assert_eq!(queue.len(), 2);
473 assert!(
474 queue
475 .iter()
476 .any(|e| request_id_of(e).as_deref() == Some("req-1")),
477 "service response with request_id was wiped by flush"
478 );
479 }
480
481 #[test]
482 fn flush_retains_goal_id_events() {
483 let (mut sched, id) = make_scheduler(10);
485
486 let mut goal_params = MetadataParameters::new();
487 goal_params.insert(GOAL_ID.into(), Parameter::String("goal-7".to_string()));
488 sched.add_event(make_input("audio", goal_params));
489 sched.add_event(make_input("audio", MetadataParameters::new()));
490
491 let mut flush_params = MetadataParameters::new();
492 flush_params.insert(FLUSH.into(), Parameter::Bool(true));
493 sched.add_event(make_input("audio", flush_params));
494
495 let queue = &sched.event_queues[&id].1;
496 assert_eq!(queue.len(), 2);
498 let has_goal = queue.iter().any(|e| {
499 let EventItem::NodeEvent {
500 event: NodeEvent::Input { metadata, .. },
501 ..
502 } = e
503 else {
504 return false;
505 };
506 get_string_param(&metadata.parameters, GOAL_ID) == Some("goal-7")
507 });
508 assert!(has_goal, "action result with goal_id was wiped by flush");
509 }
510
511 #[test]
512 fn flush_with_all_correlated_queue_keeps_everything() {
513 let (mut sched, id) = make_scheduler(10);
516
517 sched.add_event(make_input("audio", with_request_id("req-1")));
518 sched.add_event(make_input("audio", with_request_id("req-2")));
519 sched.add_event(make_input("audio", with_request_id("req-3")));
520
521 let mut flush_params = MetadataParameters::new();
522 flush_params.insert(FLUSH.into(), Parameter::Bool(true));
523 sched.add_event(make_input("audio", flush_params));
524
525 assert_eq!(sched.event_queues[&id].1.len(), 4);
527 }
528
529 fn request_id_of(event: &EventItem) -> Option<String> {
533 let EventItem::NodeEvent {
534 event: NodeEvent::Input { metadata, .. },
535 ..
536 } = event
537 else {
538 return None;
539 };
540 get_string_param(&metadata.parameters, REQUEST_ID).map(|s| s.to_string())
541 }
542
543 fn with_request_id(id: &str) -> MetadataParameters {
544 let mut params = MetadataParameters::new();
545 params.insert(REQUEST_ID.into(), Parameter::String(id.to_string()));
546 params
547 }
548
549 #[test]
550 fn drop_oldest_preserves_correlated_when_non_correlated_present() {
551 let (mut sched, id) = make_scheduler(3);
554
555 sched.add_event(make_input("audio", with_request_id("req-1")));
556 sched.add_event(make_input("audio", MetadataParameters::new()));
557 sched.add_event(make_input("audio", MetadataParameters::new()));
558 sched.add_event(make_input("audio", MetadataParameters::new()));
559
560 let queue = &sched.event_queues[&id].1;
561 assert_eq!(queue.len(), 3);
562 assert!(
564 queue
565 .iter()
566 .any(|e| request_id_of(e).as_deref() == Some("req-1")),
567 "correlated message was dropped even though non-correlated events were available"
568 );
569 }
570
571 #[test]
572 fn drop_oldest_drops_middle_non_correlated_to_save_front_correlated() {
573 let (mut sched, id) = make_scheduler(3);
576
577 sched.add_event(make_input("audio", with_request_id("req-1")));
578 sched.add_event(make_input("audio", MetadataParameters::new()));
579 sched.add_event(make_input("audio", with_request_id("req-2")));
580 sched.add_event(make_input("audio", MetadataParameters::new()));
581
582 let queue = &sched.event_queues[&id].1;
583 assert_eq!(queue.len(), 3);
584 assert!(
585 queue
586 .iter()
587 .any(|e| request_id_of(e).as_deref() == Some("req-1"))
588 );
589 assert!(
590 queue
591 .iter()
592 .any(|e| request_id_of(e).as_deref() == Some("req-2"))
593 );
594 }
595
596 #[test]
597 fn drop_oldest_drops_incoming_if_queue_is_fully_correlated_and_incoming_is_not() {
598 let (mut sched, id) = make_scheduler(2);
601
602 sched.add_event(make_input("audio", with_request_id("req-1")));
603 sched.add_event(make_input("audio", with_request_id("req-2")));
604 sched.add_event(make_input("audio", MetadataParameters::new()));
605
606 let queue = &sched.event_queues[&id].1;
607 assert_eq!(queue.len(), 2);
608 let ids: Vec<_> = queue.iter().filter_map(request_id_of).collect();
609 assert_eq!(ids, vec!["req-1".to_string(), "req-2".to_string()]);
610
611 let counts = sched.drain_drop_counts();
613 assert_eq!(counts.get(&id), Some(&1));
614 }
615
616 #[test]
617 fn drop_oldest_drops_front_loudly_when_both_queue_and_incoming_are_correlated() {
618 let (mut sched, id) = make_scheduler(2);
621
622 sched.add_event(make_input("audio", with_request_id("req-1")));
623 sched.add_event(make_input("audio", with_request_id("req-2")));
624 sched.add_event(make_input("audio", with_request_id("req-3")));
625
626 let queue = &sched.event_queues[&id].1;
627 assert_eq!(queue.len(), 2);
628 let ids: Vec<_> = queue.iter().filter_map(request_id_of).collect();
629 assert_eq!(ids, vec!["req-2".to_string(), "req-3".to_string()]);
630 }
631
632 #[test]
633 fn drop_oldest_goal_id_is_also_preserved() {
634 let (mut sched, id) = make_scheduler(2);
636
637 let mut goal_params = MetadataParameters::new();
638 goal_params.insert(GOAL_ID.into(), Parameter::String("goal-42".to_string()));
639
640 sched.add_event(make_input("audio", goal_params));
641 sched.add_event(make_input("audio", MetadataParameters::new()));
642 sched.add_event(make_input("audio", MetadataParameters::new()));
643
644 let queue = &sched.event_queues[&id].1;
645 assert_eq!(queue.len(), 2);
646 let has_goal = queue.iter().any(|e| {
647 let EventItem::NodeEvent {
648 event: NodeEvent::Input { metadata, .. },
649 ..
650 } = e
651 else {
652 return false;
653 };
654 get_string_param(&metadata.parameters, GOAL_ID) == Some("goal-42")
655 });
656 assert!(
657 has_goal,
658 "goal-42 was dropped despite having non-correlated events to drop"
659 );
660 }
661
662 #[test]
663 fn backpressure_policy_prevents_drops() {
664 let id = DataId::from("commands".to_string());
665 let mut queues = HashMap::new();
666 queues.insert(id.clone(), (2, VecDeque::new()));
667 queues.insert(
668 DataId::from(NON_INPUT_EVENT.to_string()),
669 (10, VecDeque::new()),
670 );
671 let policies = HashMap::from([(id.clone(), QueuePolicy::Backpressure)]);
672 let mut sched = Scheduler::with_policies(queues, policies);
673
674 sched.add_event(make_input("commands", MetadataParameters::new()));
676 sched.add_event(make_input("commands", MetadataParameters::new()));
677 sched.add_event(make_input("commands", MetadataParameters::new()));
678 sched.add_event(make_input("commands", MetadataParameters::new()));
679
680 assert_eq!(sched.event_queues[&id].1.len(), 4);
682
683 let counts = sched.drain_drop_counts();
685 assert!(counts.is_empty());
686 }
687
688 fn make_zenoh_input(id: &str, params: MetadataParameters) -> EventItem {
691 let ts = uhlc::HLC::default().new_timestamp();
692 let metadata = Metadata::from_parameters(ts, params);
693 use dora_arrow_convert::IntoArrow;
694 EventItem::ZenohInput {
695 id: DataId::from(id.to_string()),
696 metadata: std::sync::Arc::new(metadata),
697 data: ().into_arrow().into(),
698 }
699 }
700
701 fn request_id_of_zenoh(event: &EventItem) -> Option<String> {
703 let EventItem::ZenohInput { metadata, .. } = event else {
704 return None;
705 };
706 get_string_param(&metadata.parameters, REQUEST_ID).map(|s| s.to_string())
707 }
708
709 #[test]
710 fn zenoh_drop_oldest_drops_front_loudly_when_both_queue_and_incoming_are_correlated() {
711 let (mut sched, id) = make_scheduler(2);
718
719 sched.add_event(make_zenoh_input("audio", with_request_id("req-1")));
720 sched.add_event(make_zenoh_input("audio", with_request_id("req-2")));
721 sched.add_event(make_zenoh_input("audio", with_request_id("req-3")));
722
723 let queue = &sched.event_queues[&id].1;
724 assert_eq!(queue.len(), 2);
725 let ids: Vec<_> = queue.iter().filter_map(request_id_of_zenoh).collect();
726 assert_eq!(ids, vec!["req-2".to_string(), "req-3".to_string()]);
727 }
728
729 #[test]
730 fn zenoh_drop_oldest_preserves_correlated_when_non_correlated_present() {
731 let (mut sched, id) = make_scheduler(3);
734
735 sched.add_event(make_zenoh_input("audio", with_request_id("req-1")));
736 sched.add_event(make_zenoh_input("audio", MetadataParameters::new()));
737 sched.add_event(make_zenoh_input("audio", MetadataParameters::new()));
738 sched.add_event(make_zenoh_input("audio", MetadataParameters::new()));
739
740 let queue = &sched.event_queues[&id].1;
741 assert_eq!(queue.len(), 3);
742 assert!(
743 queue
744 .iter()
745 .any(|e| request_id_of_zenoh(e).as_deref() == Some("req-1")),
746 "correlated ZenohInput was dropped even though non-correlated events were available"
747 );
748 }
749}