add barrier when num workers > 1

This commit is contained in:
Ruben Fiszel
2023-04-19 15:16:44 +02:00
parent 17646130bc
commit 62d196ecec
6 changed files with 88 additions and 29 deletions
+13 -4
View File
@@ -14,9 +14,9 @@ use std::{
use git_version::git_version;
use monitor::handle_zombie_jobs_periodically;
use sqlx::{Pool, Postgres};
use tokio::sync::RwLock;
use tokio::{fs::DirBuilder, sync::RwLock};
use windmill_common::{utils::rd_string, IS_READY, METRICS_ADDR};
use windmill_worker::S3_CACHE_BUCKET;
use windmill_worker::{DENO_CACHE_DIR, GO_CACHE_DIR, PIP_CACHE_DIR, S3_CACHE_BUCKET};
const GIT_VERSION: &str = git_version!(args = ["--tag", "--always"], fallback = "unknown-version");
const DEFAULT_NUM_WORKERS: usize = 3;
@@ -271,6 +271,14 @@ pub async fn run_workers<R: rsmq_async::RsmqConnection + Send + Sync + Clone + '
let mut handles = Vec::with_capacity(num_workers as usize);
for x in [PIP_CACHE_DIR, DENO_CACHE_DIR, GO_CACHE_DIR] {
DirBuilder::new()
.recursive(true)
.create(x)
.await
.expect("could not create initial worker dir");
}
#[cfg(feature = "enterprise")]
if let Some(ref s) = S3_CACHE_BUCKET.clone() {
// We donwload the entire cache as tar
@@ -278,7 +286,7 @@ pub async fn run_workers<R: rsmq_async::RsmqConnection + Send + Sync + Clone + '
}
IS_READY.store(true, Ordering::Relaxed);
let sync_barrier = Arc::new(RwLock::new(None));
for i in 1..(num_workers + 1) {
let db1 = db.clone();
let instance_name = instance_name.clone();
@@ -287,6 +295,7 @@ pub async fn run_workers<R: rsmq_async::RsmqConnection + Send + Sync + Clone + '
let rx = rx.resubscribe();
let base_internal_url = base_internal_url.clone();
let rsmq2 = rsmq.clone();
let sync_barrier = sync_barrier.clone();
handles.push(tokio::spawn(monitor.instrument(async move {
tracing::info!(worker = %worker_name, "starting worker");
windmill_worker::run_worker(
@@ -299,7 +308,7 @@ pub async fn run_workers<R: rsmq_async::RsmqConnection + Send + Sync + Clone + '
rx,
&base_internal_url,
rsmq2,
RwLock::new(Arc::new(None)),
sync_barrier,
)
.await
})));
+1 -1
View File
@@ -927,7 +927,7 @@ fn spawn_test_worker(
rx,
&base_internal_url,
None,
RwLock::new(Arc::new(None)),
Arc::new(RwLock::new(None)),
)
.await
};
+2 -2
View File
@@ -6,7 +6,7 @@
* LICENSE-AGPL for a copy of the license.
*/
use std::{net::SocketAddr, sync::Arc};
use std::net::SocketAddr;
use error::Error;
@@ -43,7 +43,7 @@ lazy_static::lazy_static! {
.flatten();
pub static ref METRICS_ENABLED: bool = METRICS_ADDR.is_some();
pub static ref BASE_URL: String = std::env::var("BASE_URL").unwrap_or_else(|_| "http://localhost".to_string());
pub static ref IS_READY: Arc<std::sync::atomic::AtomicBool> = Arc::new(std::sync::atomic::AtomicBool::new(false));
pub static ref IS_READY: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
}
#[cfg(feature = "tokio")]
+20 -4
View File
@@ -45,7 +45,9 @@ pub async fn copy_cache_from_bucket(bucket: &str, tx: Sender<()>) -> error::Resu
"--size-only",
"--fast-list",
"--exclude",
&format!("\"/{TAR_CACHE_FILENAME},/deno/gen/file/**\""),
&format!("deno/gen/file/tmp/windmill/**"),
"--exclude",
&format!("{TAR_CACHE_FILENAME}"),
],
)
.await
@@ -79,7 +81,9 @@ pub async fn copy_cache_to_bucket(bucket: &str) -> error::Result<()> {
"--size-only",
"--fast-list",
"--exclude",
&format!("\"/{TAR_CACHE_FILENAME},/deno/gen/file/**\""),
&format!("deno/gen/file/tmp/windmill/**"),
"--exclude",
&format!("{TAR_CACHE_FILENAME}"),
],
)
.await
@@ -251,7 +255,13 @@ pub async fn copy_tmp_cache_to_cache() -> error::Result<()> {
execute_command(
TMP_DIR,
"rclone",
vec!["sync", ROOT_TMP_CACHE_DIR, ROOT_CACHE_DIR],
vec![
"sync",
ROOT_TMP_CACHE_DIR,
ROOT_CACHE_DIR,
"--exclude",
TAR_CACHE_FILENAME,
],
)
.await?;
tracing::info!(
@@ -267,7 +277,13 @@ pub async fn copy_cache_to_tmp_cache() -> error::Result<()> {
execute_command(
TMP_DIR,
"rclone",
vec!["sync", ROOT_CACHE_DIR, ROOT_TMP_CACHE_DIR],
vec![
"sync",
ROOT_CACHE_DIR,
ROOT_TMP_CACHE_DIR,
"--exclude",
TAR_CACHE_FILENAME,
],
)
.await?;
tracing::info!(
+1
View File
@@ -7,5 +7,6 @@ mod python_executor;
mod worker;
mod worker_flow;
#[cfg(feature = "enterprise")]
pub use global_cache::copy_cache_from_bucket_as_tar;
pub use worker::*;
+51 -18
View File
@@ -232,6 +232,8 @@ lazy_static::lazy_static! {
.map(|e| Some(e))
.unwrap_or(None);
pub static ref CAN_PULL: Arc<RwLock<()>> = Arc::new(RwLock::new(()));
}
//only matter if CLOUD_HOSTED
@@ -284,7 +286,7 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
mut rx: tokio::sync::broadcast::Receiver<()>,
base_internal_url: &str,
rsmq: Option<R>,
sync_barrier: RwLock<Arc<Option<Barrier>>>,
sync_barrier: Arc<RwLock<Option<Barrier>>>,
) {
#[cfg(not(feature = "enterprise"))]
if !*DISABLE_NSJAIL {
@@ -298,13 +300,11 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
let worker_dir = format!("{TMP_DIR}/{worker_name}");
tracing::debug!(worker_dir = %worker_dir, worker_name = %worker_name, "Creating worker dir");
for x in [&worker_dir, PIP_CACHE_DIR, DENO_CACHE_DIR, GO_CACHE_DIR] {
DirBuilder::new()
.recursive(true)
.create(x)
.await
.expect("could not create initial worker dir");
}
DirBuilder::new()
.recursive(true)
.create(&worker_dir)
.await
.expect("could not create initial worker dir");
let _ = write_file(
&worker_dir,
@@ -468,8 +468,15 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
if i_worker == 1 && S3_CACHE_BUCKET.is_some() {
if last_sync.elapsed().as_secs() > *GLOBAL_CACHE_INTERVAL &&
(copy_cache_from_bucket_handle.is_none() || copy_cache_from_bucket_handle.as_ref().unwrap().is_finished()) {
tracing::debug!("CAN PULL LOCK START");
let _lock = CAN_PULL.write().await;
tracing::info!("Started syncing cache");
last_sync = Instant::now();
if num_workers > 1 {
create_barrier_for_all_workers(num_workers, sync_barrier.clone()).await;
}
if let Err(e) = copy_cache_to_tmp_cache().await {
tracing::error!("failed to copy cache to tmp cache: {}", e);
} else {
@@ -484,12 +491,20 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
}
}
// The barrier is to avoid the sync to bucket syncing partial folders
#[cfg(feature = "enterprise")]
if num_workers > 1 {
if num_workers > 1 && S3_CACHE_BUCKET.is_some() {
let read_barrier = sync_barrier.read().await;
let barrier = read_barrier.clone();
if let Some(b) = barrier.as_ref() {
if let Some(b) = read_barrier.as_ref() {
tracing::debug!("worker #{i_worker} waiting for barrier");
b.wait().await;
tracing::debug!("worker #{i_worker} done waiting for barrier");
drop(read_barrier);
// wait for barrier to be reset
let _ = CAN_PULL.read().await;
tracing::debug!("worker #{i_worker} done waiting for lock");
} else {
tracing::debug!("worker #{i_worker} no barrier");
};
}
@@ -509,19 +524,17 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
(true, Ok(None))
},
_ = copy_to_bucket_rx.recv() => {
tracing::debug!("CAN PULL LOCK START");
let _lock = CAN_PULL.write().await;
if num_workers > 1 {
let mut barrier = sync_barrier.write().await;
let arc_barrier = Arc::new(Some(tokio::sync::Barrier::new(num_workers as usize)));
*barrier = arc_barrier.clone();
if let Some(b) = arc_barrier.as_ref() {
b.wait().await;
};
create_barrier_for_all_workers(num_workers, sync_barrier.clone()).await;
}
//Arc::new(tokio::sync::Barrier::new(num_workers as usize + 1));
#[cfg(feature = "enterprise")]
if let Err(e) = copy_tmp_cache_to_cache().await {
tracing::error!(worker = %worker_name, "failed to sync tmp cache to cache: {}", e);
}
tracing::debug!("CAN PULL LOCK END");
(false, Ok(None))
},
Some(job_id) = same_worker_rx.recv() => {
@@ -580,7 +593,12 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
} else { None };
let job_root = job.root_job.map(|x| x.to_string()).unwrap_or_else(|| "none".to_string());
tracing::info!(worker = %worker_name, id = %job.id, root_id = %job_root, "fetched job {}, root job: {}", job.id, job_root);
if job.id == Uuid::nil() {
tracing::info!(worker = %worker_name, "running warmup job");
} else {
tracing::info!(worker = %worker_name, id = %job.id, root_id = %job_root, "fetched job {}, root job: {}", job.id, job_root);
}
let job_dir = format!("{worker_dir}/{}", job.id);
@@ -684,6 +702,21 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
}
pub async fn create_barrier_for_all_workers(num_workers: u32, sync_barrier: Arc<RwLock<Option<tokio::sync::Barrier>>>) {
tracing::debug!("acquiring write lock");
let mut barrier = sync_barrier.write().await;
*barrier = Some(tokio::sync::Barrier::new(num_workers as usize));
drop(barrier);
tracing::debug!("dropped write lock");
if let Some(b) = sync_barrier.read().await.as_ref() {
tracing::debug!("leader worker waiting for barrier");
b.wait().await;
tracing::debug!("leader worker done waiting for barrier");
};
let mut barrier = sync_barrier.write().await;
*barrier = None;
tracing::debug!("leader worker done waiting for");
}
pub async fn handle_job_error<R: rsmq_async::RsmqConnection + Send + Sync + Clone>(
db: &Pool<Postgres>,
client: &AuthedClient,