runtime: canonicalize Worker aggregates
This commit is contained in:
@@ -532,12 +532,16 @@ impl Runtime {
|
||||
status: WorkerStatus::Stopped,
|
||||
workspace_id: scope.map(|scope| scope.workspace_id.clone()),
|
||||
request: request.clone(),
|
||||
run_generation: 1,
|
||||
working_directory: None,
|
||||
execution_handle: None,
|
||||
};
|
||||
state.workers.insert(worker_id, record);
|
||||
state.persist_runtime_snapshot()?;
|
||||
state.persist_worker(&worker_ref.worker_id)?;
|
||||
let spawn_request = WorkerExecutionSpawnRequest {
|
||||
worker_ref: worker_ref.clone(),
|
||||
run_generation: 1,
|
||||
request,
|
||||
context: self.execution_context(worker_ref.clone()),
|
||||
working_directory: None,
|
||||
@@ -859,35 +863,46 @@ impl Runtime {
|
||||
/// present this is idempotent; otherwise the configured backend is tried.
|
||||
pub fn restore_worker(&self, worker_ref: &WorkerRef) -> Result<WorkerDetail, RuntimeError> {
|
||||
let (backend, request) = {
|
||||
let state = self.lock()?;
|
||||
let mut state = self.lock()?;
|
||||
state.ensure_running()?;
|
||||
let worker = state.worker(worker_ref)?;
|
||||
if worker.execution_handle.is_some() {
|
||||
return Ok(worker.detail());
|
||||
}
|
||||
if worker.status == WorkerStatus::Cancelled {
|
||||
return Err(RuntimeError::InvalidRequest(format!(
|
||||
"worker {} is cancelled",
|
||||
worker_ref.worker_id
|
||||
)));
|
||||
}
|
||||
let (worker_request, previous_working_directory, config_bundle, run_generation) = {
|
||||
let worker = state.worker(worker_ref)?;
|
||||
if worker.execution_handle.is_some() {
|
||||
return Ok(worker.detail());
|
||||
}
|
||||
if worker.status == WorkerStatus::Cancelled {
|
||||
return Err(RuntimeError::InvalidRequest(format!(
|
||||
"worker {} is cancelled",
|
||||
worker_ref.worker_id
|
||||
)));
|
||||
}
|
||||
let config_bundle = worker
|
||||
.request
|
||||
.config_bundle
|
||||
.as_ref()
|
||||
.and_then(|bundle_ref| state.config_bundles.get(&bundle_ref.id))
|
||||
.cloned();
|
||||
(
|
||||
worker.request.clone(),
|
||||
worker.working_directory.clone(),
|
||||
config_bundle,
|
||||
worker.run_generation.saturating_add(1).max(1),
|
||||
)
|
||||
};
|
||||
let backend = state.execution_backend.clone().ok_or_else(|| {
|
||||
RuntimeError::WorkerExecutionUnavailable {
|
||||
worker_id: worker_ref.worker_id.clone(),
|
||||
message: "runtime has no execution backend".to_string(),
|
||||
}
|
||||
})?;
|
||||
let config_bundle = worker
|
||||
.request
|
||||
.config_bundle
|
||||
.as_ref()
|
||||
.and_then(|bundle_ref| state.config_bundles.get(&bundle_ref.id))
|
||||
.cloned();
|
||||
state.worker_mut(worker_ref)?.run_generation = run_generation;
|
||||
state.persist_worker(&worker_ref.worker_id)?;
|
||||
let request = WorkerExecutionRestoreRequest {
|
||||
worker_ref: worker_ref.clone(),
|
||||
request: worker.request.clone(),
|
||||
run_generation,
|
||||
request: worker_request,
|
||||
context: self.execution_context(worker_ref.clone()),
|
||||
previous_working_directory: worker.working_directory.clone(),
|
||||
previous_working_directory,
|
||||
working_directory: None,
|
||||
config_bundle,
|
||||
};
|
||||
@@ -1514,34 +1529,64 @@ impl Runtime {
|
||||
struct RestoreCandidate {
|
||||
worker_ref: WorkerRef,
|
||||
request: CreateWorkerRequest,
|
||||
run_generation: u64,
|
||||
previous_working_directory: Option<CatalogWorkingDirectoryStatus>,
|
||||
config_bundle: Option<ConfigBundle>,
|
||||
}
|
||||
|
||||
let candidates = {
|
||||
let state = self.lock()?;
|
||||
let mut state = self.lock()?;
|
||||
if state.execution_backend.is_none() {
|
||||
return Ok(());
|
||||
}
|
||||
state
|
||||
let worker_ids = state
|
||||
.workers
|
||||
.values()
|
||||
.filter(|worker| worker.execution_handle.is_none())
|
||||
.map(|worker| {
|
||||
.map(|worker| worker.worker_id)
|
||||
.collect::<Vec<_>>();
|
||||
let mut candidates = Vec::with_capacity(worker_ids.len());
|
||||
for worker_id in worker_ids {
|
||||
let (
|
||||
worker_ref,
|
||||
request,
|
||||
previous_working_directory,
|
||||
config_bundle,
|
||||
run_generation,
|
||||
) = {
|
||||
let worker = state
|
||||
.workers
|
||||
.get(&worker_id)
|
||||
.expect("collected Worker exists");
|
||||
let config_bundle = worker
|
||||
.request
|
||||
.config_bundle
|
||||
.as_ref()
|
||||
.and_then(|bundle_ref| state.config_bundles.get(&bundle_ref.id))
|
||||
.cloned();
|
||||
RestoreCandidate {
|
||||
worker_ref: worker.worker_ref.clone(),
|
||||
request: worker.request.clone(),
|
||||
previous_working_directory: worker.working_directory.clone(),
|
||||
(
|
||||
worker.worker_ref.clone(),
|
||||
worker.request.clone(),
|
||||
worker.working_directory.clone(),
|
||||
config_bundle,
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
worker.run_generation.saturating_add(1).max(1),
|
||||
)
|
||||
};
|
||||
state
|
||||
.workers
|
||||
.get_mut(&worker_id)
|
||||
.expect("collected Worker exists")
|
||||
.run_generation = run_generation;
|
||||
state.persist_worker(&worker_id)?;
|
||||
candidates.push(RestoreCandidate {
|
||||
worker_ref,
|
||||
request,
|
||||
run_generation,
|
||||
previous_working_directory,
|
||||
config_bundle,
|
||||
});
|
||||
}
|
||||
candidates
|
||||
};
|
||||
|
||||
for candidate in candidates {
|
||||
@@ -1554,6 +1599,7 @@ impl Runtime {
|
||||
};
|
||||
let request = WorkerExecutionRestoreRequest {
|
||||
worker_ref: candidate.worker_ref.clone(),
|
||||
run_generation: candidate.run_generation,
|
||||
request: candidate.request,
|
||||
context: self.execution_context(candidate.worker_ref.clone()),
|
||||
previous_working_directory: candidate.previous_working_directory,
|
||||
@@ -1723,6 +1769,7 @@ impl RuntimeState {
|
||||
status: WorkerStatus::Stopped,
|
||||
workspace_id: worker.workspace_id,
|
||||
request: worker.request,
|
||||
run_generation: worker.run_generation,
|
||||
working_directory: worker.working_directory,
|
||||
execution_handle: None,
|
||||
},
|
||||
@@ -2281,6 +2328,7 @@ struct WorkerRecord {
|
||||
status: WorkerStatus,
|
||||
workspace_id: Option<String>,
|
||||
request: CreateWorkerRequest,
|
||||
run_generation: u64,
|
||||
working_directory: Option<CatalogWorkingDirectoryStatus>,
|
||||
execution_handle: Option<WorkerExecutionHandle>,
|
||||
}
|
||||
@@ -2324,6 +2372,7 @@ impl WorkerRecord {
|
||||
worker_ref: self.worker_ref.clone(),
|
||||
worker_id: self.worker_id.clone(),
|
||||
request: self.request.clone(),
|
||||
run_generation: self.run_generation,
|
||||
workspace_id: self.workspace_id.clone(),
|
||||
working_directory: self.working_directory.clone(),
|
||||
}
|
||||
@@ -2627,6 +2676,7 @@ mod tests {
|
||||
dispatch_result: Mutex<Option<WorkerExecutionResult>>,
|
||||
restore_result: Mutex<Option<WorkerExecutionSpawnResult>>,
|
||||
restore_count: Mutex<u64>,
|
||||
run_generations: Mutex<Vec<u64>>,
|
||||
contexts: Mutex<BTreeMap<WorkerId, WorkerExecutionContext>>,
|
||||
dispatched_inputs: Mutex<Vec<WorkerInput>>,
|
||||
preserve_commit_ack_submission_id: AtomicBool,
|
||||
@@ -2670,6 +2720,10 @@ mod tests {
|
||||
}
|
||||
|
||||
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult {
|
||||
self.run_generations
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(request.run_generation);
|
||||
self.contexts
|
||||
.lock()
|
||||
.unwrap()
|
||||
@@ -2689,6 +2743,10 @@ mod tests {
|
||||
request: WorkerExecutionRestoreRequest,
|
||||
) -> WorkerExecutionSpawnResult {
|
||||
*self.restore_count.lock().unwrap() += 1;
|
||||
self.run_generations
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(request.run_generation);
|
||||
if let Some(result) = self.restore_result.lock().unwrap().clone() {
|
||||
return result;
|
||||
}
|
||||
@@ -3551,6 +3609,7 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(*backend.restore_count.lock().unwrap(), 1);
|
||||
assert_eq!(*backend.run_generations.lock().unwrap(), vec![1, 2]);
|
||||
assert_eq!(
|
||||
runtime.worker_detail(&detail.worker_ref).unwrap().status,
|
||||
WorkerStatus::Idle
|
||||
|
||||
Reference in New Issue
Block a user