use tokio::time::Instant; use windmill_common::error; #[cfg(all(feature = "enterprise", feature = "parquet"))] use windmill_object_store::object_store_reexports::ObjectStore; #[cfg(all(feature = "enterprise", feature = "parquet"))] use std::sync::Arc; pub const TARGET: &str = const_format::concatcp!(std::env::consts::OS, "_", std::env::consts::ARCH); #[cfg(all(feature = "enterprise", feature = "parquet"))] pub async fn build_tar_and_push( s3_client: Arc, folder: String, lang: String, custom_folder_name: Option, platform_agnostic: bool, ) -> error::Result<()> { use tokio::fs::create_dir_all; use windmill_object_store::object_store_reexports::Path; use crate::TAR_PYBASE_CACHE_DIR; tracing::info!("Started building and pushing piptar {folder}"); let start = Instant::now(); // e.g. tiny==1.0.0 let folder_name = if let Some(name) = custom_folder_name { name } else { folder.split("/").last().unwrap().to_owned() }; let prefix = &format!("{}/{}", *TAR_PYBASE_CACHE_DIR, lang); let tar_path = format!("{prefix}/{folder_name}_tar.tar"); create_dir_all(prefix).await?; let tar_file = std::fs::File::create(&tar_path)?; let mut tar = tar::Builder::new(tar_file); tar.append_dir_all(".", &folder)?; // Write the trailing zero blocks and close the inner file BEFORE std::fs::read // below. Without this, the bytes we upload to S3 are an unfinalized archive. drop(tar.into_inner()?); let tar_metadata = tokio::fs::metadata(&tar_path).await; if tar_metadata.is_err() || tar_metadata.as_ref().unwrap().len() == 0 { tracing::info!("Failed to tar cache: {folder}"); return Err(error::Error::ExecutionErr(format!( "Failed to tar cache: {folder}" ))); } // let s3_settings = S3_CACHE_SETTINGS.read().await; // let s3_client = s3_settings.as_ref().ok_or_else(|| { // error::Error::ExecutionErr("Failed to read s3 cache settings".to_string()) // })?; if let Err(e) = s3_client .put( &Path::from(format!( "/tar/{}/{lang}/{folder_name}.tar", if platform_agnostic { "" } else { TARGET } )), std::fs::read(&tar_path)?.into(), ) .await { tracing::info!("Failed to put tar to s3: {tar_path}. Error: {:?}", e); return Err(error::Error::ExecutionErr(format!( "Failed to put tar to s3: {tar_path}" ))); } tokio::fs::remove_file(&tar_path).await.map_err(|e| { tracing::error!("Failed to remove piptar {folder_name}. Error: {:?}", e); e })?; tracing::info!( "Finished copying piptar {folder} to bucket as tar, took: {:?}s. Size of tar: {}", start.elapsed().as_secs(), tar_metadata.unwrap().len(), ); Ok(()) } #[cfg(all(feature = "enterprise", feature = "parquet"))] pub async fn pull_from_tar( client: Arc, folder: String, lang: String, custom_folder_name: Option, platform_agnostic: bool, ) -> error::Result<()> { use windmill_object_store::attempt_fetch_bytes; let folder_name = if let Some(name) = custom_folder_name { name } else { folder.split("/").last().unwrap().to_owned() }; tracing::info!("Attempting to pull tar {folder_name} from bucket"); let start = Instant::now(); let tar_path = format!( "tar/{}/{lang}/{folder_name}.tar", if platform_agnostic { "" } else { TARGET } ); let bytes = attempt_fetch_bytes(client, &tar_path).await?; extract_tar(bytes, &folder).map_err(|e| { tracing::error!("Failed to extract piptar {folder_name}. Error: {:?}", e); e })?; tracing::info!( "Finished pulling and extracting {folder_name}. Took {:?}ms", start.elapsed().as_millis() ); Ok(()) } pub fn extract_tar(tar: bytes::Bytes, folder: &str) -> error::Result<()> { use bytes::Buf; let start: Instant = Instant::now(); std::fs::create_dir_all(&folder)?; let mut ar = tar::Archive::new(tar.reader()); if let Err(e) = ar.unpack(folder) { tracing::info!("Failed to untar to {folder}. Error: {:?}", e); std::fs::remove_dir_all(&folder)?; return Err(error::Error::ExecutionErr(format!( "Failed to untar tar {folder}. Error: {:?}", e ))); } tracing::info!( "Finished extracting tar to {folder}. Took {}ms", start.elapsed().as_millis(), ); Ok(()) } /// Two-tier cache load: check local disk first, then fall back to the shared object store. /// /// "Shared" is the worker group's own store when its config overrides one, the instance store /// otherwise — see [`windmill_object_store::get_cache_object_store`]. /// Returns `(hit, log_message)`. pub async fn load_cache(bin_path: &str, _remote_path: &str, is_dir: bool) -> (bool, String) { if tokio::fs::metadata(&bin_path).await.is_ok() { (true, format!("loaded from local cache: {}\n", bin_path)) } else { #[cfg(all(feature = "enterprise", feature = "parquet"))] if let Some(os) = windmill_object_store::get_cache_object_store().await { let started = std::time::Instant::now(); if let Ok(mut x) = windmill_object_store::attempt_fetch_bytes(os, _remote_path).await { if is_dir { // Extract into a sibling temp dir then atomically publish it, // so a concurrent cold-load gating on metadata(bin_path) never // observes a half-extracted cache directory. let tmp_dir = format!("{}.tmp.{}", bin_path, uuid::Uuid::new_v4()); let res = match windmill_common::worker::extract_tar(x, &tmp_dir).await { Ok(()) => windmill_common::worker::atomic_publish_dir(&tmp_dir, bin_path), Err(e) => Err(e), }; if let Err(e) = res { let _ = tokio::fs::remove_dir_all(&tmp_dir).await; tracing::error!("could not write tar archive locally: {e:?}"); return ( false, "error writing tar archive from object store".to_string(), ); } } else { if let Err(e) = windmill_common::worker::write_binary_file(bin_path, &mut x) { tracing::error!("could not write bundle/bin file locally: {e:?}"); return ( false, "error writing bundle/bin file from object store".to_string(), ); } } tracing::info!("loaded from object store {}", bin_path); return ( true, format!( "loaded bin/bundle from object store {} in {}ms", bin_path, started.elapsed().as_millis() ), ); } } let _ = is_dir; (false, "".to_string()) } } /// Whether this worker can push to the shared object store at all — the features are /// compiled in and a store is loaded. False on builds without them, where `save_cache` /// only ever writes to the worker's own disk. pub async fn object_store_available() -> bool { #[cfg(all(feature = "enterprise", feature = "parquet"))] { windmill_object_store::get_cache_object_store() .await .is_some() } #[cfg(not(all(feature = "enterprise", feature = "parquet")))] { false } } /// Whether a binary/bundle is in the shared object store, ignoring the local cache. /// /// The deploy-time prebuild asks this rather than [`exists_in_cache`]: a copy on the /// building worker's own disk is exactly the state the prebuild exists to fix, so /// answering from it would latch a failed upload into a permanent skip. pub async fn exists_in_object_store(_remote_path: &str) -> bool { #[cfg(all(feature = "enterprise", feature = "parquet"))] if let Some(os) = windmill_object_store::get_cache_object_store().await { return os .head(&windmill_object_store::object_store_reexports::Path::from( _remote_path, )) .await .is_ok(); } false } /// Fail a deploy-time prebuild whose artifact never reached the object store. `save_cache` /// logs and swallows a failed upload, which is right for a run that has the binary locally /// anyway — but for a prebuild the upload *is* the result, and a silent miss would be /// latched by the next build's existence check. pub async fn ensure_pushed_to_object_store(remote_path: &str) -> error::Result<()> { if exists_in_object_store(remote_path).await { return Ok(()); } Err(error::Error::ExecutionErr(format!( "the binary was built but did not reach the object store at {remote_path}, \ so no other worker can load it" ))) } /// Check whether a binary/bundle exists in local cache or the shared object store. pub async fn exists_in_cache(bin_path: &str, _remote_path: &str) -> bool { if tokio::fs::metadata(&bin_path).await.is_ok() { return true; } else { #[cfg(all(feature = "enterprise", feature = "parquet"))] if let Some(os) = windmill_object_store::get_cache_object_store().await { return os .get(&windmill_object_store::object_store_reexports::Path::from( _remote_path, )) .await .is_ok(); } return false; } } /// Two-tier cache write: upload to the shared object store, then copy to local disk. pub async fn save_cache( local_cache_path: &str, _remote_cache_path: &str, origin: &str, is_dir: bool, ) -> windmill_common::error::Result { use std::path::PathBuf; let mut _cached_to_s3 = false; #[cfg(all(feature = "enterprise", feature = "parquet"))] if let Some(os) = windmill_object_store::get_cache_object_store().await { use windmill_object_store::object_store_reexports::Path; let file_to_cache = if is_dir { let tar_path = format!( "{}/tar/{}_tar.tar", *windmill_common::worker::ROOT_CACHE_DIR, local_cache_path .split("/") .last() .unwrap_or(&uuid::Uuid::new_v4().to_string()) ); let tar_file = std::fs::File::create(&tar_path)?; let mut tar = tar::Builder::new(tar_file); tar.append_dir_all(".", &origin)?; drop(tar.into_inner()?); let tar_metadata = tokio::fs::metadata(&tar_path).await; if tar_metadata.is_err() || tar_metadata.as_ref().unwrap().len() == 0 { tracing::info!("Failed to tar cache: {origin}"); return Err(error::Error::ExecutionErr(format!( "Failed to tar cache: {origin}" ))); } tar_path } else { origin.to_owned() }; if let Err(e) = os .put( &Path::from(_remote_cache_path), std::fs::read(&file_to_cache)?.into(), ) .await { tracing::error!( "Failed to put bin to object store: {_remote_cache_path}. Error: {:?}", e ); } else { _cached_to_s3 = true; if is_dir { tokio::fs::remove_dir_all(&file_to_cache).await?; } } } if true { if is_dir { // Populate a sibling temp dir then atomically publish it, so a // concurrent `load_cache`/`exists_in_cache` metadata() check never // observes a half-copied cache directory. let tmp_dir = format!("{}.tmp.{}", local_cache_path, uuid::Uuid::new_v4()); if let Err(e) = windmill_common::worker::copy_dir_recursively( &PathBuf::from(origin), &PathBuf::from(&tmp_dir), ) .and_then(|_| windmill_common::worker::atomic_publish_dir(&tmp_dir, local_cache_path)) { let _ = std::fs::remove_dir_all(&tmp_dir); return Err(e); } } else { windmill_common::worker::atomic_copy_file(origin, local_cache_path)?; } Ok(format!( "\nwrote cached binary: {} (backed by EE distributed object store: {_cached_to_s3})\n", local_cache_path )) } else if _cached_to_s3 { Ok(format!( "wrote cached binary to object store {}\n", local_cache_path )) } else { Ok("".to_string()) } }