dynamo_mocker/scheduler/
source_holds.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::VecDeque;
5
6use anyhow::{Result, bail};
7use rustc_hash::FxHashMap;
8use uuid::Uuid;
9
10use crate::common::handoff::{HandoffId, HandoffTransferTiming};
11use crate::common::protocols::DirectRequest;
12
13#[allow(dead_code)]
14pub enum SchedulerCommand {
15    Submit(DirectRequest),
16    /// Remove an ordinary request from the live scheduler by its stable ID.
17    ///
18    /// Live requests use a dedicated cancellation lane. Their owned submit
19    /// task must acknowledge admission before cancellation is enqueued, so a
20    /// dropped network stream cannot race ahead of its own request admission.
21    CancelRequest {
22        request_id: Uuid,
23    },
24    SubmitHandoffPrefill {
25        handoff_id: HandoffId,
26        request: DirectRequest,
27    },
28    ReleaseSource {
29        handoff_id: HandoffId,
30    },
31    CancelSource {
32        handoff_id: HandoffId,
33    },
34    ReserveDestination {
35        handoff_id: HandoffId,
36        request: DirectRequest,
37    },
38    ActivateDestination {
39        handoff_id: HandoffId,
40    },
41    CancelDestination {
42        handoff_id: HandoffId,
43    },
44}
45
46#[derive(Clone, Copy, Debug, Eq, PartialEq)]
47pub enum SchedulerCommandResult {
48    Submitted(Uuid),
49    DestinationAccepted { request_id: Uuid },
50    Applied,
51    Noop,
52}
53
54#[derive(Clone, Copy, Debug, PartialEq)]
55pub enum SchedulerLifecycleEvent {
56    SourceHeld {
57        handoff_id: HandoffId,
58        request_id: Uuid,
59        transfer_timing: HandoffTransferTiming,
60    },
61    DestinationReserved {
62        handoff_id: HandoffId,
63        request_id: Uuid,
64        transferable_prompt_tokens: usize,
65    },
66}
67
68impl SchedulerLifecycleEvent {
69    pub fn handoff_id(&self) -> HandoffId {
70        match *self {
71            Self::SourceHeld { handoff_id, .. } | Self::DestinationReserved { handoff_id, .. } => {
72                handoff_id
73            }
74        }
75    }
76}
77
78#[derive(Debug)]
79pub struct SchedulerCommandEffects {
80    pub result: SchedulerCommandResult,
81    pub lifecycle_events: Vec<SchedulerLifecycleEvent>,
82    pub kv_events: Vec<dynamo_kv_router::protocols::RouterEvent>,
83}
84
85impl SchedulerCommandEffects {
86    pub(crate) fn new(result: SchedulerCommandResult) -> Self {
87        Self {
88            result,
89            lifecycle_events: Vec::new(),
90            kv_events: Vec::new(),
91        }
92    }
93}
94
95pub(crate) enum SourceCompletion<T> {
96    Release(T),
97    Held { handoff_id: HandoffId },
98}
99
100pub(crate) enum RemovedSource<T> {
101    Held(T),
102    Pending { request_id: Uuid },
103    Missing,
104}
105
106pub(crate) struct SourceHolds<T> {
107    pending_by_request: FxHashMap<Uuid, HandoffId>,
108    pending_by_handoff: FxHashMap<HandoffId, Uuid>,
109    held_prefills: FxHashMap<HandoffId, (Uuid, T)>,
110    held_by_request: FxHashMap<Uuid, HandoffId>,
111}
112
113impl<T> Default for SourceHolds<T> {
114    fn default() -> Self {
115        Self {
116            pending_by_request: FxHashMap::default(),
117            pending_by_handoff: FxHashMap::default(),
118            held_prefills: FxHashMap::default(),
119            held_by_request: FxHashMap::default(),
120        }
121    }
122}
123
124impl<T> SourceHolds<T> {
125    #[allow(dead_code)]
126    pub(crate) fn is_empty(&self) -> bool {
127        debug_assert_eq!(self.pending_by_request.len(), self.pending_by_handoff.len());
128        debug_assert_eq!(self.held_prefills.len(), self.held_by_request.len());
129        self.pending_by_request.is_empty()
130            && self.held_prefills.is_empty()
131            && self.held_by_request.is_empty()
132    }
133
134    pub(crate) fn register(&mut self, request_id: Uuid, handoff_id: HandoffId) -> Result<()> {
135        if self.contains_request(request_id) {
136            bail!("source hold already active for request {request_id}");
137        }
138        if self.pending_by_handoff.contains_key(&handoff_id)
139            || self.held_prefills.contains_key(&handoff_id)
140        {
141            bail!("handoff {handoff_id:?} is already active");
142        }
143
144        self.pending_by_request.insert(request_id, handoff_id);
145        self.pending_by_handoff.insert(handoff_id, request_id);
146        Ok(())
147    }
148
149    pub(crate) fn contains_request(&self, request_id: Uuid) -> bool {
150        self.pending_by_request.contains_key(&request_id)
151            || self.held_by_request.contains_key(&request_id)
152    }
153
154    pub(crate) fn complete_source(&mut self, request_id: Uuid, payload: T) -> SourceCompletion<T> {
155        let Some(handoff_id) = self.pending_by_request.remove(&request_id) else {
156            return SourceCompletion::Release(payload);
157        };
158
159        let registered_request = self
160            .pending_by_handoff
161            .remove(&handoff_id)
162            .expect("validated source-hold registration must be bidirectional");
163        debug_assert_eq!(registered_request, request_id);
164
165        let previous = self.held_prefills.insert(handoff_id, (request_id, payload));
166        debug_assert!(previous.is_none());
167        let previous = self.held_by_request.insert(request_id, handoff_id);
168        debug_assert!(previous.is_none());
169        SourceCompletion::Held { handoff_id }
170    }
171
172    pub(crate) fn remove(&mut self, handoff_id: HandoffId) -> RemovedSource<T> {
173        if let Some((request_id, payload)) = self.held_prefills.remove(&handoff_id) {
174            let removed = self.held_by_request.remove(&request_id);
175            debug_assert_eq!(removed, Some(handoff_id));
176            return RemovedSource::Held(payload);
177        }
178
179        let Some(request_id) = self.pending_by_handoff.remove(&handoff_id) else {
180            return RemovedSource::Missing;
181        };
182        let removed = self.pending_by_request.remove(&request_id);
183        debug_assert_eq!(removed, Some(handoff_id));
184        RemovedSource::Pending { request_id }
185    }
186
187    pub(crate) fn remove_request(&mut self, request_id: Uuid) {
188        let Some(handoff_id) = self.pending_by_request.remove(&request_id) else {
189            return;
190        };
191        let removed = self.pending_by_handoff.remove(&handoff_id);
192        debug_assert_eq!(removed, Some(request_id));
193    }
194
195    #[cfg(test)]
196    pub(crate) fn is_held(&self, handoff_id: HandoffId) -> bool {
197        self.held_prefills.contains_key(&handoff_id)
198    }
199
200    #[cfg(test)]
201    pub(crate) fn is_registered(&self, handoff_id: HandoffId) -> bool {
202        self.pending_by_handoff.contains_key(&handoff_id)
203    }
204}
205
206pub(crate) struct DestinationHolds<T> {
207    by_handoff: FxHashMap<HandoffId, (Uuid, T)>,
208    by_request: FxHashMap<Uuid, HandoffId>,
209}
210
211#[derive(Default)]
212pub(crate) struct ActiveHandoffRequests {
213    by_handoff: FxHashMap<HandoffId, Uuid>,
214    by_request: FxHashMap<Uuid, HandoffId>,
215}
216
217impl ActiveHandoffRequests {
218    pub(crate) fn insert(&mut self, handoff_id: HandoffId, request_id: Uuid) {
219        let previous = self.by_handoff.insert(handoff_id, request_id);
220        debug_assert!(previous.is_none());
221        let previous = self.by_request.insert(request_id, handoff_id);
222        debug_assert!(previous.is_none());
223    }
224
225    pub(crate) fn remove_handoff(&mut self, handoff_id: HandoffId) -> Option<Uuid> {
226        let request_id = self.by_handoff.remove(&handoff_id)?;
227        let removed = self.by_request.remove(&request_id);
228        debug_assert_eq!(removed, Some(handoff_id));
229        Some(request_id)
230    }
231
232    pub(crate) fn remove_request(&mut self, request_id: Uuid) -> Option<HandoffId> {
233        let handoff_id = self.by_request.remove(&request_id)?;
234        let removed = self.by_handoff.remove(&handoff_id);
235        debug_assert_eq!(removed, Some(request_id));
236        Some(handoff_id)
237    }
238
239    pub(crate) fn contains_request(&self, request_id: Uuid) -> bool {
240        self.by_request.contains_key(&request_id)
241    }
242
243    pub(crate) fn contains_handoff(&self, handoff_id: HandoffId) -> bool {
244        self.by_handoff.contains_key(&handoff_id)
245    }
246
247    pub(crate) fn is_empty(&self) -> bool {
248        debug_assert_eq!(self.by_handoff.len(), self.by_request.len());
249        self.by_handoff.is_empty()
250    }
251}
252
253pub(crate) struct PendingDestinations<T> {
254    fifo: VecDeque<HandoffId>,
255    by_handoff: FxHashMap<HandoffId, PendingDestination<T>>,
256    by_request: FxHashMap<Uuid, HandoffId>,
257}
258
259struct PendingDestination<T> {
260    request_id: Uuid,
261    payload: T,
262    last_attempt_generation: Option<u64>,
263}
264
265impl<T> Default for PendingDestinations<T> {
266    fn default() -> Self {
267        Self {
268            fifo: VecDeque::new(),
269            by_handoff: FxHashMap::default(),
270            by_request: FxHashMap::default(),
271        }
272    }
273}
274
275impl<T> PendingDestinations<T> {
276    pub(crate) fn validate(&self, request_id: Uuid, handoff_id: HandoffId) -> Result<()> {
277        if self.by_request.contains_key(&request_id) {
278            bail!("destination request {request_id} is already pending");
279        }
280        if self.by_handoff.contains_key(&handoff_id) {
281            bail!("destination handoff {handoff_id:?} is already pending");
282        }
283        Ok(())
284    }
285
286    pub(crate) fn insert(&mut self, request_id: Uuid, handoff_id: HandoffId, payload: T) {
287        debug_assert!(self.validate(request_id, handoff_id).is_ok());
288        self.fifo.push_back(handoff_id);
289        let previous = self.by_handoff.insert(
290            handoff_id,
291            PendingDestination {
292                request_id,
293                payload,
294                last_attempt_generation: None,
295            },
296        );
297        debug_assert!(previous.is_none());
298        let previous = self.by_request.insert(request_id, handoff_id);
299        debug_assert!(previous.is_none());
300    }
301
302    pub(crate) fn front_due(&mut self, generation: u64) -> Option<(HandoffId, Uuid, &T)> {
303        self.normalize_front();
304        let handoff_id = *self.fifo.front()?;
305        let pending = self
306            .by_handoff
307            .get(&handoff_id)
308            .expect("normalized pending destination must exist");
309        if pending.last_attempt_generation == Some(generation) {
310            return None;
311        }
312        Some((handoff_id, pending.request_id, &pending.payload))
313    }
314
315    pub(crate) fn front_due_mut(&mut self, generation: u64) -> Option<(HandoffId, Uuid, &mut T)> {
316        self.normalize_front();
317        let handoff_id = *self.fifo.front()?;
318        let pending = self
319            .by_handoff
320            .get_mut(&handoff_id)
321            .expect("normalized pending destination must exist");
322        if pending.last_attempt_generation == Some(generation) {
323            return None;
324        }
325        Some((handoff_id, pending.request_id, &mut pending.payload))
326    }
327
328    pub(crate) fn mark_front_attempted(&mut self, generation: u64) {
329        self.normalize_front();
330        let Some(handoff_id) = self.fifo.front() else {
331            return;
332        };
333        let pending = self
334            .by_handoff
335            .get_mut(handoff_id)
336            .expect("normalized pending destination must exist");
337        pending.last_attempt_generation = Some(generation);
338    }
339
340    pub(crate) fn pop_front(&mut self) -> Option<(HandoffId, Uuid, T)> {
341        self.normalize_front();
342        let handoff_id = self.fifo.pop_front()?;
343        let pending = self
344            .by_handoff
345            .remove(&handoff_id)
346            .expect("normalized pending destination must exist");
347        let removed = self.by_request.remove(&pending.request_id);
348        debug_assert_eq!(removed, Some(handoff_id));
349        self.compact_if_sparse();
350        Some((handoff_id, pending.request_id, pending.payload))
351    }
352
353    pub(crate) fn remove(&mut self, handoff_id: HandoffId) -> Option<(Uuid, T)> {
354        let pending = self.by_handoff.remove(&handoff_id)?;
355        let removed = self.by_request.remove(&pending.request_id);
356        debug_assert_eq!(removed, Some(handoff_id));
357        self.compact_if_sparse();
358        Some((pending.request_id, pending.payload))
359    }
360
361    pub(crate) fn contains_request(&self, request_id: Uuid) -> bool {
362        self.by_request.contains_key(&request_id)
363    }
364
365    #[cfg(test)]
366    pub(crate) fn contains_handoff(&self, handoff_id: HandoffId) -> bool {
367        self.by_handoff.contains_key(&handoff_id)
368    }
369
370    pub(crate) fn is_empty(&self) -> bool {
371        debug_assert_eq!(self.by_handoff.len(), self.by_request.len());
372        self.by_handoff.is_empty()
373    }
374
375    pub(crate) fn has_pending(&self) -> bool {
376        !self.by_handoff.is_empty()
377    }
378
379    pub(crate) fn len(&self) -> usize {
380        self.by_handoff.len()
381    }
382
383    pub(crate) fn payloads(&self) -> impl Iterator<Item = &T> {
384        self.by_handoff.values().map(|pending| &pending.payload)
385    }
386
387    fn normalize_front(&mut self) {
388        while self
389            .fifo
390            .front()
391            .is_some_and(|handoff_id| !self.by_handoff.contains_key(handoff_id))
392        {
393            self.fifo.pop_front();
394        }
395    }
396
397    fn compact_if_sparse(&mut self) {
398        if self.fifo.len() <= self.by_handoff.len().saturating_mul(2) {
399            return;
400        }
401        self.fifo
402            .retain(|handoff_id| self.by_handoff.contains_key(handoff_id));
403    }
404}
405
406impl<T> Default for DestinationHolds<T> {
407    fn default() -> Self {
408        Self {
409            by_handoff: FxHashMap::default(),
410            by_request: FxHashMap::default(),
411        }
412    }
413}
414
415impl<T> DestinationHolds<T> {
416    pub(crate) fn validate(&self, request_id: Uuid, handoff_id: HandoffId) -> Result<()> {
417        if self.by_request.contains_key(&request_id) {
418            bail!("destination reservation already exists for request {request_id}");
419        }
420        if self.by_handoff.contains_key(&handoff_id) {
421            bail!("destination handoff {handoff_id:?} is already active");
422        }
423        Ok(())
424    }
425
426    pub(crate) fn contains_request(&self, request_id: Uuid) -> bool {
427        self.by_request.contains_key(&request_id)
428    }
429
430    pub(crate) fn insert(&mut self, request_id: Uuid, handoff_id: HandoffId, payload: T) {
431        debug_assert!(self.validate(request_id, handoff_id).is_ok());
432        let previous = self.by_handoff.insert(handoff_id, (request_id, payload));
433        debug_assert!(previous.is_none());
434        let previous = self.by_request.insert(request_id, handoff_id);
435        debug_assert!(previous.is_none());
436    }
437
438    pub(crate) fn remove(&mut self, handoff_id: HandoffId) -> Option<(Uuid, T)> {
439        let (request_id, payload) = self.by_handoff.remove(&handoff_id)?;
440        let removed = self.by_request.remove(&request_id);
441        debug_assert_eq!(removed, Some(handoff_id));
442        Some((request_id, payload))
443    }
444
445    #[allow(dead_code)]
446    pub(crate) fn is_empty(&self) -> bool {
447        debug_assert_eq!(self.by_handoff.len(), self.by_request.len());
448        self.by_handoff.is_empty()
449    }
450
451    pub(crate) fn len(&self) -> usize {
452        self.by_handoff.len()
453    }
454
455    pub(crate) fn payloads(&self) -> impl Iterator<Item = &T> {
456        self.by_handoff.values().map(|(_, payload)| payload)
457    }
458
459    #[cfg(test)]
460    pub(crate) fn contains(&self, handoff_id: HandoffId) -> bool {
461        self.by_handoff.contains_key(&handoff_id)
462    }
463
464    #[cfg(test)]
465    pub(crate) fn get(&self, handoff_id: HandoffId) -> Option<&T> {
466        self.by_handoff.get(&handoff_id).map(|(_, payload)| payload)
467    }
468}
469
470#[cfg(test)]
471mod tests {
472    use super::*;
473
474    #[test]
475    fn completion_atomically_moves_registered_payload_to_held() {
476        let request_id = Uuid::from_u128(1);
477        let handoff_id = HandoffId::from(Uuid::from_u128(2));
478        let mut holds = SourceHolds::default();
479
480        holds.register(request_id, handoff_id).unwrap();
481        assert!(holds.is_registered(handoff_id));
482
483        let completion = holds.complete_source(request_id, "payload");
484        assert!(matches!(
485            completion,
486            SourceCompletion::Held { handoff_id: held } if held == handoff_id
487        ));
488        assert!(!holds.is_registered(handoff_id));
489        assert!(holds.is_held(handoff_id));
490        assert!(
491            holds
492                .register(request_id, HandoffId::from(Uuid::from_u128(7)))
493                .is_err()
494        );
495        assert!(matches!(
496            holds.remove(handoff_id),
497            RemovedSource::Held("payload")
498        ));
499        assert!(matches!(holds.remove(handoff_id), RemovedSource::Missing));
500    }
501
502    #[test]
503    fn active_ids_are_rejected_but_released_ids_can_be_reused() {
504        let request_id = Uuid::from_u128(3);
505        let handoff_id = HandoffId::from(Uuid::from_u128(4));
506        let mut holds = SourceHolds::<()>::default();
507
508        holds.register(request_id, handoff_id).unwrap();
509        assert!(holds.register(Uuid::from_u128(5), handoff_id).is_err());
510        assert!(holds.register(request_id, HandoffId::new()).is_err());
511        assert!(matches!(
512            holds.remove(handoff_id),
513            RemovedSource::Pending { request_id: pending } if pending == request_id
514        ));
515
516        holds.register(Uuid::from_u128(5), handoff_id).unwrap();
517    }
518
519    #[test]
520    fn unregistered_completion_releases_payload() {
521        let mut holds = SourceHolds::default();
522        let payload = String::from("payload");
523
524        assert!(matches!(
525            holds.complete_source(Uuid::from_u128(6), payload),
526            SourceCompletion::Release(value) if value == "payload"
527        ));
528    }
529
530    #[test]
531    fn pending_destination_preserves_fifo_after_head_cancellation() {
532        let first = HandoffId::from(Uuid::from_u128(10));
533        let second = HandoffId::from(Uuid::from_u128(11));
534        let mut pending = PendingDestinations::default();
535        pending.insert(Uuid::from_u128(20), first, "first");
536        pending.insert(Uuid::from_u128(21), second, "second");
537
538        assert_eq!(pending.remove(first), Some((Uuid::from_u128(20), "first")));
539        let (handoff_id, request_id, payload) = pending.front_due(0).unwrap();
540        assert_eq!(handoff_id, second);
541        assert_eq!(request_id, Uuid::from_u128(21));
542        assert_eq!(*payload, "second");
543
544        pending.mark_front_attempted(0);
545        assert!(pending.front_due(0).is_none());
546        assert!(pending.front_due(1).is_some());
547    }
548
549    #[test]
550    fn pending_destination_indexes_reject_conflicts_and_compact_tombstones() {
551        let mut pending = PendingDestinations::default();
552        let handoffs = (0..8)
553            .map(|index| HandoffId::from(Uuid::from_u128(100 + index)))
554            .collect::<Vec<_>>();
555        for (index, handoff_id) in handoffs.iter().copied().enumerate() {
556            pending.insert(Uuid::from_u128(200 + index as u128), handoff_id, index);
557        }
558        assert!(
559            pending
560                .validate(Uuid::from_u128(200), HandoffId::new())
561                .is_err()
562        );
563        assert!(pending.validate(Uuid::new_v4(), handoffs[0]).is_err());
564
565        for handoff_id in handoffs.iter().take(6) {
566            assert!(pending.remove(*handoff_id).is_some());
567        }
568        assert!(pending.fifo.len() <= pending.by_handoff.len() * 2);
569        let (handoff_id, request_id, payload) = pending.front_due(0).unwrap();
570        assert_eq!(handoff_id, handoffs[6]);
571        assert_eq!(request_id, Uuid::from_u128(206));
572        assert_eq!(*payload, 6);
573        assert_eq!(pending.pop_front(), Some((handoffs[6], request_id, 6)));
574        assert_eq!(pending.pop_front().map(|entry| entry.0), Some(handoffs[7]));
575        assert!(pending.is_empty());
576    }
577}
dynamo_mocker/scheduler/source_holds.rs

dynamo_mocker/scheduler/
source_holds.rs
