1use std::collections::{BTreeMap, BTreeSet};
2use std::fmt;
3use std::sync::atomic::{AtomicBool, Ordering};
4use std::sync::{Arc, Mutex};
5use std::time::{Duration, SystemTime};
6
7use oxide_batch_repository::{
8 PartitionMutationError, aggregate_partition_parent, map_partition_aggregation,
9 recovered_execution,
10};
11
12use crate::{
13 ActorRef, BatchStatus, CursorKey, DefinitionDescriptor, DefinitionIdentity, DefinitionRevision,
14 DefinitionUpgrade, DurableStateKind, ExecutionCounts, ExecutionMetadata, ExecutionTimestamps,
15 ExecutionVersion, ExitStatus, ExplorerError, ExplorerQuery, ExplorerRepository, FlowDecision,
16 FlowDecisionId, FlowDecisionRequest, FlowStepState, FlowTransitionKind, IdentifierKind,
17 JobExecution, JobExecutionId, JobExecutionProjection, JobInstance, JobInstanceId,
18 JobInstanceKey, JobInstanceProjection, JobName, LifecycleError, LifecycleTransition,
19 MAX_PARTITIONS, NodeId, OperationId, OperatorAction, OperatorRecord, OperatorRecordDraft,
20 OperatorRequestId, OwnerObservation, OwnerToken, ParameterDescriptor, PartitionPlanEntry,
21 PartitionResult, PurgeCandidate, PurgeCounts, PurgePlan, PurgePlanRequest, PurgeSurvey,
22 QueryWindow, ReasonCode, RecoveryDecisionId, RecoveryRepository, RecoverySnapshot,
23 RecoveryStepEvidence, RetentionAction, RetentionActionId, RetentionHold, RetentionRecord,
24 RetentionRecordDraft, StartLimit, StateEnvelopeDescriptor, StepExecution, StepExecutionId,
25 StepExecutionProjection, StepName, StepPartition, StepPartitionId, StepPartitionProjection,
26};
27use crate::{
28 BoxFuture, Clock, IdGenerator, JobInstanceSelection, JobRepository, RecoveryDecision,
29 RecoveryRequest, RecoveryResult, RepositoryCapability, RepositoryDescriptor, RepositoryError,
30 RepositoryUnitOfWork,
31};
32
33#[derive(Clone)]
40pub struct InMemoryJobRepository {
41 state: Arc<Mutex<MemoryState>>,
42 clock: Arc<dyn Clock>,
43 ids: Arc<dyn IdGenerator>,
44 fail_next_partition_aggregate_commit: Arc<AtomicBool>,
45}
46
47impl InMemoryJobRepository {
48 #[must_use]
50 pub fn new(clock: Arc<dyn Clock>, ids: Arc<dyn IdGenerator>) -> Self {
51 Self {
52 state: Arc::new(Mutex::new(MemoryState::default())),
53 clock,
54 ids,
55 fail_next_partition_aggregate_commit: Arc::new(AtomicBool::new(false)),
56 }
57 }
58
59 pub fn inject_next_partition_aggregate_commit_unknown(&self) {
64 self.fail_next_partition_aggregate_commit
65 .store(true, Ordering::Release);
66 }
67}
68
69impl fmt::Debug for InMemoryJobRepository {
70 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
71 let mut debug = formatter.debug_struct("InMemoryJobRepository");
72 match self.state.lock() {
73 Ok(state) => debug
74 .field("revision", &state.revision)
75 .field("job_instance_count", &state.instances_by_id.len())
76 .field("job_execution_count", &state.job_executions.len())
77 .field("step_execution_count", &state.step_executions.len()),
78 Err(_) => debug.field("state", &"<poisoned>"),
79 };
80 debug.finish_non_exhaustive()
81 }
82}
83
84impl JobRepository for InMemoryJobRepository {
85 fn connection_capacity(&self) -> u32 {
86 u32::from(crate::MAX_PARTITION_WORKERS) + 1
87 }
88
89 fn descriptor(&self) -> RepositoryDescriptor {
93 RepositoryDescriptor::new(
94 0,
95 [
96 RepositoryCapability::ExecutionOwnership,
97 RepositoryCapability::InstanceHolds,
98 RepositoryCapability::OperatorRequests,
99 RepositoryCapability::RetentionPurge,
100 RepositoryCapability::StepPartitions,
101 RepositoryCapability::StopRequests,
102 ],
103 )
104 }
105
106 fn begin<'a>(
107 &'a self,
108 ) -> BoxFuture<'a, Result<Box<dyn RepositoryUnitOfWork + 'a>, RepositoryError>> {
109 Box::pin(async move {
110 let snapshot = self
111 .state
112 .lock()
113 .map_err(|_| RepositoryError::Unavailable)?
114 .clone();
115 let base_revision = snapshot.revision;
116 Ok(Box::new(InMemoryUnitOfWork {
117 repository: self,
118 base_revision,
119 staged: snapshot,
120 definition_override: None,
121 created_partition_plans: BTreeSet::new(),
122 aggregated_partition_parent: false,
123 }) as Box<dyn RepositoryUnitOfWork + 'a>)
124 })
125 }
126}
127
128#[derive(Clone, Debug, Default)]
129struct MemoryState {
130 revision: u64,
131 instances_by_key: BTreeMap<JobInstanceKey, JobInstanceId>,
132 instances_by_id: BTreeMap<JobInstanceId, JobInstance>,
133 job_executions: BTreeMap<JobExecutionId, JobExecution>,
134 job_executions_by_instance: BTreeMap<JobInstanceId, Vec<JobExecutionId>>,
135 step_executions: BTreeMap<StepExecutionId, StepExecution>,
136 step_executions_by_job: BTreeMap<JobExecutionId, Vec<StepExecutionId>>,
137 step_logical_ids: BTreeMap<StepExecutionId, NodeId>,
138 step_partitions: BTreeMap<StepPartitionId, StepPartition>,
139 step_partitions_by_step: BTreeMap<StepExecutionId, Vec<StepPartitionId>>,
140 flow_decisions: BTreeMap<FlowDecisionId, FlowDecision>,
141 flow_decisions_by_job: BTreeMap<JobExecutionId, Vec<FlowDecisionId>>,
142 recovery_decisions: BTreeMap<JobExecutionId, Vec<RecoveryDecision>>,
143 definitions: BTreeMap<(JobName, DefinitionRevision), DefinitionIdentity>,
144 execution_definitions: BTreeMap<JobExecutionId, DefinitionIdentity>,
145 definition_upgrades: BTreeMap<(JobName, [u8; 32], [u8; 32]), DefinitionUpgrade>,
146 job_name_order: BTreeMap<JobName, u64>,
147 instance_created_at: BTreeMap<JobInstanceId, SystemTime>,
148 execution_updated_at: BTreeMap<JobExecutionId, SystemTime>,
149 holds: BTreeMap<JobInstanceId, RetentionHold>,
150 stop_requests: BTreeMap<JobExecutionId, SystemTime>,
151 owner_tokens: BTreeMap<JobExecutionId, OwnerToken>,
152 operator_requests: BTreeMap<OperatorRequestId, OperatorRecord>,
153 operator_request_keys: BTreeMap<(&'static str, String), OperatorRequestId>,
154 retention_actions: BTreeMap<RetentionActionId, RetentionRecord>,
155 retention_action_keys: BTreeMap<(&'static str, String), RetentionActionId>,
156}
157
158impl MemoryState {
159 fn register_job_name(&mut self, job_name: &JobName) {
160 if self.job_name_order.contains_key(job_name) {
161 return;
162 }
163 let next = self
164 .job_name_order
165 .values()
166 .copied()
167 .max()
168 .map_or(1, |value| value.saturating_add(1));
169 self.job_name_order.insert(job_name.clone(), next);
170 }
171
172 fn attempt_of(&self, execution: &JobExecution) -> u32 {
173 self.job_executions_by_instance
174 .get(&execution.job_instance_id())
175 .and_then(|executions| {
176 executions
177 .iter()
178 .position(|candidate| *candidate == execution.id())
179 })
180 .and_then(|position| u32::try_from(position.saturating_add(1)).ok())
181 .unwrap_or(1)
182 }
183
184 fn job_name_of(&self, instance_id: JobInstanceId) -> Option<JobName> {
185 self.instances_by_id
186 .get(&instance_id)
187 .map(|instance| instance.key().job_name().clone())
188 }
189
190 fn job_execution_projection(
191 &self,
192 execution: &JobExecution,
193 ) -> Result<JobExecutionProjection, ExplorerError> {
194 let job_name =
195 self.job_name_of(execution.job_instance_id())
196 .ok_or(ExplorerError::Repository(
197 RepositoryError::JobInstanceNotFound {
198 id: execution.job_instance_id(),
199 },
200 ))?;
201 let definition = self
202 .execution_definitions
203 .get(&execution.id())
204 .map(|definition| {
205 DefinitionDescriptor::new(
206 definition.revision().clone(),
207 definition.manifest_format(),
208 *definition.manifest_digest(),
209 )
210 });
211 Ok(JobExecutionProjection::new(
212 execution.id(),
213 execution.job_instance_id(),
214 job_name,
215 self.attempt_of(execution),
216 execution.metadata().status(),
217 execution.metadata().exit_status().clone(),
218 execution.metadata().counts(),
219 execution.version(),
220 execution.metadata().timestamps(),
221 self.execution_updated_at
222 .get(&execution.id())
223 .copied()
224 .unwrap_or_else(|| updated_at(execution)),
225 execution.metadata().failure(),
226 definition,
227 None,
228 self.stop_requests.get(&execution.id()).copied(),
229 false,
230 ))
231 }
232
233 fn job_instance_projection(&self, instance: &JobInstance) -> JobInstanceProjection {
234 let key = instance.key();
235 let parameters = key
236 .identifying_fields()
237 .map(|(name, kind)| ParameterDescriptor::new(name.clone(), kind, true))
238 .collect();
239 JobInstanceProjection::new(
240 instance.id(),
241 key.job_name().clone(),
242 key.digest(),
243 parameters,
244 self.instance_created_at.get(&instance.id()).copied(),
245 self.holds.get(&instance.id()).cloned(),
246 )
247 }
248
249 fn step_execution_projection(&self, execution: &StepExecution) -> StepExecutionProjection {
250 StepExecutionProjection::new(
251 execution.id(),
252 execution.job_execution_id(),
253 execution.step_name().clone(),
254 self.step_logical_ids.get(&execution.id()).cloned(),
255 execution.metadata().status(),
256 execution.metadata().exit_status().clone(),
257 execution.metadata().counts(),
258 execution.version(),
259 execution.metadata().timestamps(),
260 execution.metadata().failure(),
261 None,
262 None,
263 )
264 }
265}
266
267fn updated_at(execution: &JobExecution) -> SystemTime {
268 let timestamps = execution.metadata().timestamps();
269 timestamps
270 .ended_at()
271 .or_else(|| timestamps.started_at())
272 .unwrap_or_else(|| timestamps.created_at())
273}
274
275struct InMemoryUnitOfWork<'repository> {
276 repository: &'repository InMemoryJobRepository,
277 base_revision: u64,
278 staged: MemoryState,
279 definition_override: Option<DefinitionIdentity>,
280 created_partition_plans: BTreeSet<StepExecutionId>,
281 aggregated_partition_parent: bool,
282}
283
284impl InMemoryUnitOfWork<'_> {
285 fn create_starting_metadata(&self) -> Result<ExecutionMetadata, RepositoryError> {
286 let created_at = self.repository.clock.now();
287 let timestamps = ExecutionTimestamps::new(created_at, None, None)?;
288 ExecutionMetadata::new(
289 BatchStatus::Starting,
290 ExitStatus::unknown(),
291 timestamps,
292 ExecutionCounts::default(),
293 None,
294 )
295 .map_err(RepositoryError::from)
296 }
297
298 fn latest_job_execution(
299 &self,
300 instance_id: JobInstanceId,
301 ) -> Result<Option<&JobExecution>, RepositoryError> {
302 let Some(execution_ids) = self.staged.job_executions_by_instance.get(&instance_id) else {
303 if self.staged.instances_by_id.contains_key(&instance_id) {
304 return Ok(None);
305 }
306 return Err(RepositoryError::JobInstanceNotFound { id: instance_id });
307 };
308 Ok(execution_ids
309 .last()
310 .and_then(|id| self.staged.job_executions.get(id)))
311 }
312
313 fn ensure_definition(
314 &mut self,
315 job_name: &JobName,
316 definition: &DefinitionIdentity,
317 ) -> Result<(), RepositoryError> {
318 if let Some(actual) = definition.job_name()
319 && actual != job_name
320 {
321 return Err(RepositoryError::DefinitionJobMismatch {
322 expected: job_name.clone(),
323 actual: actual.clone(),
324 });
325 }
326 let key = (job_name.clone(), definition.revision().clone());
327 if let Some(existing) = self.staged.definitions.get(&key) {
328 if existing.manifest_digest() != definition.manifest_digest() {
329 return Err(RepositoryError::DefinitionDrift {
330 job_name: job_name.clone(),
331 revision: definition.revision().clone(),
332 });
333 }
334 return Ok(());
335 }
336 self.staged.register_job_name(job_name);
337 self.staged.definitions.insert(key, definition.clone());
338 Ok(())
339 }
340
341 fn instance_for_execution(
342 &self,
343 execution_id: JobExecutionId,
344 ) -> Result<JobInstanceId, RepositoryError> {
345 self.staged
346 .job_executions
347 .get(&execution_id)
348 .map(JobExecution::job_instance_id)
349 .ok_or(RepositoryError::JobExecutionNotFound { id: execution_id })
350 }
351
352 fn next_recovery_decision_id(&self) -> Result<RecoveryDecisionId, RepositoryError> {
353 let next = self
354 .staged
355 .recovery_decisions
356 .values()
357 .flatten()
358 .map(|decision| decision.id().get())
359 .max()
360 .map_or(1, |id| id.checked_add(1).unwrap_or(0));
361 RecoveryDecisionId::new(next).map_err(RepositoryError::from)
362 }
363
364 fn next_operator_request_id(&self) -> Result<OperatorRequestId, RepositoryError> {
365 let next = self
366 .staged
367 .operator_requests
368 .keys()
369 .next_back()
370 .map_or(1, |id| id.get().checked_add(1).unwrap_or(0));
371 OperatorRequestId::new(next).map_err(RepositoryError::from)
372 }
373
374 fn next_retention_action_id(&self) -> Result<RetentionActionId, RepositoryError> {
375 let next = self
376 .staged
377 .retention_actions
378 .keys()
379 .next_back()
380 .map_or(1, |id| id.get().checked_add(1).unwrap_or(0));
381 RetentionActionId::new(next).map_err(RepositoryError::from)
382 }
383
384 fn next_step_partition_id(&self) -> Result<StepPartitionId, RepositoryError> {
385 let next = self
386 .staged
387 .step_partitions
388 .keys()
389 .next_back()
390 .map_or(1, |id| id.get().checked_add(1).unwrap_or(0));
391 StepPartitionId::new(next).map_err(RepositoryError::from)
392 }
393
394 fn remove_step_execution(&mut self, step_execution_id: StepExecutionId) {
395 for partition_id in self
396 .staged
397 .step_partitions_by_step
398 .remove(&step_execution_id)
399 .unwrap_or_default()
400 {
401 self.staged.step_partitions.remove(&partition_id);
402 }
403 self.staged.step_executions.remove(&step_execution_id);
404 self.staged.step_logical_ids.remove(&step_execution_id);
405 }
406
407 fn purge_eligible(&self, request: &PurgePlanRequest, now: SystemTime) -> Vec<PurgeCandidate> {
408 let mut candidates = Vec::new();
409 for (instance_id, instance) in &self.staged.instances_by_id {
410 if instance.key().job_name() != request.job_name()
411 || self.staged.holds.contains_key(instance_id)
412 {
413 continue;
414 }
415 let executions = self
416 .staged
417 .job_executions_by_instance
418 .get(instance_id)
419 .into_iter()
420 .flatten()
421 .filter_map(|id| self.staged.job_executions.get(id))
422 .collect::<Vec<_>>();
423 if executions
424 .iter()
425 .any(|execution| !execution.metadata().status().is_finished())
426 {
427 continue;
428 }
429 for execution in executions {
430 let status = execution.metadata().status();
431 if !request.statuses().contains(status) {
432 continue;
433 }
434 let age = now
435 .duration_since(updated_at(execution))
436 .unwrap_or(Duration::ZERO);
437 if age < request.minimum_age() {
438 continue;
439 }
440 candidates.push(PurgeCandidate::new(
441 *instance_id,
442 execution.id(),
443 execution.version(),
444 ));
445 }
446 }
447 candidates.sort_unstable();
448 candidates.truncate(usize::try_from(request.batch().get()).unwrap_or(usize::MAX));
449 candidates
450 }
451
452 fn purge_counts(&self, candidates: &[PurgeCandidate]) -> PurgeCounts {
453 let mut flow_decisions = 0_u64;
454 let mut recovery_decisions = 0_u64;
455 let mut operator_requests = 0_u64;
456 let mut step_partitions = 0_u64;
457 let mut step_executions = 0_u64;
458 let mut instances = BTreeMap::new();
459 for candidate in candidates {
460 let execution_id = candidate.job_execution_id();
461 flow_decisions = flow_decisions.saturating_add(count_of(
462 self.staged.flow_decisions_by_job.get(&execution_id),
463 ));
464 recovery_decisions = recovery_decisions
465 .saturating_add(count_of(self.staged.recovery_decisions.get(&execution_id)));
466 step_executions = step_executions.saturating_add(count_of(
467 self.staged.step_executions_by_job.get(&execution_id),
468 ));
469 step_partitions = step_partitions.saturating_add(
470 self.staged
471 .step_executions_by_job
472 .get(&execution_id)
473 .into_iter()
474 .flatten()
475 .map(|id| count_of(self.staged.step_partitions_by_step.get(id)))
476 .fold(0_u64, u64::saturating_add),
477 );
478 operator_requests = operator_requests.saturating_add(
479 u64::try_from(
480 self.staged
481 .operator_requests
482 .values()
483 .filter(|record| record.job_execution_id() == Some(execution_id))
484 .count(),
485 )
486 .unwrap_or(u64::MAX),
487 );
488 *instances
489 .entry(candidate.job_instance_id())
490 .or_insert(0_u64) += 1;
491 }
492 let job_instances = instances
493 .iter()
494 .filter(|(instance_id, purged)| {
495 count_of(self.staged.job_executions_by_instance.get(instance_id)) == **purged
496 })
497 .count();
498 PurgeCounts::new(
499 flow_decisions,
500 recovery_decisions,
501 operator_requests,
502 step_partitions,
503 step_executions,
504 u64::try_from(candidates.len()).unwrap_or(u64::MAX),
505 u64::try_from(job_instances).unwrap_or(u64::MAX),
506 )
507 }
508
509 fn next_flow_decision_id(&self) -> Result<FlowDecisionId, RepositoryError> {
510 let next = self
511 .staged
512 .flow_decisions
513 .keys()
514 .next_back()
515 .map_or(1, |id| id.get().checked_add(1).unwrap_or(0));
516 FlowDecisionId::new(next).map_err(RepositoryError::from)
517 }
518
519 fn latest_flow_step_snapshot(
520 &self,
521 instance_id: JobInstanceId,
522 node_id: &NodeId,
523 ) -> Result<Option<FlowStepState>, RepositoryError> {
524 let executions = self
525 .staged
526 .job_executions_by_instance
527 .get(&instance_id)
528 .ok_or(RepositoryError::JobInstanceNotFound { id: instance_id })?;
529 for execution_id in executions.iter().rev() {
530 let step_ids = self
531 .staged
532 .step_executions_by_job
533 .get(execution_id)
534 .into_iter()
535 .flatten();
536 for step_id in step_ids.rev() {
537 if self.staged.step_logical_ids.get(step_id) == Some(node_id) {
538 let execution = self
539 .staged
540 .step_executions
541 .get(step_id)
542 .cloned()
543 .ok_or(RepositoryError::FlowStateCorrupt)?;
544 return Ok(Some(FlowStepState::new(node_id.clone(), execution, None)));
545 }
546 }
547 }
548 Ok(None)
549 }
550}
551
552impl RepositoryUnitOfWork for InMemoryUnitOfWork<'_> {
553 fn register_definition_upgrade<'a>(
554 &'a mut self,
555 job_name: &'a JobName,
556 upgrade: &'a DefinitionUpgrade,
557 ) -> BoxFuture<'a, Result<(), RepositoryError>> {
558 Box::pin(async move {
559 self.ensure_definition(job_name, upgrade.from())?;
560 self.ensure_definition(job_name, upgrade.to())?;
561 let key = (
562 job_name.clone(),
563 *upgrade.from().manifest_digest(),
564 *upgrade.to().manifest_digest(),
565 );
566 if let Some(existing) = self.staged.definition_upgrades.get(&key) {
567 if existing != upgrade {
568 return Err(RepositoryError::DefinitionUpgradeConflict {
569 job_name: job_name.clone(),
570 });
571 }
572 return Ok(());
573 }
574 self.staged.definition_upgrades.insert(key, upgrade.clone());
575 Ok(())
576 })
577 }
578
579 fn select_or_create_job_instance<'a>(
580 &'a mut self,
581 key: &'a JobInstanceKey,
582 ) -> BoxFuture<'a, Result<JobInstanceSelection, RepositoryError>> {
583 Box::pin(async move {
584 if let Some(id) = self.staged.instances_by_key.get(key) {
585 let instance = self
586 .staged
587 .instances_by_id
588 .get(id)
589 .cloned()
590 .ok_or(RepositoryError::JobInstanceNotFound { id: *id })?;
591 return Ok(JobInstanceSelection::Existing(instance));
592 }
593
594 let id = self.repository.ids.next_job_instance_id()?;
595 if self.staged.instances_by_id.contains_key(&id) {
596 return Err(RepositoryError::DuplicateIdentifier {
597 kind: IdentifierKind::JobInstance,
598 value: id.get(),
599 });
600 }
601 let instance = JobInstance::new(id, key.clone());
602 let created_at = self.repository.clock.now();
603 self.staged.instances_by_key.insert(key.clone(), id);
604 self.staged.instances_by_id.insert(id, instance.clone());
605 self.staged.instance_created_at.insert(id, created_at);
606 self.staged.register_job_name(key.job_name());
607 self.staged
608 .job_executions_by_instance
609 .insert(id, Vec::new());
610 Ok(JobInstanceSelection::Created(instance))
611 })
612 }
613
614 fn create_job_execution(
615 &mut self,
616 job_instance_id: JobInstanceId,
617 ) -> BoxFuture<'_, Result<JobExecution, RepositoryError>> {
618 Box::pin(async move {
619 let definition = self
620 .definition_override
621 .take()
622 .unwrap_or_else(DefinitionIdentity::legacy);
623 let job_name = self
624 .staged
625 .instances_by_id
626 .get(&job_instance_id)
627 .ok_or(RepositoryError::JobInstanceNotFound {
628 id: job_instance_id,
629 })?
630 .key()
631 .job_name()
632 .clone();
633 self.ensure_definition(&job_name, &definition)?;
634 if let Some(latest) = self.latest_job_execution(job_instance_id)? {
635 match latest.metadata().status() {
636 BatchStatus::Stopped | BatchStatus::Failed => {}
637 BatchStatus::Completed => {
638 return Err(RepositoryError::CompletedInstance {
639 id: job_instance_id,
640 });
641 }
642 BatchStatus::Abandoned => {
643 return Err(RepositoryError::AbandonedInstance {
644 id: job_instance_id,
645 });
646 }
647 status => {
648 return Err(RepositoryError::ExecutionAlreadyActive {
649 instance_id: job_instance_id,
650 execution_id: latest.id(),
651 status,
652 });
653 }
654 }
655 let previous_definition =
656 self.staged.execution_definitions.get(&latest.id()).ok_or(
657 RepositoryError::IncompatibleDefinition {
658 instance_id: job_instance_id,
659 },
660 )?;
661 if previous_definition.manifest_digest() != definition.manifest_digest()
662 && !self.staged.definition_upgrades.contains_key(&(
663 job_name,
664 *previous_definition.manifest_digest(),
665 *definition.manifest_digest(),
666 ))
667 {
668 return Err(RepositoryError::IncompatibleDefinition {
669 instance_id: job_instance_id,
670 });
671 }
672 }
673
674 let id = self.repository.ids.next_job_execution_id()?;
675 if self.staged.job_executions.contains_key(&id) {
676 return Err(RepositoryError::DuplicateIdentifier {
677 kind: IdentifierKind::JobExecution,
678 value: id.get(),
679 });
680 }
681 let execution =
682 JobExecution::new(id, job_instance_id, self.create_starting_metadata()?);
683 self.staged
684 .execution_updated_at
685 .insert(id, execution.metadata().timestamps().created_at());
686 self.staged.job_executions.insert(id, execution.clone());
687 self.staged.execution_definitions.insert(id, definition);
688 self.staged
689 .job_executions_by_instance
690 .get_mut(&job_instance_id)
691 .ok_or(RepositoryError::JobInstanceNotFound {
692 id: job_instance_id,
693 })?
694 .push(id);
695 Ok(execution)
696 })
697 }
698
699 fn create_job_execution_with_definition<'a>(
700 &'a mut self,
701 job_instance_id: JobInstanceId,
702 definition: &'a DefinitionIdentity,
703 ) -> BoxFuture<'a, Result<JobExecution, RepositoryError>> {
704 Box::pin(async move {
705 self.definition_override = Some(definition.clone());
706 self.create_job_execution(job_instance_id).await
707 })
708 }
709
710 fn create_step_execution<'a>(
711 &'a mut self,
712 job_execution_id: JobExecutionId,
713 step_name: &'a StepName,
714 ) -> BoxFuture<'a, Result<StepExecution, RepositoryError>> {
715 Box::pin(async move {
716 let node_id =
717 NodeId::new(step_name.as_str()).map_err(|_| RepositoryError::FlowStateCorrupt)?;
718 self.create_flow_step_execution(
719 job_execution_id,
720 step_name,
721 &node_id,
722 StartLimit::UNRESTRICTED,
723 )
724 .await
725 })
726 }
727
728 fn create_flow_step_execution<'a>(
729 &'a mut self,
730 job_execution_id: JobExecutionId,
731 step_name: &'a StepName,
732 node_id: &'a NodeId,
733 start_limit: StartLimit,
734 ) -> BoxFuture<'a, Result<StepExecution, RepositoryError>> {
735 Box::pin(async move {
736 let instance_id = self.instance_for_execution(job_execution_id)?;
737 let historical_starts = self
738 .staged
739 .job_executions_by_instance
740 .get(&instance_id)
741 .into_iter()
742 .flatten()
743 .flat_map(|execution_id| {
744 self.staged
745 .step_executions_by_job
746 .get(execution_id)
747 .into_iter()
748 .flatten()
749 })
750 .filter(|step_id| self.staged.step_logical_ids.get(step_id) == Some(node_id))
751 .count();
752 if u64::try_from(historical_starts).unwrap_or(u64::MAX) >= u64::from(start_limit.get())
753 {
754 return Err(RepositoryError::StartLimitExceeded {
755 instance_id,
756 node_id: node_id.clone(),
757 limit: start_limit,
758 });
759 }
760 let id = self.repository.ids.next_step_execution_id()?;
761 if self.staged.step_executions.contains_key(&id) {
762 return Err(RepositoryError::DuplicateIdentifier {
763 kind: IdentifierKind::StepExecution,
764 value: id.get(),
765 });
766 }
767 let counts = self
768 .latest_flow_step_snapshot(instance_id, node_id)?
769 .map_or_else(ExecutionCounts::default, |state| {
770 state.execution().metadata().counts()
771 });
772 let created_at = self.repository.clock.now();
773 let metadata = ExecutionMetadata::new(
774 BatchStatus::Starting,
775 ExitStatus::unknown(),
776 ExecutionTimestamps::new(created_at, None, None)?,
777 counts,
778 None,
779 )?;
780 let execution = StepExecution::new(id, job_execution_id, step_name.clone(), metadata);
781 self.staged.step_executions.insert(id, execution.clone());
782 self.staged.step_logical_ids.insert(id, node_id.clone());
783 self.staged
784 .step_executions_by_job
785 .entry(job_execution_id)
786 .or_default()
787 .push(id);
788 Ok(execution)
789 })
790 }
791
792 fn transition_job_execution(
793 &mut self,
794 id: JobExecutionId,
795 expected_version: ExecutionVersion,
796 transition: LifecycleTransition,
797 ) -> BoxFuture<'_, Result<JobExecution, RepositoryError>> {
798 Box::pin(async move {
799 let execution = self
800 .staged
801 .job_executions
802 .get_mut(&id)
803 .ok_or(RepositoryError::JobExecutionNotFound { id })?;
804 execution.transition(expected_version, transition)?;
805 self.staged
806 .execution_updated_at
807 .insert(id, transition.transitioned_at());
808 Ok(execution.clone())
809 })
810 }
811
812 fn enrich_job_exit_status<'a>(
813 &'a mut self,
814 id: JobExecutionId,
815 expected_version: ExecutionVersion,
816 exit_status: &'a ExitStatus,
817 ) -> BoxFuture<'a, Result<JobExecution, RepositoryError>> {
818 Box::pin(async move {
819 let execution = self
820 .staged
821 .job_executions
822 .get_mut(&id)
823 .ok_or(RepositoryError::JobExecutionNotFound { id })?;
824 execution.enrich_exit_status(expected_version, exit_status.clone())?;
825 Ok(execution.clone())
826 })
827 }
828
829 fn transition_step_execution(
830 &mut self,
831 id: StepExecutionId,
832 expected_version: ExecutionVersion,
833 transition: LifecycleTransition,
834 ) -> BoxFuture<'_, Result<StepExecution, RepositoryError>> {
835 Box::pin(async move {
836 let execution = self
837 .staged
838 .step_executions
839 .get_mut(&id)
840 .ok_or(RepositoryError::StepExecutionNotFound { id })?;
841 execution.transition(expected_version, transition)?;
842 Ok(execution.clone())
843 })
844 }
845
846 fn enrich_step_exit_status<'a>(
847 &'a mut self,
848 id: StepExecutionId,
849 expected_version: ExecutionVersion,
850 exit_status: &'a ExitStatus,
851 ) -> BoxFuture<'a, Result<StepExecution, RepositoryError>> {
852 Box::pin(async move {
853 let execution = self
854 .staged
855 .step_executions
856 .get_mut(&id)
857 .ok_or(RepositoryError::StepExecutionNotFound { id })?;
858 execution.enrich_exit_status(expected_version, exit_status.clone())?;
859 Ok(execution.clone())
860 })
861 }
862
863 fn find_job_instance<'a>(
864 &'a mut self,
865 key: &'a JobInstanceKey,
866 ) -> BoxFuture<'a, Result<Option<JobInstance>, RepositoryError>> {
867 Box::pin(async move {
868 Ok(self
869 .staged
870 .instances_by_key
871 .get(key)
872 .and_then(|id| self.staged.instances_by_id.get(id))
873 .cloned())
874 })
875 }
876
877 fn get_job_execution(
878 &mut self,
879 id: JobExecutionId,
880 ) -> BoxFuture<'_, Result<Option<JobExecution>, RepositoryError>> {
881 Box::pin(async move { Ok(self.staged.job_executions.get(&id).cloned()) })
882 }
883
884 fn get_job_instance(
885 &mut self,
886 id: JobInstanceId,
887 ) -> BoxFuture<'_, Result<Option<JobInstance>, RepositoryError>> {
888 Box::pin(async move { Ok(self.staged.instances_by_id.get(&id).cloned()) })
889 }
890
891 fn job_executions(
892 &mut self,
893 job_instance_id: JobInstanceId,
894 ) -> BoxFuture<'_, Result<Vec<JobExecution>, RepositoryError>> {
895 Box::pin(async move {
896 if !self.staged.instances_by_id.contains_key(&job_instance_id) {
897 return Err(RepositoryError::JobInstanceNotFound {
898 id: job_instance_id,
899 });
900 }
901 Ok(self
902 .staged
903 .job_executions_by_instance
904 .get(&job_instance_id)
905 .into_iter()
906 .flatten()
907 .filter_map(|id| self.staged.job_executions.get(id))
908 .cloned()
909 .collect())
910 })
911 }
912
913 fn get_step_execution(
914 &mut self,
915 id: StepExecutionId,
916 ) -> BoxFuture<'_, Result<Option<StepExecution>, RepositoryError>> {
917 Box::pin(async move { Ok(self.staged.step_executions.get(&id).cloned()) })
918 }
919
920 fn step_executions(
921 &mut self,
922 job_execution_id: JobExecutionId,
923 ) -> BoxFuture<'_, Result<Vec<StepExecution>, RepositoryError>> {
924 Box::pin(async move {
925 if !self.staged.job_executions.contains_key(&job_execution_id) {
926 return Err(RepositoryError::JobExecutionNotFound {
927 id: job_execution_id,
928 });
929 }
930 Ok(self
931 .staged
932 .step_executions_by_job
933 .get(&job_execution_id)
934 .into_iter()
935 .flatten()
936 .filter_map(|id| self.staged.step_executions.get(id))
937 .cloned()
938 .collect())
939 })
940 }
941
942 fn latest_flow_step<'a>(
943 &'a mut self,
944 job_instance_id: JobInstanceId,
945 node_id: &'a NodeId,
946 ) -> BoxFuture<'a, Result<Option<FlowStepState>, RepositoryError>> {
947 Box::pin(async move { self.latest_flow_step_snapshot(job_instance_id, node_id) })
948 }
949
950 fn append_flow_decision<'a>(
951 &'a mut self,
952 request: &'a FlowDecisionRequest,
953 ) -> BoxFuture<'a, Result<FlowDecision, RepositoryError>> {
954 Box::pin(async move {
955 let instance_id = self.instance_for_execution(request.job_execution_id())?;
956 let definition = self
957 .staged
958 .execution_definitions
959 .get(&request.job_execution_id())
960 .ok_or(RepositoryError::FlowStateCorrupt)?;
961 if definition.manifest_digest() != request.plan_fingerprint() {
962 return Err(RepositoryError::FlowStateCorrupt);
963 }
964 let manifest = serde_json::from_slice(definition.canonical_manifest())
965 .map_err(|_| RepositoryError::FlowStateCorrupt)?;
966 if !crate::flow::decision_matches_manifest(&manifest, request) {
967 return Err(RepositoryError::FlowStateCorrupt);
968 }
969 let existing = self
970 .staged
971 .flow_decisions_by_job
972 .get(&request.job_execution_id())
973 .cloned()
974 .unwrap_or_default();
975 let expected_sequence = u64::try_from(existing.len())
976 .ok()
977 .and_then(|value| value.checked_add(1))
978 .ok_or(RepositoryError::FlowStateCorrupt)?;
979 if request.sequence().get() != expected_sequence
980 || existing.iter().any(|id| {
981 self.staged.flow_decisions.get(id).is_some_and(|decision| {
982 decision.source_node_id() == request.source_node_id()
983 })
984 })
985 {
986 return Err(RepositoryError::ConcurrentModification);
987 }
988 if let Some(step_id) = request.source_step_execution_id() {
989 let step = self
990 .staged
991 .step_executions
992 .get(&step_id)
993 .ok_or(RepositoryError::FlowStateCorrupt)?;
994 if self.instance_for_execution(step.job_execution_id())? != instance_id
995 || self.staged.step_logical_ids.get(&step_id) != Some(request.source_node_id())
996 {
997 return Err(RepositoryError::FlowStateCorrupt);
998 }
999 } else if !matches!(
1000 request.kind(),
1001 FlowTransitionKind::Decider | FlowTransitionKind::SplitAggregate
1002 ) {
1003 return Err(RepositoryError::FlowStateCorrupt);
1004 }
1005 if let Some(reused_id) = request.reused_decision_id() {
1006 let reused = self
1007 .staged
1008 .flow_decisions
1009 .get(&reused_id)
1010 .ok_or(RepositoryError::FlowStateCorrupt)?;
1011 let reused_instance = self.instance_for_execution(reused.job_execution_id())?;
1012 if reused_instance != instance_id
1013 || reused.source_node_id() != request.source_node_id()
1014 || reused.plan_fingerprint() != request.plan_fingerprint()
1015 || reused.input_digest() != request.input_digest()
1016 || reused.observed_outcome() != request.observed_outcome()
1017 || reused.target() != request.target()
1018 {
1019 return Err(RepositoryError::FlowStateCorrupt);
1020 }
1021 }
1022 let id = self.next_flow_decision_id()?;
1023 let decision = FlowDecision::new(
1024 id,
1025 request.job_execution_id(),
1026 request.sequence(),
1027 request.source_node_id().clone(),
1028 request.source_step_execution_id(),
1029 request.kind(),
1030 request.observed_outcome().clone(),
1031 request.target().clone(),
1032 *request.plan_fingerprint(),
1033 *request.input_digest(),
1034 request.reused_decision_id(),
1035 request.decided_at(),
1036 );
1037 self.staged.flow_decisions.insert(id, decision.clone());
1038 self.staged
1039 .flow_decisions_by_job
1040 .entry(request.job_execution_id())
1041 .or_default()
1042 .push(id);
1043 Ok(decision)
1044 })
1045 }
1046
1047 fn find_reusable_flow_decision<'a>(
1048 &'a mut self,
1049 job_instance_id: JobInstanceId,
1050 node_id: &'a NodeId,
1051 plan_fingerprint: &'a [u8; 32],
1052 input_digest: &'a [u8; 32],
1053 kind: FlowTransitionKind,
1054 ) -> BoxFuture<'a, Result<Option<FlowDecision>, RepositoryError>> {
1055 Box::pin(async move {
1056 let executions = self
1057 .staged
1058 .job_executions_by_instance
1059 .get(&job_instance_id)
1060 .ok_or(RepositoryError::JobInstanceNotFound {
1061 id: job_instance_id,
1062 })?;
1063 for execution_id in executions.iter().rev() {
1064 for decision_id in self
1065 .staged
1066 .flow_decisions_by_job
1067 .get(execution_id)
1068 .into_iter()
1069 .flatten()
1070 .rev()
1071 {
1072 let decision = self
1073 .staged
1074 .flow_decisions
1075 .get(decision_id)
1076 .ok_or(RepositoryError::FlowStateCorrupt)?;
1077 if decision.source_node_id() == node_id
1078 && decision.plan_fingerprint() == plan_fingerprint
1079 && decision.input_digest() == input_digest
1080 && decision.kind() == kind
1081 {
1082 return Ok(Some(decision.clone()));
1083 }
1084 }
1085 }
1086 Ok(None)
1087 })
1088 }
1089
1090 fn flow_decisions(
1091 &mut self,
1092 job_execution_id: JobExecutionId,
1093 ) -> BoxFuture<'_, Result<Vec<FlowDecision>, RepositoryError>> {
1094 Box::pin(async move {
1095 if !self.staged.job_executions.contains_key(&job_execution_id) {
1096 return Err(RepositoryError::JobExecutionNotFound {
1097 id: job_execution_id,
1098 });
1099 }
1100 self.staged
1101 .flow_decisions_by_job
1102 .get(&job_execution_id)
1103 .into_iter()
1104 .flatten()
1105 .map(|id| {
1106 self.staged
1107 .flow_decisions
1108 .get(id)
1109 .cloned()
1110 .ok_or(RepositoryError::FlowStateCorrupt)
1111 })
1112 .collect()
1113 })
1114 }
1115
1116 fn create_step_partition_plan<'a>(
1117 &'a mut self,
1118 step_execution_id: StepExecutionId,
1119 entries: &'a [PartitionPlanEntry],
1120 ) -> BoxFuture<'a, Result<Vec<StepPartition>, RepositoryError>> {
1121 Box::pin(async move {
1122 let parent = self.staged.step_executions.get(&step_execution_id).ok_or(
1123 RepositoryError::StepExecutionNotFound {
1124 id: step_execution_id,
1125 },
1126 )?;
1127 if !matches!(
1128 parent.metadata().status(),
1129 BatchStatus::Starting | BatchStatus::Started
1130 ) {
1131 return Err(RepositoryError::PartitionParentNotActive {
1132 step_execution_id,
1133 status: parent.metadata().status(),
1134 });
1135 }
1136 if entries.is_empty() {
1137 return Err(RepositoryError::EmptyPartitionPlan);
1138 }
1139 if entries.len() > usize::from(MAX_PARTITIONS) {
1140 return Err(RepositoryError::PartitionPlanTooLarge {
1141 max: usize::from(MAX_PARTITIONS),
1142 });
1143 }
1144 if self
1145 .staged
1146 .step_partitions_by_step
1147 .contains_key(&step_execution_id)
1148 {
1149 return Err(RepositoryError::PartitionPlanExists { step_execution_id });
1150 }
1151 let mut keys = BTreeSet::new();
1152 for entry in entries {
1153 if !keys.insert(entry.key().clone()) {
1154 return Err(RepositoryError::DuplicatePartitionKey);
1155 }
1156 }
1157
1158 let mut partitions = Vec::with_capacity(entries.len());
1159 for (index, entry) in entries.iter().cloned().enumerate() {
1160 let id = self.next_step_partition_id()?;
1161 let ordinal = u32::try_from(index.saturating_add(1))
1162 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
1163 let partition = StepPartition::starting(id, step_execution_id, ordinal, entry);
1164 self.staged.step_partitions.insert(id, partition.clone());
1165 partitions.push(partition);
1166 }
1167 self.staged.step_partitions_by_step.insert(
1168 step_execution_id,
1169 partitions.iter().map(StepPartition::id).collect(),
1170 );
1171 self.created_partition_plans.insert(step_execution_id);
1172 Ok(partitions)
1173 })
1174 }
1175
1176 fn step_partition_plan(
1177 &mut self,
1178 step_execution_id: StepExecutionId,
1179 ) -> BoxFuture<'_, Result<Vec<StepPartition>, RepositoryError>> {
1180 Box::pin(async move {
1181 if !self.staged.step_executions.contains_key(&step_execution_id) {
1182 return Err(RepositoryError::StepExecutionNotFound {
1183 id: step_execution_id,
1184 });
1185 }
1186 let mut partitions = self
1187 .staged
1188 .step_partitions_by_step
1189 .get(&step_execution_id)
1190 .into_iter()
1191 .flatten()
1192 .map(|id| {
1193 self.staged
1194 .step_partitions
1195 .get(id)
1196 .cloned()
1197 .ok_or(RepositoryError::PartitionStateCorrupt)
1198 })
1199 .collect::<Result<Vec<_>, _>>()?;
1200 partitions.sort_by(|left, right| left.key().cmp(right.key()));
1201 Ok(partitions)
1202 })
1203 }
1204
1205 fn restart_step_partition_plan(
1206 &mut self,
1207 source_step_execution_id: StepExecutionId,
1208 target_step_execution_id: StepExecutionId,
1209 ) -> BoxFuture<'_, Result<Vec<StepPartition>, RepositoryError>> {
1210 Box::pin(async move {
1211 if self
1212 .staged
1213 .step_partitions_by_step
1214 .contains_key(&target_step_execution_id)
1215 {
1216 return Err(RepositoryError::PartitionPlanExists {
1217 step_execution_id: target_step_execution_id,
1218 });
1219 }
1220 let source_parent = self
1221 .staged
1222 .step_executions
1223 .get(&source_step_execution_id)
1224 .ok_or(RepositoryError::StepExecutionNotFound {
1225 id: source_step_execution_id,
1226 })?;
1227 let source_job = self
1228 .staged
1229 .job_executions
1230 .get(&source_parent.job_execution_id())
1231 .ok_or(RepositoryError::PartitionStateCorrupt)?;
1232 if !matches!(
1233 source_job.metadata().status(),
1234 BatchStatus::Failed | BatchStatus::Stopped
1235 ) {
1236 return Err(RepositoryError::PartitionStateCorrupt);
1237 }
1238 let target_parent = self
1239 .staged
1240 .step_executions
1241 .get(&target_step_execution_id)
1242 .ok_or(RepositoryError::StepExecutionNotFound {
1243 id: target_step_execution_id,
1244 })?;
1245 if !matches!(
1246 target_parent.metadata().status(),
1247 BatchStatus::Starting | BatchStatus::Started
1248 ) {
1249 return Err(RepositoryError::PartitionParentNotActive {
1250 step_execution_id: target_step_execution_id,
1251 status: target_parent.metadata().status(),
1252 });
1253 }
1254 let source_ids = self
1255 .staged
1256 .step_partitions_by_step
1257 .get(&source_step_execution_id)
1258 .cloned()
1259 .ok_or(RepositoryError::PartitionStateCorrupt)?;
1260 let mut copied = Vec::with_capacity(source_ids.len());
1261 for source_id in source_ids {
1262 let source = self
1263 .staged
1264 .step_partitions
1265 .get(&source_id)
1266 .cloned()
1267 .ok_or(RepositoryError::PartitionStateCorrupt)?;
1268 let id = self.next_step_partition_id()?;
1269 let partition = if source.status() == BatchStatus::Completed {
1270 if source.worker_step_execution_id().is_none() {
1271 return Err(RepositoryError::PartitionStateCorrupt);
1272 }
1273 StepPartition::from_snapshot(
1274 id,
1275 target_step_execution_id,
1276 source.worker_step_execution_id(),
1277 source.key().clone(),
1278 source.ordinal(),
1279 source.status(),
1280 source.exit_status().clone(),
1281 source.counts(),
1282 source.context().clone(),
1283 ExecutionVersion::INITIAL,
1284 )
1285 } else {
1286 StepPartition::starting(
1287 id,
1288 target_step_execution_id,
1289 source.ordinal(),
1290 PartitionPlanEntry::new(source.key().clone(), source.context().clone())
1291 .map_err(|_| RepositoryError::PartitionStateCorrupt)?,
1292 )
1293 };
1294 self.staged.step_partitions.insert(id, partition.clone());
1295 copied.push(partition);
1296 }
1297 self.staged.step_partitions_by_step.insert(
1298 target_step_execution_id,
1299 copied.iter().map(StepPartition::id).collect(),
1300 );
1301 self.created_partition_plans
1302 .insert(target_step_execution_id);
1303 Ok(copied)
1304 })
1305 }
1306
1307 fn assign_step_partition(
1308 &mut self,
1309 id: StepPartitionId,
1310 expected_version: ExecutionVersion,
1311 worker_step_execution_id: StepExecutionId,
1312 ) -> BoxFuture<'_, Result<StepPartition, RepositoryError>> {
1313 Box::pin(async move {
1314 let mut partition = self
1315 .staged
1316 .step_partitions
1317 .get(&id)
1318 .cloned()
1319 .ok_or(RepositoryError::StepPartitionNotFound { id })?;
1320 if self
1321 .created_partition_plans
1322 .contains(&partition.step_execution_id())
1323 {
1324 return Err(RepositoryError::PartitionPlanNotCommitted {
1325 step_execution_id: partition.step_execution_id(),
1326 });
1327 }
1328 let parent = self
1329 .staged
1330 .step_executions
1331 .get(&partition.step_execution_id())
1332 .ok_or(RepositoryError::PartitionStateCorrupt)?;
1333 if parent.metadata().status() != BatchStatus::Started {
1334 return Err(RepositoryError::PartitionParentNotActive {
1335 step_execution_id: parent.id(),
1336 status: parent.metadata().status(),
1337 });
1338 }
1339 partition
1340 .assign(expected_version, worker_step_execution_id)
1341 .map_err(|error| map_partition_mutation(id, error))?;
1342 let worker = self
1343 .staged
1344 .step_executions
1345 .get(&worker_step_execution_id)
1346 .ok_or(RepositoryError::StepExecutionNotFound {
1347 id: worker_step_execution_id,
1348 })?;
1349 if partition.step_execution_id() == worker_step_execution_id
1350 || parent.job_execution_id() != worker.job_execution_id()
1351 {
1352 return Err(RepositoryError::PartitionWorkerMismatch {
1353 partition_id: id,
1354 worker_step_execution_id,
1355 });
1356 }
1357 if self.staged.step_partitions.values().any(|candidate| {
1358 candidate.worker_step_execution_id() == Some(worker_step_execution_id)
1359 }) {
1360 return Err(RepositoryError::PartitionWorkerAlreadyAssigned {
1361 worker_step_execution_id,
1362 });
1363 }
1364 self.staged.step_partitions.insert(id, partition.clone());
1365 Ok(partition)
1366 })
1367 }
1368
1369 fn complete_step_partition(
1370 &mut self,
1371 id: StepPartitionId,
1372 expected_version: ExecutionVersion,
1373 worker_step_execution_id: StepExecutionId,
1374 ) -> BoxFuture<'_, Result<StepPartition, RepositoryError>> {
1375 Box::pin(async move {
1376 let mut partition = self
1377 .staged
1378 .step_partitions
1379 .get(&id)
1380 .cloned()
1381 .ok_or(RepositoryError::StepPartitionNotFound { id })?;
1382 let parent = self
1383 .staged
1384 .step_executions
1385 .get(&partition.step_execution_id())
1386 .ok_or(RepositoryError::PartitionStateCorrupt)?;
1387 if !matches!(
1388 parent.metadata().status(),
1389 BatchStatus::Started | BatchStatus::Stopping
1390 ) {
1391 return Err(RepositoryError::PartitionParentNotActive {
1392 step_execution_id: parent.id(),
1393 status: parent.metadata().status(),
1394 });
1395 }
1396 if partition.worker_step_execution_id() != Some(worker_step_execution_id) {
1397 return Err(RepositoryError::PartitionWorkerStale {
1398 partition_id: id,
1399 worker_step_execution_id,
1400 });
1401 }
1402 let worker = self
1403 .staged
1404 .step_executions
1405 .get(&worker_step_execution_id)
1406 .ok_or(RepositoryError::StepExecutionNotFound {
1407 id: worker_step_execution_id,
1408 })?;
1409 if worker.job_execution_id() != parent.job_execution_id() {
1410 return Err(RepositoryError::PartitionWorkerMismatch {
1411 partition_id: id,
1412 worker_step_execution_id,
1413 });
1414 }
1415 let result = PartitionResult::from_worker(worker).map_err(|_| {
1416 RepositoryError::PartitionAggregationIncomplete {
1417 step_execution_id: parent.id(),
1418 status: worker.metadata().status(),
1419 }
1420 })?;
1421 partition
1422 .complete(expected_version, &result)
1423 .map_err(|error| map_partition_mutation(id, error))?;
1424 self.staged.step_partitions.insert(id, partition.clone());
1425 Ok(partition.clone())
1426 })
1427 }
1428
1429 fn aggregate_step_partitions(
1430 &mut self,
1431 step_execution_id: StepExecutionId,
1432 expected_version: ExecutionVersion,
1433 transitioned_at: SystemTime,
1434 ) -> BoxFuture<'_, Result<StepExecution, RepositoryError>> {
1435 Box::pin(async move {
1436 let partitions = self.step_partition_plan(step_execution_id).await?;
1437 let aggregate = crate::aggregate_step_partitions(&partitions)
1438 .map_err(|error| map_partition_aggregation(step_execution_id, error))?;
1439 for partition in &partitions {
1440 let worker_id = partition.worker_step_execution_id().ok_or(
1441 RepositoryError::PartitionAggregationIncomplete {
1442 step_execution_id,
1443 status: partition.status(),
1444 },
1445 )?;
1446 let worker = self
1447 .staged
1448 .step_executions
1449 .get(&worker_id)
1450 .ok_or(RepositoryError::PartitionStateCorrupt)?;
1451 if worker.metadata().status() != partition.status()
1452 || worker.metadata().exit_status() != partition.exit_status()
1453 || worker.metadata().counts() != partition.counts()
1454 {
1455 return Err(RepositoryError::PartitionStateCorrupt);
1456 }
1457 }
1458 let parent = self
1459 .staged
1460 .step_executions
1461 .get(&step_execution_id)
1462 .cloned()
1463 .ok_or(RepositoryError::StepExecutionNotFound {
1464 id: step_execution_id,
1465 })?;
1466 let selected_worker = self
1467 .staged
1468 .step_executions
1469 .get(&aggregate.selected_worker_step_execution_id())
1470 .ok_or(RepositoryError::PartitionStateCorrupt)?;
1471 let failure = selected_worker.metadata().failure();
1472 if let Some(next) = expected_version.get().checked_add(1)
1473 && parent.version().get() == next
1474 && parent.metadata().status() == aggregate.status()
1475 && parent.metadata().exit_status() == aggregate.exit_status()
1476 && parent.metadata().counts() == aggregate.counts()
1477 && parent.metadata().failure() == failure
1478 {
1479 return Ok(parent);
1480 }
1481 if !matches!(
1482 parent.metadata().status(),
1483 BatchStatus::Started | BatchStatus::Stopping
1484 ) {
1485 return Err(RepositoryError::PartitionParentNotActive {
1486 step_execution_id,
1487 status: parent.metadata().status(),
1488 });
1489 }
1490 let aggregated = aggregate_partition_parent(
1491 &parent,
1492 expected_version,
1493 &aggregate,
1494 transitioned_at,
1495 failure,
1496 )?;
1497 self.staged
1498 .step_executions
1499 .insert(step_execution_id, aggregated.clone());
1500 self.aggregated_partition_parent = true;
1501 Ok(aggregated)
1502 })
1503 }
1504
1505 fn recover_job_execution<'a>(
1506 &'a mut self,
1507 id: JobExecutionId,
1508 request: &'a RecoveryRequest,
1509 ) -> BoxFuture<'a, Result<RecoveryResult, RepositoryError>> {
1510 Box::pin(async move {
1511 let prior = self
1512 .staged
1513 .job_executions
1514 .get(&id)
1515 .cloned()
1516 .ok_or(RepositoryError::JobExecutionNotFound { id })?;
1517 if self
1518 .staged
1519 .recovery_decisions
1520 .get(&id)
1521 .is_some_and(|decisions| {
1522 decisions
1523 .iter()
1524 .any(|decision| decision.execution_version() == request.expected_version())
1525 })
1526 {
1527 return Err(RepositoryError::ConcurrentModification);
1528 }
1529 let decided_at = self.repository.clock.now();
1530 let recovered = recovered_execution(&prior, request, decided_at)?;
1531 let decision_id = self.next_recovery_decision_id()?;
1532 let decision = RecoveryDecision::new(
1533 decision_id,
1534 id,
1535 request.expected_version(),
1536 prior.metadata().status(),
1537 recovered.metadata().status(),
1538 request.reason_code().to_owned(),
1539 request.operator_reference().to_owned(),
1540 *request.evidence_digest(),
1541 decided_at,
1542 );
1543 self.staged.job_executions.insert(id, recovered.clone());
1544 self.staged.execution_updated_at.insert(id, decided_at);
1545 self.staged
1546 .recovery_decisions
1547 .entry(id)
1548 .or_default()
1549 .push(decision.clone());
1550 Ok(RecoveryResult::new(recovered, decision))
1551 })
1552 }
1553
1554 fn recovery_decision(
1555 &mut self,
1556 id: JobExecutionId,
1557 ) -> BoxFuture<'_, Result<Option<RecoveryDecision>, RepositoryError>> {
1558 Box::pin(async move {
1559 if !self.staged.job_executions.contains_key(&id) {
1560 return Err(RepositoryError::JobExecutionNotFound { id });
1561 }
1562 Ok(self
1563 .staged
1564 .recovery_decisions
1565 .get(&id)
1566 .and_then(|decisions| decisions.first())
1567 .cloned())
1568 })
1569 }
1570
1571 fn find_operator_request<'a>(
1572 &'a mut self,
1573 action: OperatorAction,
1574 operation_id: &'a OperationId,
1575 ) -> BoxFuture<'a, Result<Option<OperatorRecord>, RepositoryError>> {
1576 Box::pin(async move {
1577 Ok(self
1578 .staged
1579 .operator_request_keys
1580 .get(&(action.as_str(), operation_id.as_str().to_owned()))
1581 .and_then(|id| self.staged.operator_requests.get(id))
1582 .cloned())
1583 })
1584 }
1585
1586 fn append_operator_request<'a>(
1587 &'a mut self,
1588 draft: &'a OperatorRecordDraft,
1589 ) -> BoxFuture<'a, Result<OperatorRecord, RepositoryError>> {
1590 Box::pin(async move {
1591 let key = (
1592 draft.action().as_str(),
1593 draft.operation_id().as_str().to_owned(),
1594 );
1595 if self.staged.operator_request_keys.contains_key(&key) {
1596 return Err(RepositoryError::ConcurrentModification);
1597 }
1598 let id = self.next_operator_request_id()?;
1599 let record = OperatorRecord::from_parts(id, draft.clone());
1600 self.staged.operator_requests.insert(id, record.clone());
1601 self.staged.operator_request_keys.insert(key, id);
1602 Ok(record)
1603 })
1604 }
1605
1606 fn request_execution_stop<'a>(
1607 &'a mut self,
1608 id: JobExecutionId,
1609 expected_version: ExecutionVersion,
1610 _actor: &'a ActorRef,
1611 requested_at: SystemTime,
1612 ) -> BoxFuture<'a, Result<JobExecution, RepositoryError>> {
1613 Box::pin(async move {
1614 let execution = self
1615 .staged
1616 .job_executions
1617 .get(&id)
1618 .cloned()
1619 .ok_or(RepositoryError::JobExecutionNotFound { id })?;
1620 if execution.version() != expected_version {
1621 return Err(RepositoryError::Lifecycle(LifecycleError::StaleVersion {
1622 expected: expected_version,
1623 actual: execution.version(),
1624 }));
1625 }
1626 self.staged.stop_requests.insert(id, requested_at);
1627 self.staged.execution_updated_at.insert(id, requested_at);
1628 Ok(execution)
1629 })
1630 }
1631
1632 fn claim_execution_owner<'a>(
1633 &'a mut self,
1634 id: JobExecutionId,
1635 expected_version: ExecutionVersion,
1636 owner: &'a OwnerToken,
1637 claimed_at: SystemTime,
1638 ) -> BoxFuture<'a, Result<JobExecution, RepositoryError>> {
1639 Box::pin(async move {
1640 let execution = self
1641 .staged
1642 .job_executions
1643 .get(&id)
1644 .cloned()
1645 .ok_or(RepositoryError::JobExecutionNotFound { id })?;
1646 if execution.version() != expected_version {
1647 return Err(RepositoryError::Lifecycle(LifecycleError::StaleVersion {
1648 expected: expected_version,
1649 actual: execution.version(),
1650 }));
1651 }
1652 if execution.metadata().status() != BatchStatus::Starting {
1653 return Err(RepositoryError::ExecutionOwnershipNotAllowed {
1654 id,
1655 status: execution.metadata().status(),
1656 });
1657 }
1658 if self
1659 .staged
1660 .owner_tokens
1661 .get(&id)
1662 .is_some_and(|recorded| recorded != owner)
1663 {
1664 return Err(RepositoryError::ExecutionOwned { id });
1665 }
1666 self.staged.owner_tokens.insert(id, *owner);
1667 self.staged.execution_updated_at.insert(id, claimed_at);
1668 Ok(execution)
1669 })
1670 }
1671
1672 fn observe_execution_control<'a>(
1673 &'a mut self,
1674 id: JobExecutionId,
1675 owner: &'a OwnerToken,
1676 observed_at: SystemTime,
1677 ) -> BoxFuture<'a, Result<crate::ExecutionControl, RepositoryError>> {
1678 Box::pin(async move {
1679 let owner_matches = self.staged.owner_tokens.get(&id) == Some(owner);
1680 let stop_requested = self.staged.stop_requests.contains_key(&id);
1681 let execution = self
1682 .staged
1683 .job_executions
1684 .get_mut(&id)
1685 .ok_or(RepositoryError::JobExecutionNotFound { id })?;
1686 if owner_matches
1687 && stop_requested
1688 && matches!(
1689 execution.metadata().status(),
1690 BatchStatus::Starting | BatchStatus::Started
1691 )
1692 {
1693 execution.transition(
1694 execution.version(),
1695 LifecycleTransition::new(BatchStatus::Stopping, observed_at),
1696 )?;
1697 self.staged.execution_updated_at.insert(id, observed_at);
1698 }
1699 Ok(crate::ExecutionControl::new(
1700 execution.clone(),
1701 owner_matches,
1702 stop_requested,
1703 ))
1704 })
1705 }
1706
1707 fn job_instance_hold(
1708 &mut self,
1709 id: JobInstanceId,
1710 ) -> BoxFuture<'_, Result<Option<RetentionHold>, RepositoryError>> {
1711 Box::pin(async move {
1712 if !self.staged.instances_by_id.contains_key(&id) {
1713 return Err(RepositoryError::JobInstanceNotFound { id });
1714 }
1715 Ok(self.staged.holds.get(&id).cloned())
1716 })
1717 }
1718
1719 fn place_instance_hold<'a>(
1720 &'a mut self,
1721 id: JobInstanceId,
1722 actor: &'a ActorRef,
1723 reason: &'a ReasonCode,
1724 placed_at: SystemTime,
1725 ) -> BoxFuture<'a, Result<RetentionHold, RepositoryError>> {
1726 Box::pin(async move {
1727 if !self.staged.instances_by_id.contains_key(&id) {
1728 return Err(RepositoryError::JobInstanceNotFound { id });
1729 }
1730 let hold = RetentionHold::new(id, actor.clone(), reason.clone(), placed_at);
1731 self.staged.holds.insert(id, hold.clone());
1732 Ok(hold)
1733 })
1734 }
1735
1736 fn release_instance_hold(
1737 &mut self,
1738 id: JobInstanceId,
1739 ) -> BoxFuture<'_, Result<Option<RetentionHold>, RepositoryError>> {
1740 Box::pin(async move {
1741 if !self.staged.instances_by_id.contains_key(&id) {
1742 return Err(RepositoryError::JobInstanceNotFound { id });
1743 }
1744 Ok(self.staged.holds.remove(&id))
1745 })
1746 }
1747
1748 fn find_retention_action<'a>(
1749 &'a mut self,
1750 action: RetentionAction,
1751 operation_id: &'a OperationId,
1752 ) -> BoxFuture<'a, Result<Option<RetentionRecord>, RepositoryError>> {
1753 Box::pin(async move {
1754 Ok(self
1755 .staged
1756 .retention_action_keys
1757 .get(&(action.as_str(), operation_id.as_str().to_owned()))
1758 .and_then(|id| self.staged.retention_actions.get(id))
1759 .cloned())
1760 })
1761 }
1762
1763 fn append_retention_action<'a>(
1764 &'a mut self,
1765 draft: &'a RetentionRecordDraft,
1766 ) -> BoxFuture<'a, Result<RetentionRecord, RepositoryError>> {
1767 Box::pin(async move {
1768 let key = (
1769 draft.action().as_str(),
1770 draft.operation_id().as_str().to_owned(),
1771 );
1772 if self.staged.retention_action_keys.contains_key(&key) {
1773 return Err(RepositoryError::ConcurrentModification);
1774 }
1775 let id = self.next_retention_action_id()?;
1776 let record = RetentionRecord::from_parts(id, draft.clone());
1777 self.staged.retention_actions.insert(id, record.clone());
1778 self.staged.retention_action_keys.insert(key, id);
1779 Ok(record)
1780 })
1781 }
1782
1783 fn purge_survey<'a>(
1784 &'a mut self,
1785 request: &'a PurgePlanRequest,
1786 ) -> BoxFuture<'a, Result<PurgeSurvey, RepositoryError>> {
1787 Box::pin(async move {
1788 let now = self.repository.clock.now();
1789 let candidates = self.purge_eligible(request, now);
1790 let counts = self.purge_counts(&candidates);
1791 Ok(PurgeSurvey::new(candidates, counts))
1792 })
1793 }
1794
1795 fn apply_purge<'a>(
1796 &'a mut self,
1797 plan: &'a PurgePlan,
1798 ) -> BoxFuture<'a, Result<PurgeCounts, RepositoryError>> {
1799 Box::pin(async move {
1800 for candidate in plan.candidates() {
1801 let execution = self
1802 .staged
1803 .job_executions
1804 .get(&candidate.job_execution_id())
1805 .ok_or(RepositoryError::RetentionPlanStale)?;
1806 let status = execution.metadata().status();
1807 if execution.version() != candidate.version()
1808 || !plan.request().statuses().contains(status)
1809 || self.staged.holds.contains_key(&candidate.job_instance_id())
1810 {
1811 return Err(RepositoryError::RetentionPlanStale);
1812 }
1813 let siblings_resolved = self
1814 .staged
1815 .job_executions_by_instance
1816 .get(&candidate.job_instance_id())
1817 .into_iter()
1818 .flatten()
1819 .filter_map(|id| self.staged.job_executions.get(id))
1820 .all(|sibling| sibling.metadata().status().is_finished());
1821 if !siblings_resolved {
1822 return Err(RepositoryError::RetentionPlanStale);
1823 }
1824 }
1825 let counts = self.purge_counts(plan.candidates());
1826 for candidate in plan.candidates() {
1827 let execution_id = candidate.job_execution_id();
1828 for decision_id in self
1829 .staged
1830 .flow_decisions_by_job
1831 .remove(&execution_id)
1832 .unwrap_or_default()
1833 {
1834 self.staged.flow_decisions.remove(&decision_id);
1835 }
1836 self.staged.recovery_decisions.remove(&execution_id);
1837 let request_ids = self
1838 .staged
1839 .operator_requests
1840 .iter()
1841 .filter(|(_, record)| record.job_execution_id() == Some(execution_id))
1842 .map(|(id, _)| *id)
1843 .collect::<Vec<_>>();
1844 for request_id in request_ids {
1845 if let Some(record) = self.staged.operator_requests.remove(&request_id) {
1846 self.staged.operator_request_keys.remove(&(
1847 record.action().as_str(),
1848 record.operation_id().as_str().to_owned(),
1849 ));
1850 }
1851 }
1852 for step_id in self
1853 .staged
1854 .step_executions_by_job
1855 .remove(&execution_id)
1856 .unwrap_or_default()
1857 {
1858 self.remove_step_execution(step_id);
1859 }
1860 self.staged.job_executions.remove(&execution_id);
1861 self.staged.execution_updated_at.remove(&execution_id);
1862 self.staged.owner_tokens.remove(&execution_id);
1863 self.staged.execution_definitions.remove(&execution_id);
1864 self.staged.stop_requests.remove(&execution_id);
1865 if let Some(executions) = self
1866 .staged
1867 .job_executions_by_instance
1868 .get_mut(&candidate.job_instance_id())
1869 {
1870 executions.retain(|id| *id != execution_id);
1871 }
1872 }
1873 let mut touched = plan
1874 .candidates()
1875 .iter()
1876 .map(PurgeCandidate::job_instance_id)
1877 .collect::<Vec<_>>();
1878 touched.dedup();
1879 for instance_id in touched {
1880 if self
1881 .staged
1882 .job_executions_by_instance
1883 .get(&instance_id)
1884 .is_none_or(|executions| !executions.is_empty())
1885 {
1886 continue;
1887 }
1888 let Some(instance) = self.staged.instances_by_id.remove(&instance_id) else {
1889 continue;
1890 };
1891 self.staged.instances_by_key.remove(instance.key());
1892 self.staged.job_executions_by_instance.remove(&instance_id);
1893 self.staged.instance_created_at.remove(&instance_id);
1894 self.staged.holds.remove(&instance_id);
1895 }
1896 Ok(counts)
1897 })
1898 }
1899
1900 fn commit<'a>(self: Box<Self>) -> BoxFuture<'a, Result<(), RepositoryError>>
1901 where
1902 Self: 'a,
1903 {
1904 Box::pin(async move {
1905 let mut current = self
1906 .repository
1907 .state
1908 .lock()
1909 .map_err(|_| RepositoryError::Unavailable)?;
1910 if current.revision != self.base_revision {
1911 return Err(RepositoryError::ConcurrentModification);
1912 }
1913 let mut staged = self.staged;
1914 staged.revision = current
1915 .revision
1916 .checked_add(1)
1917 .ok_or(RepositoryError::ConcurrentModification)?;
1918 *current = staged;
1919 if self.aggregated_partition_parent
1920 && self
1921 .repository
1922 .fail_next_partition_aggregate_commit
1923 .swap(false, Ordering::AcqRel)
1924 {
1925 return Err(RepositoryError::CommitOutcomeUnknown);
1926 }
1927 Ok(())
1928 })
1929 }
1930
1931 fn rollback<'a>(self: Box<Self>) -> BoxFuture<'a, Result<(), RepositoryError>>
1932 where
1933 Self: 'a,
1934 {
1935 Box::pin(async move { Ok(()) })
1936 }
1937}
1938
1939fn count_of<T>(values: Option<&Vec<T>>) -> u64 {
1940 values.map_or(0, |values| u64::try_from(values.len()).unwrap_or(u64::MAX))
1941}
1942
1943fn map_partition_mutation(id: StepPartitionId, error: PartitionMutationError) -> RepositoryError {
1944 match error {
1945 PartitionMutationError::StaleVersion { expected, actual } => {
1946 RepositoryError::Lifecycle(LifecycleError::StaleVersion { expected, actual })
1947 }
1948 PartitionMutationError::InvalidState { status } => {
1949 RepositoryError::PartitionUpdateNotAllowed { id, status }
1950 }
1951 PartitionMutationError::VersionExhausted => RepositoryError::PartitionStateCorrupt,
1952 }
1953}
1954
1955#[derive(Clone)]
1964pub struct InMemoryExplorer {
1965 state: Arc<Mutex<MemoryState>>,
1966 clock: Arc<dyn Clock>,
1967}
1968
1969impl InMemoryExplorer {
1970 #[must_use]
1972 pub fn new(repository: &InMemoryJobRepository) -> Self {
1973 Self {
1974 state: Arc::clone(&repository.state),
1975 clock: Arc::clone(&repository.clock),
1976 }
1977 }
1978
1979 fn snapshot(&self) -> Result<MemoryState, ExplorerError> {
1980 self.state
1981 .lock()
1982 .map(|state| state.clone())
1983 .map_err(|_| ExplorerError::Repository(RepositoryError::Unavailable))
1984 }
1985}
1986
1987impl fmt::Debug for InMemoryExplorer {
1988 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1989 formatter
1990 .debug_struct("InMemoryExplorer")
1991 .finish_non_exhaustive()
1992 }
1993}
1994
1995const fn identity_after(window: &QueryWindow) -> Option<u64> {
1996 match window.after() {
1997 Some(CursorKey::Identity(value)) => Some(*value),
1998 _ => None,
1999 }
2000}
2001
2002const fn ordered_after(window: &QueryWindow) -> Option<(u64, u64)> {
2003 match window.after() {
2004 Some(CursorKey::Ordered { primary, identity }) => Some((*primary, *identity)),
2005 _ => None,
2006 }
2007}
2008
2009fn name_after(window: &QueryWindow) -> Option<&str> {
2010 match window.after() {
2011 Some(CursorKey::Name(value)) => Some(value.as_str()),
2012 _ => None,
2013 }
2014}
2015
2016fn limit_of(window: &QueryWindow) -> usize {
2017 usize::from(window.limit())
2018}
2019
2020impl ExplorerRepository for InMemoryExplorer {
2021 fn identity_ceiling<'a>(
2022 &'a self,
2023 query: &'a ExplorerQuery,
2024 ) -> BoxFuture<'a, Result<u64, ExplorerError>> {
2025 Box::pin(async move {
2026 let state = self.snapshot()?;
2027 let ceiling = match query {
2028 ExplorerQuery::JobNames => state.job_name_order.values().copied().max(),
2029 ExplorerQuery::Instances { .. } => {
2030 state.instances_by_id.keys().next_back().map(|id| id.get())
2031 }
2032 ExplorerQuery::Executions { .. } | ExplorerQuery::UnresolvedExecutions { .. } => {
2033 state.job_executions.keys().next_back().map(|id| id.get())
2034 }
2035 ExplorerQuery::StepExecutions { .. } => {
2036 state.step_executions.keys().next_back().map(|id| id.get())
2037 }
2038 ExplorerQuery::RecoveryDecisions { .. } => state
2039 .recovery_decisions
2040 .values()
2041 .flatten()
2042 .map(|decision| decision.id().get())
2043 .max(),
2044 ExplorerQuery::FlowDecisions { .. } => {
2045 state.flow_decisions.keys().next_back().map(|id| id.get())
2046 }
2047 ExplorerQuery::StepPartitions { .. } => {
2048 state.step_partitions.keys().next_back().map(|id| id.get())
2049 }
2050 ExplorerQuery::OperatorRequests { .. } => state
2051 .operator_requests
2052 .keys()
2053 .next_back()
2054 .map(|id| id.get()),
2055 _ => return Err(ExplorerError::UnsupportedCapability),
2059 };
2060 Ok(ceiling.unwrap_or(0))
2061 })
2062 }
2063
2064 fn job_names<'a>(
2065 &'a self,
2066 window: &'a QueryWindow,
2067 ) -> BoxFuture<'a, Result<Vec<JobName>, ExplorerError>> {
2068 Box::pin(async move {
2069 let state = self.snapshot()?;
2070 let after = name_after(window);
2071 Ok(state
2072 .job_name_order
2073 .iter()
2074 .filter(|(_, order)| **order <= window.ceiling())
2075 .map(|(name, _)| name)
2076 .filter(|name| after.is_none_or(|after| name.as_str() > after))
2077 .take(limit_of(window))
2078 .cloned()
2079 .collect())
2080 })
2081 }
2082
2083 fn instances<'a>(
2084 &'a self,
2085 job_name: &'a JobName,
2086 window: &'a QueryWindow,
2087 ) -> BoxFuture<'a, Result<Vec<JobInstanceProjection>, ExplorerError>> {
2088 Box::pin(async move {
2089 let state = self.snapshot()?;
2090 let after = identity_after(window);
2091 let rows = state
2092 .instances_by_id
2093 .values()
2094 .rev()
2095 .filter(|instance| instance.key().job_name() == job_name)
2096 .filter(|instance| instance.id().get() <= window.ceiling())
2097 .filter(|instance| after.is_none_or(|after| instance.id().get() < after))
2098 .take(limit_of(window))
2099 .map(|instance| state.job_instance_projection(instance))
2100 .collect::<Vec<_>>();
2101 Ok(rows)
2102 })
2103 }
2104
2105 fn executions<'a>(
2106 &'a self,
2107 job_instance_id: JobInstanceId,
2108 window: &'a QueryWindow,
2109 ) -> BoxFuture<'a, Result<Vec<JobExecutionProjection>, ExplorerError>> {
2110 Box::pin(async move {
2111 let state = self.snapshot()?;
2112 if !state.instances_by_id.contains_key(&job_instance_id) {
2113 return Err(ExplorerError::Repository(
2114 RepositoryError::JobInstanceNotFound {
2115 id: job_instance_id,
2116 },
2117 ));
2118 }
2119 let after = ordered_after(window);
2120 state
2121 .job_executions_by_instance
2122 .get(&job_instance_id)
2123 .into_iter()
2124 .flatten()
2125 .rev()
2126 .filter_map(|id| state.job_executions.get(id))
2127 .filter(|execution| execution.id().get() <= window.ceiling())
2128 .filter(|execution| {
2129 after.is_none_or(|after| {
2130 (u64::from(state.attempt_of(execution)), execution.id().get()) < after
2131 })
2132 })
2133 .take(limit_of(window))
2134 .map(|execution| state.job_execution_projection(execution))
2135 .collect()
2136 })
2137 }
2138
2139 fn execution(
2140 &self,
2141 job_execution_id: JobExecutionId,
2142 ) -> BoxFuture<'_, Result<Option<JobExecutionProjection>, ExplorerError>> {
2143 Box::pin(async move {
2144 let state = self.snapshot()?;
2145 state
2146 .job_executions
2147 .get(&job_execution_id)
2148 .map(|execution| state.job_execution_projection(execution))
2149 .transpose()
2150 })
2151 }
2152
2153 fn step_executions<'a>(
2154 &'a self,
2155 job_execution_id: JobExecutionId,
2156 window: &'a QueryWindow,
2157 ) -> BoxFuture<'a, Result<Vec<StepExecutionProjection>, ExplorerError>> {
2158 Box::pin(async move {
2159 let state = self.snapshot()?;
2160 if !state.job_executions.contains_key(&job_execution_id) {
2161 return Err(ExplorerError::Repository(
2162 RepositoryError::JobExecutionNotFound {
2163 id: job_execution_id,
2164 },
2165 ));
2166 }
2167 let after = identity_after(window);
2168 Ok(state
2169 .step_executions_by_job
2170 .get(&job_execution_id)
2171 .into_iter()
2172 .flatten()
2173 .filter_map(|id| state.step_executions.get(id))
2174 .filter(|execution| execution.id().get() <= window.ceiling())
2175 .filter(|execution| after.is_none_or(|after| execution.id().get() > after))
2176 .take(limit_of(window))
2177 .map(|execution| state.step_execution_projection(execution))
2178 .collect())
2179 })
2180 }
2181
2182 fn unresolved_executions<'a>(
2183 &'a self,
2184 minimum_age: Duration,
2185 window: &'a QueryWindow,
2186 ) -> BoxFuture<'a, Result<Vec<JobExecutionProjection>, ExplorerError>> {
2187 Box::pin(async move {
2188 let state = self.snapshot()?;
2189 let now = self.clock.now();
2190 let after = identity_after(window);
2191 state
2192 .job_executions
2193 .values()
2194 .filter(|execution| !execution.metadata().status().is_finished())
2195 .filter(|execution| execution.id().get() <= window.ceiling())
2196 .filter(|execution| after.is_none_or(|after| execution.id().get() > after))
2197 .filter(|execution| {
2198 now.duration_since(updated_at(execution))
2199 .unwrap_or(Duration::ZERO)
2200 >= minimum_age
2201 })
2202 .take(limit_of(window))
2203 .map(|execution| state.job_execution_projection(execution))
2204 .collect()
2205 })
2206 }
2207
2208 fn recovery_decisions<'a>(
2209 &'a self,
2210 job_execution_id: JobExecutionId,
2211 window: &'a QueryWindow,
2212 ) -> BoxFuture<'a, Result<Vec<RecoveryDecision>, ExplorerError>> {
2213 Box::pin(async move {
2214 let state = self.snapshot()?;
2215 let after = identity_after(window);
2216 Ok(state
2217 .recovery_decisions
2218 .get(&job_execution_id)
2219 .into_iter()
2220 .flatten()
2221 .filter(|decision| decision.id().get() <= window.ceiling())
2222 .filter(|decision| after.is_none_or(|after| decision.id().get() > after))
2223 .take(limit_of(window))
2224 .cloned()
2225 .collect())
2226 })
2227 }
2228
2229 fn flow_decisions<'a>(
2230 &'a self,
2231 job_execution_id: JobExecutionId,
2232 window: &'a QueryWindow,
2233 ) -> BoxFuture<'a, Result<Vec<FlowDecision>, ExplorerError>> {
2234 Box::pin(async move {
2235 let state = self.snapshot()?;
2236 let after = ordered_after(window);
2237 Ok(state
2238 .flow_decisions_by_job
2239 .get(&job_execution_id)
2240 .into_iter()
2241 .flatten()
2242 .filter_map(|id| state.flow_decisions.get(id))
2243 .filter(|decision| decision.id().get() <= window.ceiling())
2244 .filter(|decision| {
2245 after.is_none_or(|after| {
2246 (decision.sequence().get(), decision.id().get()) > after
2247 })
2248 })
2249 .take(limit_of(window))
2250 .cloned()
2251 .collect())
2252 })
2253 }
2254
2255 fn step_partitions<'a>(
2256 &'a self,
2257 step_execution_id: StepExecutionId,
2258 window: &'a QueryWindow,
2259 ) -> BoxFuture<'a, Result<Vec<StepPartitionProjection>, ExplorerError>> {
2260 Box::pin(async move {
2261 let state = self.snapshot()?;
2262 if !state.step_executions.contains_key(&step_execution_id) {
2263 return Err(ExplorerError::Repository(
2264 RepositoryError::StepExecutionNotFound {
2265 id: step_execution_id,
2266 },
2267 ));
2268 }
2269 let after = identity_after(window);
2270 state
2271 .step_partitions_by_step
2272 .get(&step_execution_id)
2273 .into_iter()
2274 .flatten()
2275 .filter_map(|id| state.step_partitions.get(id))
2276 .filter(|partition| partition.id().get() <= window.ceiling())
2277 .filter(|partition| after.is_none_or(|after| partition.id().get() > after))
2278 .take(limit_of(window))
2279 .map(|partition| {
2280 Ok(StepPartitionProjection::new(
2281 partition.id(),
2282 partition.step_execution_id(),
2283 partition.key().as_str().to_owned(),
2284 partition.ordinal(),
2285 partition.status(),
2286 partition.exit_status().clone(),
2287 partition.counts(),
2288 partition.version(),
2289 partition.worker_step_execution_id(),
2290 Some(StateEnvelopeDescriptor::new(
2291 DurableStateKind::ExecutionContext,
2292 partition.context().format_version(),
2293 partition.context().schema_id().clone(),
2294 partition.context().schema_version(),
2295 partition.context().encoded_len(),
2296 )),
2297 ))
2298 })
2299 .collect()
2300 })
2301 }
2302
2303 fn operator_requests<'a>(
2304 &'a self,
2305 job_execution_id: JobExecutionId,
2306 window: &'a QueryWindow,
2307 ) -> BoxFuture<'a, Result<Vec<OperatorRecord>, ExplorerError>> {
2308 Box::pin(async move {
2309 let state = self.snapshot()?;
2310 let after = identity_after(window);
2311 Ok(state
2312 .operator_requests
2313 .values()
2314 .filter(|record| record.job_execution_id() == Some(job_execution_id))
2315 .filter(|record| record.id().get() <= window.ceiling())
2316 .filter(|record| after.is_none_or(|after| record.id().get() > after))
2317 .take(limit_of(window))
2318 .cloned()
2319 .collect())
2320 })
2321 }
2322}
2323
2324impl RecoveryRepository for InMemoryExplorer {
2325 fn recovery_snapshot<'a>(
2326 &'a self,
2327 execution_id: JobExecutionId,
2328 current_owner: &'a OwnerToken,
2329 ) -> BoxFuture<'a, Result<RecoverySnapshot, RepositoryError>> {
2330 Box::pin(async move {
2331 let state = self
2332 .state
2333 .lock()
2334 .map_err(|_| RepositoryError::Unavailable)?
2335 .clone();
2336 let execution = state
2337 .job_executions
2338 .get(&execution_id)
2339 .ok_or(RepositoryError::JobExecutionNotFound { id: execution_id })?;
2340 let owner = match state.owner_tokens.get(&execution_id) {
2341 None => OwnerObservation::Absent,
2342 Some(recorded) if recorded == current_owner => OwnerObservation::CurrentProcess,
2343 Some(_) => OwnerObservation::OtherProcess,
2344 };
2345 let latest_step = state
2346 .step_executions_by_job
2347 .get(&execution_id)
2348 .and_then(|ids| ids.last())
2349 .and_then(|id| state.step_executions.get(id))
2350 .map(|step| RecoveryStepEvidence::new(step.id(), step.metadata().status(), None));
2351 let unknown_commit = execution.metadata().status() == BatchStatus::Unknown
2352 || execution.metadata().failure().is_some_and(|failure| {
2353 failure.category() == crate::FailureCategory::UnknownCommit
2354 })
2355 || latest_step
2356 .as_ref()
2357 .is_some_and(|step| step.status() == BatchStatus::Unknown);
2358 let committed_flow_decision = state
2359 .flow_decisions_by_job
2360 .get(&execution_id)
2361 .is_some_and(|decisions| !decisions.is_empty());
2362 let ambiguous_external_effect = state
2363 .execution_definitions
2364 .get(&execution_id)
2365 .is_none_or(definition_has_ambiguous_effect);
2366 Ok(RecoverySnapshot::new(
2367 execution_id,
2368 execution.metadata().status(),
2369 state.attempt_of(execution),
2370 execution.version(),
2371 owner,
2372 state
2373 .execution_updated_at
2374 .get(&execution_id)
2375 .copied()
2376 .unwrap_or_else(|| updated_at(execution)),
2377 self.clock.now(),
2378 latest_step,
2379 crate::RecoveryMarkers::new()
2380 .with_unknown_commit(unknown_commit)
2381 .with_committed_flow_decision(committed_flow_decision)
2382 .with_ambiguous_external_effect(ambiguous_external_effect),
2383 ))
2384 })
2385 }
2386}
2387
2388fn definition_has_ambiguous_effect(definition: &DefinitionIdentity) -> bool {
2389 let Ok(document) = serde_json::from_slice::<serde_json::Value>(definition.canonical_manifest())
2390 else {
2391 return true;
2392 };
2393 document
2394 .get("delivery_mode")
2395 .and_then(serde_json::Value::as_str)
2396 != Some("atomic_same_resource")
2397}