feat: make dedicated workers for flows able to share runtime for the same scripts

This commit is contained in:
Ruben Fiszel
2024-01-12 00:21:52 +01:00
parent ad6a17e43d
commit 08a3f79642
+41 -20
View File
@@ -1183,8 +1183,8 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
if let Ok(flow) = value {
let workers = spawn_dedicated_workers_for_flow(
&flow.modules,
&_wp.path,
&_wp.workspace_id,
&_wp.path,
killpill_tx.clone(),
&killpill_rx,
db,
@@ -1195,7 +1195,13 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
)
.await;
workers.into_iter().for_each(|(path, sender, handle)| {
dedicated_handles.push(handle);
tracing::info!(
"spawned dedicated worker for flow: {}",
path.as_str()
);
if let Some(h) = handle {
dedicated_handles.push(h);
}
hm.insert(path, sender);
});
}
@@ -1224,7 +1230,9 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
)
.await
{
dedicated_handles.push(handle);
if let Some(h) = handle {
dedicated_handles.push(h);
}
hm.insert(path, sender);
} else {
tracing::error!(
@@ -1754,7 +1762,7 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
println!("worker {} exited", i_worker);
}
type DedicatedWorker = (String, Sender<Arc<QueuedJob>>, JoinHandle<()>);
type DedicatedWorker = (String, Sender<Arc<QueuedJob>>, Option<JoinHandle<()>>);
// spawn one dedicated worker per compatible steps of the flow, associating the node id to the dedicated worker channel send
#[async_recursion]
@@ -1771,24 +1779,37 @@ async fn spawn_dedicated_workers_for_flow(
job_completed_tx: &JobCompletedSender,
) -> Vec<DedicatedWorker> {
let mut workers = vec![];
let mut script_path_to_worker: HashMap<String, Sender<Arc<QueuedJob>>> = HashMap::new();
for module in modules.iter() {
match &module.value {
FlowModuleValue::Script { path, hash, .. } => {
if let Some(dedi_w) = spawn_dedicated_worker(
SpawnWorker::Script { path: path.to_string(), hash: hash.clone() },
w_id,
killpill_tx.clone(),
killpill_rx,
db,
worker_dir,
base_internal_url,
worker_name,
job_completed_tx,
Some(module.id.clone()),
)
.await
{
workers.push(dedi_w);
let key = format!(
"{}:{}",
path,
hash.clone()
.map(|x| x.to_string())
.unwrap_or_else(|| "".to_string())
);
if let Some(sender) = script_path_to_worker.get(&key) {
workers.push((module.id.clone(), sender.clone(), None));
} else {
if let Some(dedi_w) = spawn_dedicated_worker(
SpawnWorker::Script { path: path.to_string(), hash: hash.clone() },
w_id,
killpill_tx.clone(),
killpill_rx,
db,
worker_dir,
base_internal_url,
worker_name,
job_completed_tx,
Some(module.id.clone()),
)
.await
{
script_path_to_worker.insert(key, dedi_w.1.clone());
workers.push(dedi_w);
}
}
}
FlowModuleValue::ForloopFlow { modules, .. } => {
@@ -2051,7 +2072,7 @@ async fn spawn_dedicated_worker(
tracing::error!("error in dedicated worker: {:?}", e);
};
});
return Some((node_id.unwrap_or(path2), dedicated_worker_tx, handle));
return Some((node_id.unwrap_or(path2), dedicated_worker_tx, Some(handle)));
// (Some(dedi_path), Some(dedicated_worker_tx), Some(handle))
}
}