1use crate::consensus::{
4 resolve_completions, validate_disposition, validate_schedule, CompletionObservation,
5 CompletionResolution, ConsensusTransport, ScheduledWork,
6};
7use serde::{Deserialize, Serialize};
8use std::{
9 collections::{BTreeMap, BTreeSet, VecDeque},
10 time::{Duration, Instant},
11};
12
13#[derive(Debug, Clone, Copy, Eq, Ord, PartialEq, PartialOrd, Hash, Serialize, Deserialize)]
15pub struct RequestId(u64);
16
17impl RequestId {
18 pub const fn new(value: u64) -> Self {
20 Self(value)
21 }
22
23 pub const fn value(self) -> u64 {
25 self.0
26 }
27}
28
29#[derive(Debug, Clone, Copy, Eq, Ord, PartialEq, PartialOrd, Hash, Serialize, Deserialize)]
31pub struct WorkId {
32 request: RequestId,
33 sequence: u64,
34}
35
36impl WorkId {
37 pub const fn new(request: RequestId, sequence: u64) -> Self {
39 Self { request, sequence }
40 }
41
42 pub const fn request(self) -> RequestId {
44 self.request
45 }
46
47 pub const fn sequence(self) -> u64 {
49 self.sequence
50 }
51}
52
53#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
55#[serde(rename_all = "snake_case")]
56pub enum WorkLifecycle {
57 Queued,
59 Prepared,
61 Submitted,
63 Completing,
65 Committed,
67 Abandoned,
69 Failed,
71}
72
73#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
75#[serde(rename_all = "snake_case")]
76pub enum RequestStatus {
77 Active,
79 Finished,
81 Cancelled,
83 DeadlineExceeded,
85 Failed,
87}
88
89#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
91#[serde(rename_all = "snake_case")]
92pub enum CancellationCause {
93 Explicit,
95 Deadline,
97}
98
99#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
101pub struct SchedulerLimits {
102 pub max_active_requests: usize,
104 pub max_queued_work: usize,
106 pub max_new_submissions_per_turn: usize,
108 pub max_in_flight_global: usize,
110 pub max_in_flight_per_request: usize,
112 pub execution_slice: usize,
114}
115
116impl SchedulerLimits {
117 pub fn new(max_active_requests: usize, max_queued_work: usize) -> Result<Self, SchedulerError> {
119 Self::with_execution_bounds(
120 max_active_requests,
121 max_queued_work,
122 1,
123 max_active_requests,
124 1,
125 usize::MAX,
126 )
127 }
128
129 pub fn with_execution_bounds(
131 max_active_requests: usize,
132 max_queued_work: usize,
133 max_new_submissions_per_turn: usize,
134 max_in_flight_global: usize,
135 max_in_flight_per_request: usize,
136 execution_slice: usize,
137 ) -> Result<Self, SchedulerError> {
138 let values = [
139 max_active_requests,
140 max_queued_work,
141 max_new_submissions_per_turn,
142 max_in_flight_global,
143 max_in_flight_per_request,
144 execution_slice,
145 ];
146 if values.contains(&0) {
147 return Err(SchedulerError::InvalidLimits(values));
148 }
149 Ok(Self {
150 max_active_requests,
151 max_queued_work,
152 max_new_submissions_per_turn,
153 max_in_flight_global,
154 max_in_flight_per_request,
155 execution_slice,
156 })
157 }
158}
159
160impl Default for SchedulerLimits {
161 fn default() -> Self {
162 Self {
163 max_active_requests: 64,
164 max_queued_work: 256,
165 max_new_submissions_per_turn: 1,
166 max_in_flight_global: 64,
167 max_in_flight_per_request: 1,
168 execution_slice: usize::MAX,
169 }
170 }
171}
172
173pub trait WorkDescriptor {
175 type Error: std::error::Error;
177
178 fn encode_descriptor(&self, output: &mut Vec<u32>) -> Result<(), Self::Error>;
180
181 fn execution_slice_size(&self) -> usize {
183 1
184 }
185}
186
187pub trait SemanticStateTransaction {
189 type Branch;
191 type Error: std::error::Error;
193
194 fn branch(&self) -> Result<Self::Branch, Self::Error>;
196
197 fn commit_branch(&mut self, branch: Self::Branch) -> Result<(), Self::Error>;
199
200 fn discard_branch(branch: Self::Branch) -> Result<(), Self::Error> {
202 drop(branch);
203 Ok(())
204 }
205
206 fn permits_parallel_branches(&self) -> bool {
208 false
209 }
210}
211
212pub trait TransitionOutput {
214 type Error: std::error::Error;
216
217 fn is_complete(&self) -> Result<bool, Self::Error>;
219
220 fn backend_name(&self) -> Option<String> {
222 None
223 }
224
225 fn physically_preemptible(&self) -> bool {
227 false
228 }
229
230 fn retained_resources(&self) -> usize;
232}
233
234#[derive(Debug)]
235struct Request<W, S> {
236 state: S,
237 next: u64,
238 pending: VecDeque<Queued<W>>,
239}
240
241#[derive(Debug)]
242struct Queued<W> {
243 id: WorkId,
244 work: W,
245 deadline: Option<Instant>,
246}
247
248#[derive(Debug)]
249struct Prepared<W, B> {
250 id: WorkId,
251 work: W,
252 descriptor: Vec<u32>,
253 branch: B,
254 deadline: Option<Instant>,
255}
256
257#[derive(Debug, Clone, Copy)]
258enum Disposition {
259 Publish,
260 Abandon { cancelled_at: Instant },
261 Fail,
262}
263
264#[derive(Debug)]
265struct Submitted<W, B, O> {
266 id: WorkId,
267 work: W,
268 branch: B,
269 output: O,
270 disposition: Disposition,
271}
272
273#[derive(Debug)]
275pub struct SchedulerProgress<W, O> {
276 pub newly_submitted: usize,
278 pub committed: Vec<(WorkId, W, O)>,
280 pub failed: Vec<(WorkId, SchedulerError)>,
282}
283
284impl<W, O> Default for SchedulerProgress<W, O> {
285 fn default() -> Self {
286 Self {
287 newly_submitted: 0,
288 committed: Vec::new(),
289 failed: Vec::new(),
290 }
291 }
292}
293
294#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
296pub struct SchedulerCapabilities {
297 pub limits: SchedulerLimits,
299 pub observed_backends: Vec<String>,
301 pub executing_work_physically_preemptible: bool,
303 pub non_preemptible_interval: String,
305}
306
307#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Serialize, Deserialize)]
309pub struct SchedulerReport {
310 pub active_requests: usize,
312 pub queued_work: usize,
314 pub prepared_work: usize,
316 pub submitted_in_flight_work: usize,
318 pub completing_work: usize,
320 pub abandoned_in_flight_work: usize,
322 pub failed_in_flight_work: usize,
324 pub current_in_flight_work: usize,
326 pub peak_in_flight_work: usize,
328 pub peak_queued_work: usize,
330 pub submitted_work: u64,
332 pub completed_work: u64,
334 pub failed_work: u64,
336 pub discarded_work: u64,
338 pub cancellation_before_submission: u64,
340 pub cancellation_after_submission: u64,
342 pub abandoned_released_work: u64,
344 pub abandoned_retained_resources: usize,
346 pub peak_abandoned_retained_resources: usize,
348 pub last_cancellation_to_release_ns: Option<u128>,
350 pub max_cancellation_to_release_ns: Option<u128>,
352 pub finished_requests: u64,
354 pub cancelled_requests: u64,
356 pub deadline_expired_requests: u64,
358 pub drain_cycles: u64,
360 pub configured_submission_bound: usize,
362 pub configured_slice_bound: usize,
364 pub poisoned: bool,
366}
367
368#[derive(Debug)]
370pub struct Scheduler<W, S: SemanticStateTransaction, O: TransitionOutput> {
371 limits: SchedulerLimits,
372 requests: BTreeMap<RequestId, Request<W, S>>,
373 terminal: BTreeMap<RequestId, RequestStatus>,
374 ready: VecDeque<RequestId>,
375 prepared: VecDeque<Prepared<W, S::Branch>>,
376 submitted: Vec<Submitted<W, S::Branch, O>>,
377 lifecycle: BTreeMap<WorkId, WorkLifecycle>,
378 accepted_work: usize,
379 peak_accepted_work: usize,
380 peak_in_flight_work: usize,
381 submitted_work: u64,
382 completed_work: u64,
383 failed_work: u64,
384 discarded_work: u64,
385 cancellation_before_submission: u64,
386 cancellation_after_submission: u64,
387 abandoned_released_work: u64,
388 peak_abandoned_retained_resources: usize,
389 last_cancellation_to_release: Option<Duration>,
390 max_cancellation_to_release: Option<Duration>,
391 finished_requests: u64,
392 cancelled_requests: u64,
393 deadline_expired_requests: u64,
394 drain_cycles: u64,
395 observed_backends: BTreeSet<String>,
396 all_outputs_preemptible: bool,
397 poisoned: Option<String>,
398}
399
400impl<W, S: SemanticStateTransaction, O: TransitionOutput> Scheduler<W, S, O> {
401 pub fn new(limits: SchedulerLimits) -> Result<Self, SchedulerError> {
403 let limits = SchedulerLimits::with_execution_bounds(
404 limits.max_active_requests,
405 limits.max_queued_work,
406 limits.max_new_submissions_per_turn,
407 limits.max_in_flight_global,
408 limits.max_in_flight_per_request,
409 limits.execution_slice,
410 )?;
411 Ok(Self {
412 limits,
413 requests: BTreeMap::new(),
414 terminal: BTreeMap::new(),
415 ready: VecDeque::new(),
416 prepared: VecDeque::new(),
417 submitted: Vec::new(),
418 lifecycle: BTreeMap::new(),
419 accepted_work: 0,
420 peak_accepted_work: 0,
421 peak_in_flight_work: 0,
422 submitted_work: 0,
423 completed_work: 0,
424 failed_work: 0,
425 discarded_work: 0,
426 cancellation_before_submission: 0,
427 cancellation_after_submission: 0,
428 abandoned_released_work: 0,
429 peak_abandoned_retained_resources: 0,
430 last_cancellation_to_release: None,
431 max_cancellation_to_release: None,
432 finished_requests: 0,
433 cancelled_requests: 0,
434 deadline_expired_requests: 0,
435 drain_cycles: 0,
436 observed_backends: BTreeSet::new(),
437 all_outputs_preemptible: true,
438 poisoned: None,
439 })
440 }
441
442 pub fn validate_registration(&self, id: RequestId) -> Result<(), SchedulerError> {
444 self.ensure_ready()?;
445 if self.requests.contains_key(&id) || self.terminal.contains_key(&id) {
446 return Err(SchedulerError::DuplicateRequest(id));
447 }
448 if self.requests.len() >= self.limits.max_active_requests {
449 return Err(SchedulerError::Capacity(format!(
450 "scheduler active-request capacity {} is exhausted",
451 self.limits.max_active_requests
452 )));
453 }
454 Ok(())
455 }
456
457 pub fn register(&mut self, id: RequestId, state: S) -> Result<(), SchedulerError> {
459 self.validate_registration(id)?;
460 self.requests.insert(
461 id,
462 Request {
463 state,
464 next: 0,
465 pending: VecDeque::new(),
466 },
467 );
468 Ok(())
469 }
470
471 pub fn request_state(&self, id: RequestId) -> Option<&S> {
473 self.requests.get(&id).map(|entry| &entry.state)
474 }
475
476 pub fn request_state_mut(&mut self, id: RequestId) -> Result<&mut S, SchedulerError> {
478 self.ensure_ready()?;
479 if self.branch_count(id) != 0 {
480 return Err(SchedulerError::State(format!(
481 "request {} has prepared or submitted state branches",
482 id.value()
483 )));
484 }
485 self.requests
486 .get_mut(&id)
487 .map(|entry| &mut entry.state)
488 .ok_or(SchedulerError::UnknownRequest(id))
489 }
490
491 pub fn enqueue(&mut self, request: RequestId, work: W) -> Result<WorkId, SchedulerError> {
493 self.enqueue_with_deadline(request, work, None)
494 }
495
496 pub fn enqueue_with_deadline(
498 &mut self,
499 request: RequestId,
500 work: W,
501 deadline: Option<Instant>,
502 ) -> Result<WorkId, SchedulerError> {
503 Ok(self
504 .enqueue_batch_with_deadlines(request, vec![(work, deadline)])?
505 .pop()
506 .expect("one work item was supplied"))
507 }
508
509 pub fn enqueue_batch(
511 &mut self,
512 request: RequestId,
513 work: Vec<W>,
514 ) -> Result<Vec<WorkId>, SchedulerError> {
515 self.enqueue_batch_with_deadlines(
516 request,
517 work.into_iter().map(|work| (work, None)).collect(),
518 )
519 }
520
521 fn enqueue_batch_with_deadlines(
522 &mut self,
523 request: RequestId,
524 work: Vec<(W, Option<Instant>)>,
525 ) -> Result<Vec<WorkId>, SchedulerError> {
526 self.ensure_ready()?;
527 let requested = work.len();
528 let accepted_after = self
529 .accepted_work
530 .checked_add(requested)
531 .ok_or_else(|| SchedulerError::Capacity("scheduler occupancy overflow".into()))?;
532 if accepted_after > self.limits.max_queued_work {
533 return Err(SchedulerError::Capacity(format!(
534 "scheduler queue capacity {} cannot accept {requested} items with {} outstanding",
535 self.limits.max_queued_work, self.accepted_work
536 )));
537 }
538 let entry = self
539 .requests
540 .get_mut(&request)
541 .ok_or(SchedulerError::UnknownRequest(request))?;
542 let count = u64::try_from(requested)
543 .map_err(|_| SchedulerError::Capacity("work batch length exceeds u64".into()))?;
544 let next = entry
545 .next
546 .checked_add(count)
547 .ok_or_else(|| SchedulerError::Capacity("work identity space exhausted".into()))?;
548 let was_empty = entry.pending.is_empty();
549 let mut ids = Vec::with_capacity(requested);
550 for (offset, (work, deadline)) in work.into_iter().enumerate() {
551 let id = WorkId::new(request, entry.next + offset as u64);
552 entry.pending.push_back(Queued { id, work, deadline });
553 self.lifecycle.insert(id, WorkLifecycle::Queued);
554 ids.push(id);
555 }
556 entry.next = next;
557 if was_empty && requested != 0 {
558 self.ready.push_back(request);
559 }
560 self.accepted_work = accepted_after;
561 self.peak_accepted_work = self.peak_accepted_work.max(accepted_after);
562 self.submitted_work = self.submitted_work.saturating_add(count);
563 Ok(ids)
564 }
565
566 pub fn prepare_bounded(&mut self, limit: usize, now: Instant) -> Result<usize, SchedulerError>
568 where
569 W: WorkDescriptor,
570 {
571 self.ensure_ready()?;
572 if limit == 0 {
573 return Err(SchedulerError::Capacity(
574 "scheduler preparation bound must be positive".into(),
575 ));
576 }
577 self.expire_deadlines(now)?;
578 let mut count = 0;
579 let mut stalled = 0;
580 while count < limit && !self.ready.is_empty() {
581 let request = self.ready.pop_front().expect("ready queue is nonempty");
582 let branches = self.branch_count(request);
583 let can_branch = self
584 .requests
585 .get(&request)
586 .is_some_and(|entry| branches == 0 || entry.state.permits_parallel_branches());
587 if !can_branch || branches >= self.limits.max_in_flight_per_request {
588 self.ready.push_back(request);
589 stalled += 1;
590 if stalled >= self.ready.len() {
591 break;
592 }
593 continue;
594 }
595 stalled = 0;
596 let queued = self
597 .requests
598 .get_mut(&request)
599 .and_then(|entry| entry.pending.pop_front())
600 .expect("ready request owns queued work");
601 if self
602 .requests
603 .get(&request)
604 .is_some_and(|entry| !entry.pending.is_empty())
605 {
606 self.ready.push_back(request);
607 }
608 let slice = queued.work.execution_slice_size();
609 if slice == 0 || slice > self.limits.execution_slice {
610 self.fail_before_submission(queued.id);
611 return Err(SchedulerError::Descriptor(format!(
612 "work {:?} execution slice {slice} exceeds configured bound {}",
613 queued.id, self.limits.execution_slice
614 )));
615 }
616 let mut descriptor = Vec::new();
617 if let Err(error) = queued.work.encode_descriptor(&mut descriptor) {
618 self.fail_before_submission(queued.id);
619 return Err(SchedulerError::Descriptor(error.to_string()));
620 }
621 let branch = match self
622 .requests
623 .get(&request)
624 .expect("active request exists")
625 .state
626 .branch()
627 {
628 Ok(branch) => branch,
629 Err(error) => {
630 self.fail_before_submission(queued.id);
631 return Err(SchedulerError::State(error.to_string()));
632 }
633 };
634 self.lifecycle.insert(queued.id, WorkLifecycle::Prepared);
635 self.prepared.push_back(Prepared {
636 id: queued.id,
637 work: queued.work,
638 descriptor,
639 branch,
640 deadline: queued.deadline,
641 });
642 count += 1;
643 }
644 Ok(count)
645 }
646
647 pub fn submit_prepared<E>(
649 &mut self,
650 now: Instant,
651 mut execute: impl FnMut(WorkId, &W, &mut S::Branch) -> Result<O, E>,
652 ) -> Result<usize, SchedulerError>
653 where
654 E: std::error::Error,
655 {
656 self.ensure_ready()?;
657 self.expire_deadlines(now)?;
658 let capacity = self
659 .limits
660 .max_in_flight_global
661 .saturating_sub(self.submitted.len())
662 .min(self.limits.max_new_submissions_per_turn);
663 let mut count = 0;
664 while count < capacity {
665 let Some(mut prepared) = self.prepared.pop_front() else {
666 break;
667 };
668 if prepared.deadline.is_some_and(|deadline| deadline <= now) {
669 let request = prepared.id.request();
670 self.prepared.push_front(prepared);
671 self.cancel_internal(request, CancellationCause::Deadline, now)?;
672 continue;
673 }
674 let output = match execute(prepared.id, &prepared.work, &mut prepared.branch) {
675 Ok(output) => output,
676 Err(error) => {
677 let id = prepared.id;
678 let discard = S::discard_branch(prepared.branch).err();
679 self.lifecycle.insert(id, WorkLifecycle::Failed);
680 self.failed_work = self.failed_work.saturating_add(1);
681 self.accepted_work = self.accepted_work.saturating_sub(1);
682 self.fail_request(id.request());
683 let message = discard.map_or_else(
684 || error.to_string(),
685 |discard| format!("{error}; branch discard also failed: {discard}"),
686 );
687 return Err(SchedulerError::Submission(message));
688 }
689 };
690 if let Some(backend) = output.backend_name() {
691 self.observed_backends.insert(backend);
692 }
693 self.all_outputs_preemptible &= output.physically_preemptible();
694 self.lifecycle.insert(prepared.id, WorkLifecycle::Submitted);
695 self.submitted.push(Submitted {
696 id: prepared.id,
697 work: prepared.work,
698 branch: prepared.branch,
699 output,
700 disposition: Disposition::Publish,
701 });
702 count += 1;
703 self.peak_in_flight_work = self.peak_in_flight_work.max(self.submitted.len());
704 }
705 if count != 0 {
706 self.drain_cycles = self.drain_cycles.saturating_add(1);
707 }
708 self.update_abandoned_resource_peak();
709 Ok(count)
710 }
711
712 pub fn poll_completions(&mut self, now: Instant) -> SchedulerProgress<W, O> {
714 let mut progress = SchedulerProgress::default();
715 let mut retained = Vec::with_capacity(self.submitted.len());
716 for submitted in std::mem::take(&mut self.submitted) {
717 match submitted.output.is_complete() {
718 Ok(false) => retained.push(submitted),
719 Ok(true) => self.resolve_completed(submitted, now, &mut progress),
720 Err(error) => {
721 let id = submitted.id;
722 let already_failed = matches!(submitted.disposition, Disposition::Fail);
723 let discard = S::discard_branch(submitted.branch).err();
724 self.lifecycle.insert(id, WorkLifecycle::Failed);
725 if !already_failed {
726 self.failed_work = self.failed_work.saturating_add(1);
727 }
728 self.accepted_work = self.accepted_work.saturating_sub(1);
729 self.fail_request(id.request());
730 progress.failed.push((
731 id,
732 SchedulerError::Completion(discard.map_or_else(
733 || error.to_string(),
734 |discard| format!("{error}; branch discard also failed: {discard}"),
735 )),
736 ));
737 }
738 }
739 }
740 for submitted in &mut retained {
745 if !self.requests.contains_key(&submitted.id.request())
746 && matches!(submitted.disposition, Disposition::Publish)
747 {
748 submitted.disposition = Disposition::Abandon { cancelled_at: now };
749 self.lifecycle
750 .insert(submitted.id, WorkLifecycle::Abandoned);
751 }
752 }
753 self.submitted = retained;
754 self.update_abandoned_resource_peak();
755 progress
756 }
757
758 pub fn run_local_turn<E>(
760 &mut self,
761 now: Instant,
762 execute: impl FnMut(WorkId, &W, &mut S::Branch) -> Result<O, E>,
763 ) -> Result<SchedulerProgress<W, O>, SchedulerError>
764 where
765 W: WorkDescriptor,
766 E: std::error::Error,
767 {
768 self.ensure_ready()?;
769 let mut progress = self.poll_completions(now);
770 self.prepare_bounded(self.limits.max_new_submissions_per_turn, now)?;
771 progress.newly_submitted = self.submit_prepared(now, execute)?;
772 let after_submit = self.poll_completions(now);
773 progress.committed.extend(after_submit.committed);
774 progress.failed.extend(after_submit.failed);
775 Ok(progress)
776 }
777
778 pub fn run_distributed_turn<T, E>(
780 &mut self,
781 protocol: u64,
782 transport: &T,
783 now: Instant,
784 execute: impl FnMut(WorkId, &W, &mut S::Branch) -> Result<O, E>,
785 ) -> Result<SchedulerProgress<W, O>, SchedulerError>
786 where
787 W: WorkDescriptor,
788 T: ConsensusTransport,
789 E: std::error::Error,
790 {
791 self.ensure_ready()?;
792 self.expire_deadlines_distributed(protocol, transport, now)?;
793 let mut progress = self.poll_distributed(protocol, transport, now)?;
794 self.prepare_bounded(self.limits.max_new_submissions_per_turn, now)?;
795 let plan = self
796 .prepared
797 .iter()
798 .take(
799 self.limits
800 .max_in_flight_global
801 .saturating_sub(self.submitted.len())
802 .min(self.limits.max_new_submissions_per_turn),
803 )
804 .map(|work| ScheduledWork {
805 id: work.id,
806 descriptor: &work.descriptor,
807 })
808 .collect::<Vec<_>>();
809 if let Err(error) = validate_schedule(transport, &plan, self.drain_cycles, protocol) {
810 self.poison(error.to_string(), now);
811 return Err(SchedulerError::Consensus(error.to_string()));
812 }
813 progress.newly_submitted = match self.submit_prepared(now, execute) {
814 Ok(count) => count,
815 Err(error) => {
816 self.poison(error.to_string(), now);
817 return Err(error);
818 }
819 };
820 Ok(progress)
821 }
822
823 pub fn cancel_distributed<T: ConsensusTransport>(
825 &mut self,
826 protocol: u64,
827 request: RequestId,
828 transport: &T,
829 now: Instant,
830 ) -> Result<(), SchedulerError> {
831 self.ensure_ready()?;
832 if let Err(error) =
833 validate_disposition(transport, protocol, request, CancellationCause::Explicit)
834 {
835 self.poison(error.to_string(), now);
836 return Err(SchedulerError::Consensus(error.to_string()));
837 }
838 self.cancel_internal(request, CancellationCause::Explicit, now)
839 }
840
841 pub fn finish(&mut self, request: RequestId) -> Result<(), SchedulerError> {
843 self.ensure_ready()?;
844 if self
845 .submitted
846 .iter()
847 .any(|work| work.id.request() == request)
848 {
849 return Err(SchedulerError::State(format!(
850 "request {} still has submitted work",
851 request.value()
852 )));
853 }
854 let entry = self
855 .requests
856 .remove(&request)
857 .ok_or(SchedulerError::UnknownRequest(request))?;
858 let queued = entry.pending.len();
859 for work in entry.pending {
860 self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
861 }
862 self.ready.retain(|candidate| *candidate != request);
863 let (prepared, discard_error) = self.discard_prepared_for_request(request);
864 let discarded = queued + prepared;
865 self.accepted_work = self.accepted_work.saturating_sub(discarded);
866 self.discarded_work = self.discarded_work.saturating_add(discarded as u64);
867 self.terminal.insert(
868 request,
869 if discard_error.is_some() {
870 RequestStatus::Failed
871 } else {
872 RequestStatus::Finished
873 },
874 );
875 self.finished_requests = self.finished_requests.saturating_add(1);
876 discard_error.map_or(Ok(()), |error| Err(SchedulerError::State(error)))
877 }
878
879 pub fn cancel(&mut self, request: RequestId) -> Result<(), SchedulerError> {
881 self.cancel_internal(request, CancellationCause::Explicit, Instant::now())
882 }
883
884 pub fn release(&mut self, request: RequestId) -> Result<S, SchedulerError> {
886 self.ensure_ready()?;
887 let entry = self
888 .requests
889 .get(&request)
890 .ok_or(SchedulerError::UnknownRequest(request))?;
891 if !entry.pending.is_empty() || self.branch_count(request) != 0 {
892 return Err(SchedulerError::State(format!(
893 "request {} still owns unpublished work",
894 request.value()
895 )));
896 }
897 Ok(self
898 .requests
899 .remove(&request)
900 .expect("checked active request")
901 .state)
902 }
903
904 pub fn forget_terminal(&mut self, request: RequestId) -> Result<RequestStatus, SchedulerError> {
906 self.ensure_ready()?;
907 if self
908 .submitted
909 .iter()
910 .any(|work| work.id.request() == request)
911 {
912 return Err(SchedulerError::State(format!(
913 "request {} still has retained abandoned work",
914 request.value()
915 )));
916 }
917 self.terminal
918 .remove(&request)
919 .ok_or(SchedulerError::UnknownRequest(request))
920 }
921
922 pub fn request_status(&self, id: RequestId) -> Option<RequestStatus> {
924 self.requests
925 .contains_key(&id)
926 .then_some(RequestStatus::Active)
927 .or_else(|| self.terminal.get(&id).copied())
928 }
929
930 pub fn work_lifecycle(&self, id: WorkId) -> Option<WorkLifecycle> {
932 self.lifecycle.get(&id).copied()
933 }
934
935 pub fn queued_for_request(&self, request: RequestId) -> usize {
937 self.requests
938 .get(&request)
939 .map_or(0, |entry| entry.pending.len())
940 }
941
942 pub fn capabilities(&self) -> SchedulerCapabilities {
944 SchedulerCapabilities {
945 limits: self.limits,
946 observed_backends: self.observed_backends.iter().cloned().collect(),
947 executing_work_physically_preemptible: !self.observed_backends.is_empty()
948 && self.all_outputs_preemptible,
949 non_preemptible_interval:
950 "from exact backend submission until that transition's completion resolves".into(),
951 }
952 }
953
954 pub fn report(&self) -> SchedulerReport {
956 let abandoned = self
957 .submitted
958 .iter()
959 .filter(|work| matches!(work.disposition, Disposition::Abandon { .. }))
960 .count();
961 let failed = self
962 .submitted
963 .iter()
964 .filter(|work| matches!(work.disposition, Disposition::Fail))
965 .count();
966 SchedulerReport {
967 active_requests: self.requests.len(),
968 queued_work: self
969 .requests
970 .values()
971 .map(|entry| entry.pending.len())
972 .sum(),
973 prepared_work: self.prepared.len(),
974 submitted_in_flight_work: self.submitted.len() - abandoned - failed,
975 completing_work: 0,
976 abandoned_in_flight_work: abandoned,
977 failed_in_flight_work: failed,
978 current_in_flight_work: self.submitted.len(),
979 peak_in_flight_work: self.peak_in_flight_work,
980 peak_queued_work: self.peak_accepted_work,
981 submitted_work: self.submitted_work,
982 completed_work: self.completed_work,
983 failed_work: self.failed_work,
984 discarded_work: self.discarded_work,
985 cancellation_before_submission: self.cancellation_before_submission,
986 cancellation_after_submission: self.cancellation_after_submission,
987 abandoned_released_work: self.abandoned_released_work,
988 abandoned_retained_resources: self.abandoned_retained_resources(),
989 peak_abandoned_retained_resources: self.peak_abandoned_retained_resources,
990 last_cancellation_to_release_ns: self
991 .last_cancellation_to_release
992 .map(|duration| duration.as_nanos()),
993 max_cancellation_to_release_ns: self
994 .max_cancellation_to_release
995 .map(|duration| duration.as_nanos()),
996 finished_requests: self.finished_requests,
997 cancelled_requests: self.cancelled_requests,
998 deadline_expired_requests: self.deadline_expired_requests,
999 drain_cycles: self.drain_cycles,
1000 configured_submission_bound: self.limits.max_new_submissions_per_turn,
1001 configured_slice_bound: self.limits.execution_slice,
1002 poisoned: self.poisoned.is_some(),
1003 }
1004 }
1005
1006 pub fn poison_reason(&self) -> Option<&str> {
1008 self.poisoned.as_deref()
1009 }
1010
1011 fn resolve_completed(
1012 &mut self,
1013 submitted: Submitted<W, S::Branch, O>,
1014 now: Instant,
1015 progress: &mut SchedulerProgress<W, O>,
1016 ) {
1017 match submitted.disposition {
1018 Disposition::Abandon { cancelled_at } => {
1019 if let Err(error) = S::discard_branch(submitted.branch) {
1020 self.lifecycle.insert(submitted.id, WorkLifecycle::Failed);
1021 self.failed_work = self.failed_work.saturating_add(1);
1022 progress.failed.push((
1023 submitted.id,
1024 SchedulerError::State(format!(
1025 "failed to discard abandoned state branch: {error}"
1026 )),
1027 ));
1028 }
1029 self.abandoned_released_work = self.abandoned_released_work.saturating_add(1);
1030 self.accepted_work = self.accepted_work.saturating_sub(1);
1031 let latency = now.saturating_duration_since(cancelled_at);
1032 self.last_cancellation_to_release = Some(latency);
1033 self.max_cancellation_to_release = Some(
1034 self.max_cancellation_to_release
1035 .map_or(latency, |previous| previous.max(latency)),
1036 );
1037 }
1038 Disposition::Publish => {
1039 self.lifecycle
1040 .insert(submitted.id, WorkLifecycle::Completing);
1041 let Some(request) = self.requests.get_mut(&submitted.id.request()) else {
1042 if let Err(error) = S::discard_branch(submitted.branch) {
1043 self.lifecycle.insert(submitted.id, WorkLifecycle::Failed);
1044 self.failed_work = self.failed_work.saturating_add(1);
1045 progress.failed.push((
1046 submitted.id,
1047 SchedulerError::State(format!(
1048 "failed to discard unpublished state branch: {error}"
1049 )),
1050 ));
1051 } else {
1052 self.lifecycle
1053 .insert(submitted.id, WorkLifecycle::Abandoned);
1054 }
1055 self.accepted_work = self.accepted_work.saturating_sub(1);
1056 return;
1057 };
1058 if let Err(error) = request.state.commit_branch(submitted.branch) {
1059 self.lifecycle.insert(submitted.id, WorkLifecycle::Failed);
1060 self.failed_work = self.failed_work.saturating_add(1);
1061 self.accepted_work = self.accepted_work.saturating_sub(1);
1062 self.fail_request(submitted.id.request());
1063 progress
1064 .failed
1065 .push((submitted.id, SchedulerError::State(error.to_string())));
1066 return;
1067 }
1068 self.lifecycle
1069 .insert(submitted.id, WorkLifecycle::Committed);
1070 self.completed_work = self.completed_work.saturating_add(1);
1071 self.accepted_work = self.accepted_work.saturating_sub(1);
1072 progress
1073 .committed
1074 .push((submitted.id, submitted.work, submitted.output));
1075 }
1076 Disposition::Fail => {
1077 let discard = S::discard_branch(submitted.branch).err();
1078 self.lifecycle.insert(submitted.id, WorkLifecycle::Failed);
1079 self.accepted_work = self.accepted_work.saturating_sub(1);
1080 self.fail_request(submitted.id.request());
1081 progress.failed.push((
1082 submitted.id,
1083 SchedulerError::DistributedCompletion(discard.map_or_else(
1084 || "backend completion failed on at least one rank".into(),
1085 |error| {
1086 format!(
1087 "backend completion failed on at least one rank; branch discard also failed: {error}"
1088 )
1089 },
1090 )),
1091 ));
1092 }
1093 }
1094 }
1095
1096 fn poll_distributed<T: ConsensusTransport>(
1097 &mut self,
1098 protocol: u64,
1099 transport: &T,
1100 now: Instant,
1101 ) -> Result<SchedulerProgress<W, O>, SchedulerError> {
1102 let local = self
1103 .submitted
1104 .iter()
1105 .map(|work| {
1106 let status = match work.output.is_complete() {
1107 Ok(false) => CompletionObservation::Incomplete,
1108 Ok(true) => CompletionObservation::Complete,
1109 Err(_) => CompletionObservation::Failed,
1110 };
1111 (work.id, status)
1112 })
1113 .collect::<Vec<_>>();
1114 let global = match resolve_completions(transport, protocol, &local) {
1115 Ok(global) => global,
1116 Err(error) => {
1117 self.poison(error.to_string(), now);
1118 return Err(SchedulerError::Consensus(error.to_string()));
1119 }
1120 };
1121
1122 let mut progress = SchedulerProgress::default();
1123 let mut retained = Vec::with_capacity(self.submitted.len());
1124 for (mut work, status) in std::mem::take(&mut self.submitted).into_iter().zip(global) {
1125 match status {
1126 CompletionResolution::Incomplete => retained.push(work),
1127 CompletionResolution::Complete => {
1128 self.resolve_completed(work, now, &mut progress);
1129 }
1130 CompletionResolution::FailedPending => {
1131 if !matches!(work.disposition, Disposition::Fail) {
1132 work.disposition = Disposition::Fail;
1133 self.lifecycle.insert(work.id, WorkLifecycle::Failed);
1134 self.failed_work = self.failed_work.saturating_add(1);
1135 }
1136 retained.push(work);
1137 }
1138 CompletionResolution::FailedComplete => {
1139 if !matches!(work.disposition, Disposition::Fail) {
1140 work.disposition = Disposition::Fail;
1141 self.failed_work = self.failed_work.saturating_add(1);
1142 }
1143 self.resolve_completed(work, now, &mut progress);
1144 }
1145 }
1146 }
1147 for work in &mut retained {
1148 if !self.requests.contains_key(&work.id.request())
1149 && matches!(work.disposition, Disposition::Publish)
1150 {
1151 work.disposition = Disposition::Abandon { cancelled_at: now };
1152 self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
1153 }
1154 }
1155 self.submitted = retained;
1156 self.update_abandoned_resource_peak();
1157 Ok(progress)
1158 }
1159
1160 fn expire_deadlines(&mut self, now: Instant) -> Result<(), SchedulerError> {
1161 let mut expired = BTreeSet::new();
1162 for (request, entry) in &self.requests {
1163 if entry
1164 .pending
1165 .iter()
1166 .any(|work| work.deadline.is_some_and(|deadline| deadline <= now))
1167 {
1168 expired.insert(*request);
1169 }
1170 }
1171 for work in &self.prepared {
1172 if work.deadline.is_some_and(|deadline| deadline <= now) {
1173 expired.insert(work.id.request());
1174 }
1175 }
1176 for request in expired {
1177 self.cancel_internal(request, CancellationCause::Deadline, now)?;
1178 }
1179 Ok(())
1180 }
1181
1182 fn expire_deadlines_distributed<T: ConsensusTransport>(
1183 &mut self,
1184 protocol: u64,
1185 transport: &T,
1186 now: Instant,
1187 ) -> Result<(), SchedulerError> {
1188 let mut expired = BTreeSet::new();
1189 for (request, entry) in &self.requests {
1190 if entry
1191 .pending
1192 .iter()
1193 .any(|work| work.deadline.is_some_and(|deadline| deadline <= now))
1194 {
1195 expired.insert(*request);
1196 }
1197 }
1198 for work in &self.prepared {
1199 if work.deadline.is_some_and(|deadline| deadline <= now) {
1200 expired.insert(work.id.request());
1201 }
1202 }
1203 for request in expired {
1204 if let Err(error) =
1205 validate_disposition(transport, protocol, request, CancellationCause::Deadline)
1206 {
1207 self.poison(error.to_string(), now);
1208 return Err(SchedulerError::Consensus(error.to_string()));
1209 }
1210 self.cancel_internal(request, CancellationCause::Deadline, now)?;
1211 }
1212 Ok(())
1213 }
1214
1215 fn cancel_internal(
1216 &mut self,
1217 request: RequestId,
1218 cause: CancellationCause,
1219 now: Instant,
1220 ) -> Result<(), SchedulerError> {
1221 self.ensure_ready()?;
1222 if self.terminal.contains_key(&request) {
1223 return Err(SchedulerError::State(format!(
1224 "request {} is already terminal",
1225 request.value()
1226 )));
1227 }
1228 let entry = self
1229 .requests
1230 .remove(&request)
1231 .ok_or(SchedulerError::UnknownRequest(request))?;
1232 let queued = entry.pending.len();
1233 for work in entry.pending {
1234 self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
1235 }
1236 self.ready.retain(|candidate| *candidate != request);
1237 let (prepared, discard_error) = self.discard_prepared_for_request(request);
1238 let before_submission = queued + prepared;
1239 self.accepted_work = self.accepted_work.saturating_sub(before_submission);
1240 self.discarded_work = self.discarded_work.saturating_add(before_submission as u64);
1241 self.cancellation_before_submission = self
1242 .cancellation_before_submission
1243 .saturating_add(before_submission as u64);
1244 let mut after_submission = 0u64;
1245 for work in &mut self.submitted {
1246 if work.id.request() == request && matches!(work.disposition, Disposition::Publish) {
1247 work.disposition = Disposition::Abandon { cancelled_at: now };
1248 self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
1249 after_submission += 1;
1250 }
1251 }
1252 self.cancellation_after_submission = self
1253 .cancellation_after_submission
1254 .saturating_add(after_submission);
1255 let status = match cause {
1256 CancellationCause::Explicit => {
1257 self.cancelled_requests = self.cancelled_requests.saturating_add(1);
1258 RequestStatus::Cancelled
1259 }
1260 CancellationCause::Deadline => {
1261 self.deadline_expired_requests = self.deadline_expired_requests.saturating_add(1);
1262 RequestStatus::DeadlineExceeded
1263 }
1264 };
1265 self.terminal.insert(
1266 request,
1267 if discard_error.is_some() {
1268 RequestStatus::Failed
1269 } else {
1270 status
1271 },
1272 );
1273 self.update_abandoned_resource_peak();
1274 discard_error.map_or(Ok(()), |error| Err(SchedulerError::State(error)))
1275 }
1276
1277 fn discard_prepared_for_request(&mut self, request: RequestId) -> (usize, Option<String>) {
1278 let mut retained = VecDeque::with_capacity(self.prepared.len());
1279 let mut discarded = 0;
1280 let mut errors = Vec::new();
1281 for work in std::mem::take(&mut self.prepared) {
1282 if work.id.request() == request {
1283 discarded += 1;
1284 match S::discard_branch(work.branch) {
1285 Ok(()) => {
1286 self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
1287 }
1288 Err(error) => {
1289 self.lifecycle.insert(work.id, WorkLifecycle::Failed);
1290 self.failed_work = self.failed_work.saturating_add(1);
1291 errors.push(format!("work {:?}: {error}", work.id));
1292 }
1293 }
1294 } else {
1295 retained.push_back(work);
1296 }
1297 }
1298 self.prepared = retained;
1299 let error = (!errors.is_empty()).then(|| errors.join("; "));
1300 (discarded, error)
1301 }
1302
1303 fn branch_count(&self, request: RequestId) -> usize {
1304 self.prepared
1305 .iter()
1306 .filter(|work| work.id.request() == request)
1307 .count()
1308 + self
1309 .submitted
1310 .iter()
1311 .filter(|work| work.id.request() == request)
1312 .count()
1313 }
1314
1315 fn fail_before_submission(&mut self, id: WorkId) {
1316 self.lifecycle.insert(id, WorkLifecycle::Failed);
1317 self.failed_work = self.failed_work.saturating_add(1);
1318 self.accepted_work = self.accepted_work.saturating_sub(1);
1319 self.fail_request(id.request());
1320 }
1321
1322 fn fail_request(&mut self, request: RequestId) {
1323 let Some(entry) = self.requests.remove(&request) else {
1324 return;
1325 };
1326 let queued = entry.pending.len();
1327 for work in entry.pending {
1328 self.lifecycle.insert(work.id, WorkLifecycle::Failed);
1329 }
1330 self.ready.retain(|candidate| *candidate != request);
1331 let (prepared, _) = self.discard_prepared_for_request(request);
1332 let discarded = queued + prepared;
1333 self.accepted_work = self.accepted_work.saturating_sub(discarded);
1334 self.discarded_work = self.discarded_work.saturating_add(discarded as u64);
1335 for work in &mut self.submitted {
1336 if work.id.request() == request && matches!(work.disposition, Disposition::Publish) {
1337 work.disposition = Disposition::Abandon {
1338 cancelled_at: Instant::now(),
1339 };
1340 self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
1341 }
1342 }
1343 self.terminal.insert(request, RequestStatus::Failed);
1344 }
1345
1346 fn ensure_ready(&self) -> Result<(), SchedulerError> {
1347 self.poisoned.as_ref().map_or(Ok(()), |reason| {
1348 Err(SchedulerError::Poisoned(reason.clone()))
1349 })
1350 }
1351
1352 fn poison(&mut self, reason: String, now: Instant) {
1353 if self.poisoned.is_some() {
1354 return;
1355 }
1356 let mut discarded = 0usize;
1357 for (request, entry) in std::mem::take(&mut self.requests) {
1358 self.terminal.insert(request, RequestStatus::Failed);
1359 for work in entry.pending {
1360 self.lifecycle.insert(work.id, WorkLifecycle::Failed);
1361 discarded += 1;
1362 }
1363 }
1364 self.ready.clear();
1365 let mut cleanup_errors = Vec::new();
1366 for work in std::mem::take(&mut self.prepared) {
1367 let id = work.id;
1368 if let Err(error) = S::discard_branch(work.branch) {
1369 cleanup_errors.push(format!("work {id:?}: {error}"));
1370 self.failed_work = self.failed_work.saturating_add(1);
1371 }
1372 self.lifecycle.insert(id, WorkLifecycle::Failed);
1373 discarded += 1;
1374 }
1375 self.accepted_work = self.accepted_work.saturating_sub(discarded);
1376 self.discarded_work = self.discarded_work.saturating_add(discarded as u64);
1377 for work in &mut self.submitted {
1378 if matches!(work.disposition, Disposition::Publish) {
1379 work.disposition = Disposition::Abandon { cancelled_at: now };
1380 self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
1381 }
1382 }
1383 self.poisoned = Some(if cleanup_errors.is_empty() {
1384 reason
1385 } else {
1386 format!(
1387 "{reason}; prepared branch cleanup failed: {}",
1388 cleanup_errors.join("; ")
1389 )
1390 });
1391 self.update_abandoned_resource_peak();
1392 }
1393
1394 fn abandoned_retained_resources(&self) -> usize {
1395 self.submitted
1396 .iter()
1397 .filter(|work| matches!(work.disposition, Disposition::Abandon { .. }))
1398 .map(|work| work.output.retained_resources())
1399 .sum()
1400 }
1401
1402 fn update_abandoned_resource_peak(&mut self) {
1403 self.peak_abandoned_retained_resources = self
1404 .peak_abandoned_retained_resources
1405 .max(self.abandoned_retained_resources());
1406 }
1407}
1408
1409#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
1411pub enum SchedulerError {
1412 #[error("scheduler limits must be positive, got {0:?}")]
1414 InvalidLimits([usize; 6]),
1415 #[error("request {} is already registered", .0.value())]
1417 DuplicateRequest(RequestId),
1418 #[error("request {} is not active", .0.value())]
1420 UnknownRequest(RequestId),
1421 #[error("{0}")]
1423 Capacity(String),
1424 #[error("work descriptor failed: {0}")]
1426 Descriptor(String),
1427 #[error("semantic state transaction failed: {0}")]
1429 State(String),
1430 #[error("backend submission failed: {0}")]
1432 Submission(String),
1433 #[error("exact completion observation failed: {0}")]
1435 Completion(String),
1436 #[error("distributed scheduler consensus failed: {0}")]
1438 Consensus(String),
1439 #[error("distributed {0}")]
1441 DistributedCompletion(String),
1442 #[error("scheduler is poisoned after unsafe distributed ordering: {0}")]
1444 Poisoned(String),
1445}
1446
1447#[cfg(test)]
1448mod tests {
1449 use super::*;
1450 use std::{
1451 cell::{Cell, RefCell},
1452 collections::VecDeque,
1453 convert::Infallible,
1454 rc::Rc,
1455 };
1456
1457 #[derive(Default)]
1458 struct State(u32);
1459
1460 impl SemanticStateTransaction for State {
1461 type Branch = u32;
1462 type Error = Infallible;
1463
1464 fn branch(&self) -> Result<u32, Infallible> {
1465 Ok(self.0 + 1)
1466 }
1467
1468 fn commit_branch(&mut self, branch: u32) -> Result<(), Infallible> {
1469 self.0 = branch;
1470 Ok(())
1471 }
1472
1473 fn permits_parallel_branches(&self) -> bool {
1474 true
1475 }
1476 }
1477
1478 impl WorkDescriptor for u32 {
1479 type Error = Infallible;
1480
1481 fn encode_descriptor(&self, output: &mut Vec<u32>) -> Result<(), Self::Error> {
1482 output.push(*self);
1483 Ok(())
1484 }
1485 }
1486
1487 #[derive(Debug)]
1488 struct Output {
1489 complete: Rc<Cell<bool>>,
1490 fail: bool,
1491 }
1492
1493 impl TransitionOutput for Output {
1494 type Error = std::io::Error;
1495
1496 fn is_complete(&self) -> Result<bool, Self::Error> {
1497 if self.fail {
1498 Err(std::io::Error::other("mock failure"))
1499 } else {
1500 Ok(self.complete.get())
1501 }
1502 }
1503
1504 fn backend_name(&self) -> Option<String> {
1505 Some("mock".into())
1506 }
1507
1508 fn retained_resources(&self) -> usize {
1509 2
1510 }
1511 }
1512
1513 #[derive(Default)]
1514 struct GatherStep {
1515 replacements: Vec<(usize, usize, u32)>,
1516 }
1517
1518 struct ScriptedTransport {
1519 participants: usize,
1520 steps: RefCell<VecDeque<GatherStep>>,
1521 }
1522
1523 impl ScriptedTransport {
1524 fn new(participants: usize, steps: Vec<GatherStep>) -> Self {
1525 Self {
1526 participants,
1527 steps: RefCell::new(steps.into()),
1528 }
1529 }
1530 }
1531
1532 impl ConsensusTransport for ScriptedTransport {
1533 type Error = Infallible;
1534
1535 fn participant_count(&self) -> usize {
1536 self.participants
1537 }
1538
1539 fn all_gather_words(&self, local: &[u32]) -> Result<Vec<u32>, Self::Error> {
1540 let mut gathered = local.repeat(self.participants);
1541 let step = self.steps.borrow_mut().pop_front().unwrap_or_default();
1542 for (rank, offset, value) in step.replacements {
1543 gathered[rank * local.len() + offset] = value;
1544 }
1545 Ok(gathered)
1546 }
1547 }
1548
1549 fn scheduler() -> Scheduler<u32, State, Output> {
1550 Scheduler::new(SchedulerLimits::default()).unwrap()
1551 }
1552
1553 #[test]
1554 fn scheduler_construction_revalidates_public_limit_fields() {
1555 let invalid = SchedulerLimits {
1556 max_active_requests: 0,
1557 ..SchedulerLimits::default()
1558 };
1559 assert!(matches!(
1560 Scheduler::<u32, State, Output>::new(invalid),
1561 Err(SchedulerError::InvalidLimits(_))
1562 ));
1563 }
1564
1565 #[test]
1566 fn queued_prepared_submitted_committed_exactly() {
1567 let done = Rc::new(Cell::new(false));
1568 let mut scheduler = scheduler();
1569 let request = RequestId::new(1);
1570 scheduler.register(request, State::default()).unwrap();
1571 let id = scheduler.enqueue(request, 7).unwrap();
1572 assert_eq!(scheduler.work_lifecycle(id), Some(WorkLifecycle::Queued));
1573 scheduler.prepare_bounded(1, Instant::now()).unwrap();
1574 assert_eq!(scheduler.work_lifecycle(id), Some(WorkLifecycle::Prepared));
1575 scheduler
1576 .submit_prepared(Instant::now(), |_, _, _| {
1577 Ok::<_, Infallible>(Output {
1578 complete: done.clone(),
1579 fail: false,
1580 })
1581 })
1582 .unwrap();
1583 assert!(scheduler
1584 .poll_completions(Instant::now())
1585 .committed
1586 .is_empty());
1587 done.set(true);
1588 assert_eq!(
1589 scheduler.poll_completions(Instant::now()).committed.len(),
1590 1
1591 );
1592 assert_eq!(scheduler.request_state(request).unwrap().0, 1);
1593 }
1594
1595 #[test]
1596 fn cancellation_retains_submitted_resources_until_completion() {
1597 let done = Rc::new(Cell::new(false));
1598 let mut scheduler = scheduler();
1599 let request = RequestId::new(2);
1600 scheduler.register(request, State::default()).unwrap();
1601 let id = scheduler.enqueue(request, 1).unwrap();
1602 scheduler.prepare_bounded(1, Instant::now()).unwrap();
1603 scheduler
1604 .submit_prepared(Instant::now(), |_, _, _| {
1605 Ok::<_, Infallible>(Output {
1606 complete: done.clone(),
1607 fail: false,
1608 })
1609 })
1610 .unwrap();
1611 scheduler.cancel(request).unwrap();
1612 assert_eq!(scheduler.work_lifecycle(id), Some(WorkLifecycle::Abandoned));
1613 assert_eq!(scheduler.report().abandoned_retained_resources, 2);
1614 assert_eq!(
1615 scheduler.request_status(request),
1616 Some(RequestStatus::Cancelled)
1617 );
1618 done.set(true);
1619 scheduler.poll_completions(Instant::now());
1620 assert_eq!(scheduler.report().current_in_flight_work, 0);
1621 assert_eq!(scheduler.report().abandoned_released_work, 1);
1622 }
1623
1624 #[test]
1625 fn exact_completion_failure_does_not_commit() {
1626 let mut scheduler = scheduler();
1627 let request = RequestId::new(3);
1628 scheduler.register(request, State::default()).unwrap();
1629 let id = scheduler.enqueue(request, 1).unwrap();
1630 scheduler.prepare_bounded(1, Instant::now()).unwrap();
1631 scheduler
1632 .submit_prepared(Instant::now(), |_, _, _| {
1633 Ok::<_, Infallible>(Output {
1634 complete: Rc::new(Cell::new(true)),
1635 fail: true,
1636 })
1637 })
1638 .unwrap();
1639 assert_eq!(scheduler.poll_completions(Instant::now()).failed.len(), 1);
1640 assert_eq!(scheduler.work_lifecycle(id), Some(WorkLifecycle::Failed));
1641 assert_eq!(
1642 scheduler.request_status(request),
1643 Some(RequestStatus::Failed)
1644 );
1645 }
1646
1647 #[test]
1648 fn request_failure_abandons_incomplete_sibling_until_exact_completion() {
1649 let waiting = Rc::new(Cell::new(false));
1650 let limits = SchedulerLimits::with_execution_bounds(1, 2, 2, 2, 2, 1).unwrap();
1651 let mut scheduler = Scheduler::new(limits).unwrap();
1652 let request = RequestId::new(31);
1653 scheduler.register(request, State::default()).unwrap();
1654 let ids = scheduler.enqueue_batch(request, vec![1, 2]).unwrap();
1655 scheduler.prepare_bounded(2, Instant::now()).unwrap();
1656 scheduler
1657 .submit_prepared(Instant::now(), |id, _, _| {
1658 Ok::<_, Infallible>(Output {
1659 complete: if id.sequence() == 0 {
1660 Rc::new(Cell::new(true))
1661 } else {
1662 waiting.clone()
1663 },
1664 fail: id.sequence() == 0,
1665 })
1666 })
1667 .unwrap();
1668
1669 let progress = scheduler.poll_completions(Instant::now());
1670 assert_eq!(progress.failed.len(), 1);
1671 assert_eq!(
1672 scheduler.work_lifecycle(ids[0]),
1673 Some(WorkLifecycle::Failed)
1674 );
1675 assert_eq!(
1676 scheduler.work_lifecycle(ids[1]),
1677 Some(WorkLifecycle::Abandoned)
1678 );
1679 assert_eq!(scheduler.report().abandoned_retained_resources, 2);
1680
1681 waiting.set(true);
1682 scheduler.poll_completions(Instant::now());
1683 assert_eq!(scheduler.report().current_in_flight_work, 0);
1684 assert_eq!(scheduler.report().abandoned_released_work, 1);
1685 }
1686
1687 #[test]
1688 fn submission_failure_marks_work_and_request_failed() {
1689 let mut scheduler = scheduler();
1690 let request = RequestId::new(4);
1691 scheduler.register(request, State::default()).unwrap();
1692 let id = scheduler.enqueue(request, 1).unwrap();
1693 scheduler.prepare_bounded(1, Instant::now()).unwrap();
1694 let error = scheduler
1695 .submit_prepared(Instant::now(), |_, _, _| {
1696 Err::<Output, _>(std::io::Error::other("mock submit"))
1697 })
1698 .unwrap_err();
1699 assert!(matches!(error, SchedulerError::Submission(_)));
1700 assert_eq!(scheduler.work_lifecycle(id), Some(WorkLifecycle::Failed));
1701 assert_eq!(
1702 scheduler.request_status(request),
1703 Some(RequestStatus::Failed)
1704 );
1705 }
1706
1707 #[test]
1708 fn deadlines_batches_and_fairness_are_backend_neutral() {
1709 let mut scheduler = scheduler();
1710 let first = RequestId::new(10);
1711 let second = RequestId::new(20);
1712 scheduler.register(first, State::default()).unwrap();
1713 scheduler.register(second, State::default()).unwrap();
1714 scheduler.enqueue_batch(first, vec![1, 2]).unwrap();
1715 scheduler.enqueue(second, 3).unwrap();
1716 scheduler.prepare_bounded(3, Instant::now()).unwrap();
1717 let mut order = Vec::new();
1718 scheduler
1719 .submit_prepared(Instant::now(), |id, _, _| {
1720 order.push(id.request());
1721 Ok::<_, Infallible>(Output {
1722 complete: Rc::new(Cell::new(true)),
1723 fail: false,
1724 })
1725 })
1726 .unwrap();
1727 assert_eq!(order, vec![first]);
1728
1729 let expired = RequestId::new(30);
1730 scheduler.register(expired, State::default()).unwrap();
1731 scheduler
1732 .enqueue_with_deadline(
1733 expired,
1734 4,
1735 Some(Instant::now().checked_sub(Duration::from_secs(1)).unwrap()),
1736 )
1737 .unwrap();
1738 scheduler.prepare_bounded(1, Instant::now()).unwrap();
1739 assert_eq!(
1740 scheduler.request_status(expired),
1741 Some(RequestStatus::DeadlineExceeded)
1742 );
1743 }
1744
1745 #[test]
1746 fn distributed_schedule_mismatch_poisons_the_canonical_machine() {
1747 let transport = ScriptedTransport::new(
1748 2,
1749 vec![
1750 GatherStep::default(),
1751 GatherStep {
1752 replacements: vec![(1, 10, 99)],
1753 },
1754 ],
1755 );
1756 let mut scheduler = scheduler();
1757 let request = RequestId::new(50);
1758 scheduler.register(request, State::default()).unwrap();
1759 let work = scheduler.enqueue(request, 7).unwrap();
1760
1761 let error = scheduler
1762 .run_distributed_turn(0xAA, &transport, Instant::now(), |_, _, _| {
1763 Ok::<_, Infallible>(Output {
1764 complete: Rc::new(Cell::new(false)),
1765 fail: false,
1766 })
1767 })
1768 .unwrap_err();
1769 assert!(matches!(error, SchedulerError::Consensus(_)));
1770 assert!(scheduler.poison_reason().is_some());
1771 assert!(scheduler.report().poisoned);
1772 assert_eq!(
1773 scheduler.request_status(request),
1774 Some(RequestStatus::Failed)
1775 );
1776 assert_eq!(scheduler.work_lifecycle(work), Some(WorkLifecycle::Failed));
1777 assert!(matches!(
1778 scheduler.enqueue(request, 8),
1779 Err(SchedulerError::Poisoned(_))
1780 ));
1781 }
1782
1783 #[test]
1784 fn poisoned_submissions_remain_retained_until_local_exact_completion() {
1785 let transport = ScriptedTransport::new(
1786 2,
1787 vec![
1788 GatherStep::default(),
1789 GatherStep::default(),
1790 GatherStep::default(),
1791 GatherStep {
1792 replacements: vec![(1, 4, 99)],
1793 },
1794 ],
1795 );
1796 let done = Rc::new(Cell::new(false));
1797 let mut scheduler = scheduler();
1798 let request = RequestId::new(54);
1799 scheduler.register(request, State::default()).unwrap();
1800 scheduler.enqueue(request, 1).unwrap();
1801 scheduler
1802 .run_distributed_turn(0xEE, &transport, Instant::now(), |_, _, _| {
1803 Ok::<_, Infallible>(Output {
1804 complete: done.clone(),
1805 fail: false,
1806 })
1807 })
1808 .unwrap();
1809
1810 let error = scheduler
1811 .run_distributed_turn(0xEE, &transport, Instant::now(), |_, _, _| {
1812 Ok::<_, Infallible>(Output {
1813 complete: done.clone(),
1814 fail: false,
1815 })
1816 })
1817 .unwrap_err();
1818 assert!(matches!(error, SchedulerError::Consensus(_)));
1819 assert_eq!(scheduler.report().abandoned_in_flight_work, 1);
1820 assert_eq!(scheduler.report().abandoned_retained_resources, 2);
1821
1822 assert!(scheduler
1823 .poll_completions(Instant::now())
1824 .committed
1825 .is_empty());
1826 assert_eq!(scheduler.report().current_in_flight_work, 1);
1827 done.set(true);
1828 assert!(scheduler
1829 .poll_completions(Instant::now())
1830 .committed
1831 .is_empty());
1832 assert_eq!(scheduler.report().current_in_flight_work, 0);
1833 assert_eq!(scheduler.report().abandoned_released_work, 1);
1834 }
1835
1836 #[test]
1837 fn distributed_failure_retains_output_until_every_rank_is_terminal() {
1838 let transport = ScriptedTransport::new(
1839 3,
1840 vec![
1841 GatherStep::default(),
1842 GatherStep::default(),
1843 GatherStep {
1844 replacements: vec![(1, 7, 2), (2, 7, 0)],
1845 },
1846 GatherStep::default(),
1847 GatherStep {
1848 replacements: vec![(1, 7, 2), (2, 7, 1)],
1849 },
1850 GatherStep::default(),
1851 ],
1852 );
1853 let mut scheduler = scheduler();
1854 let request = RequestId::new(51);
1855 scheduler.register(request, State::default()).unwrap();
1856 scheduler.enqueue(request, 1).unwrap();
1857 let output = || Output {
1858 complete: Rc::new(Cell::new(true)),
1859 fail: false,
1860 };
1861
1862 scheduler
1863 .run_distributed_turn(0xBB, &transport, Instant::now(), |_, _, _| {
1864 Ok::<_, Infallible>(output())
1865 })
1866 .unwrap();
1867 let pending = scheduler
1868 .run_distributed_turn(0xBB, &transport, Instant::now(), |_, _, _| {
1869 Ok::<_, Infallible>(output())
1870 })
1871 .unwrap();
1872 assert!(pending.failed.is_empty());
1873 assert_eq!(scheduler.report().failed_in_flight_work, 1);
1874 assert_eq!(scheduler.report().current_in_flight_work, 1);
1875
1876 let terminal = scheduler
1877 .run_distributed_turn(0xBB, &transport, Instant::now(), |_, _, _| {
1878 Ok::<_, Infallible>(output())
1879 })
1880 .unwrap();
1881 assert_eq!(terminal.failed.len(), 1);
1882 assert_eq!(scheduler.report().current_in_flight_work, 0);
1883 assert_eq!(
1884 scheduler.request_status(request),
1885 Some(RequestStatus::Failed)
1886 );
1887 }
1888
1889 #[test]
1890 fn distributed_cancellation_and_deadline_use_the_same_core_lifecycle() {
1891 let transport = ScriptedTransport::new(2, vec![GatherStep::default()]);
1892 let mut cancelled = scheduler();
1893 let request = RequestId::new(52);
1894 cancelled.register(request, State::default()).unwrap();
1895 cancelled.enqueue(request, 1).unwrap();
1896 cancelled
1897 .cancel_distributed(0xCC, request, &transport, Instant::now())
1898 .unwrap();
1899 assert_eq!(
1900 cancelled.request_status(request),
1901 Some(RequestStatus::Cancelled)
1902 );
1903
1904 let transport = ScriptedTransport::new(
1905 2,
1906 vec![
1907 GatherStep::default(),
1908 GatherStep::default(),
1909 GatherStep::default(),
1910 ],
1911 );
1912 let mut scheduler = scheduler();
1913 let request = RequestId::new(53);
1914 scheduler.register(request, State::default()).unwrap();
1915 scheduler
1916 .enqueue_with_deadline(
1917 request,
1918 1,
1919 Some(Instant::now().checked_sub(Duration::from_secs(1)).unwrap()),
1920 )
1921 .unwrap();
1922 scheduler
1923 .run_distributed_turn(0xDD, &transport, Instant::now(), |_, _, _| {
1924 Ok::<_, Infallible>(Output {
1925 complete: Rc::new(Cell::new(true)),
1926 fail: false,
1927 })
1928 })
1929 .unwrap();
1930 assert_eq!(
1931 scheduler.request_status(request),
1932 Some(RequestStatus::DeadlineExceeded)
1933 );
1934 }
1935
1936 #[test]
1937 fn production_telemetry_round_trips_without_backend_types() {
1938 let mut scheduler = scheduler();
1939 let request = RequestId::new(40);
1940 scheduler.register(request, State::default()).unwrap();
1941 scheduler.enqueue(request, 1).unwrap();
1942 scheduler.prepare_bounded(1, Instant::now()).unwrap();
1943 scheduler
1944 .submit_prepared(Instant::now(), |_, _, _| {
1945 Ok::<_, Infallible>(Output {
1946 complete: Rc::new(Cell::new(false)),
1947 fail: false,
1948 })
1949 })
1950 .unwrap();
1951
1952 let report = scheduler.report();
1953 let report_json = serde_json::to_string(&report).unwrap();
1954 assert_eq!(
1955 serde_json::from_str::<SchedulerReport>(&report_json).unwrap(),
1956 report
1957 );
1958
1959 let capabilities = scheduler.capabilities();
1960 assert_eq!(capabilities.observed_backends, ["mock"]);
1961 let capabilities_json = serde_json::to_string(&capabilities).unwrap();
1962 assert_eq!(
1963 serde_json::from_str::<SchedulerCapabilities>(&capabilities_json).unwrap(),
1964 capabilities
1965 );
1966 }
1967}