diff --git a/Cargo.lock b/Cargo.lock index 97515ca24d..23741e2a3d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -892,6 +892,8 @@ version = "0.1.0" dependencies = [ "anyhow", "chrono", + "regex", + "remote_storage", "serde", "serde_json", "serde_with", diff --git a/compute_tools/src/bin/compute_ctl.rs b/compute_tools/src/bin/compute_ctl.rs index ad747470c2..343bb41d3b 100644 --- a/compute_tools/src/bin/compute_ctl.rs +++ b/compute_tools/src/bin/compute_ctl.rs @@ -38,7 +38,7 @@ use std::fs::File; use std::panic; use std::path::Path; use std::process::exit; -use std::sync::{mpsc, Arc, Condvar, Mutex, OnceLock, RwLock}; +use std::sync::{mpsc, Arc, Condvar, Mutex, RwLock}; use std::{thread, time::Duration}; use anyhow::{Context, Result}; @@ -147,6 +147,7 @@ fn main() -> Result<()> { match spec_json { // First, try to get cluster spec from the cli argument Some(json) => { + info!("got spec from cli argument {}", json); spec = Some(serde_json::from_str(json)?); } None => { @@ -182,6 +183,7 @@ fn main() -> Result<()> { if let Some(spec) = spec { let pspec = ParsedSpec::try_from(spec).map_err(|msg| anyhow::anyhow!(msg))?; + info!("new pspec.spec: {:?}", pspec.spec); new_state.pspec = Some(pspec); spec_set = true; } else { @@ -196,9 +198,7 @@ fn main() -> Result<()> { state: Mutex::new(new_state), state_changed: Condvar::new(), ext_remote_storage, - ext_remote_paths: OnceLock::new(), ext_download_progress: RwLock::new(HashMap::new()), - library_index: OnceLock::new(), build_tag, }; let compute = Arc::new(compute_node); diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 03ffedc0e3..d9c71621c2 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -5,7 +5,7 @@ use std::os::unix::fs::PermissionsExt; use std::path::Path; use std::process::{Command, Stdio}; use std::str::FromStr; -use std::sync::{Condvar, Mutex, OnceLock, RwLock}; +use std::sync::{Condvar, Mutex, RwLock}; use std::time::Instant; use anyhow::{Context, Result}; @@ -14,7 +14,6 @@ use futures::future::join_all; use futures::stream::FuturesUnordered; use futures::StreamExt; use postgres::{Client, NoTls}; -use regex::Regex; use tokio; use tokio_postgres; use tracing::{error, info, instrument, warn}; @@ -60,10 +59,6 @@ pub struct ComputeNode { pub state_changed: Condvar, /// the S3 bucket that we search for extensions in pub ext_remote_storage: Option, - // (key: extension name, value: path to extension archive in remote storage) - pub ext_remote_paths: OnceLock>, - // (key: library name, value: name of extension containing this library) - pub library_index: OnceLock>, // key: ext_archive_name, value: started download time, download_completed? pub ext_download_progress: RwLock, bool)>>, pub build_tag: String, @@ -75,7 +70,6 @@ pub struct RemoteExtensionMetrics { num_ext_downloaded: u64, largest_ext_size: u64, total_ext_download_size: u64, - prep_extensions_ms: u64, } #[derive(Clone, Debug)] @@ -745,11 +739,19 @@ impl ComputeNode { pspec.timeline_id, ); + info!( + "start_compute spec.remote_extensions {:?}", + pspec.spec.remote_extensions + ); + // This part is sync, because we need to download // remote shared_preload_libraries before postgres start (if any) - { + if let Some(remote_extensions) = &pspec.spec.remote_extensions { + // First, create control files for all availale extensions + extension_server::create_control_files(remote_extensions, &self.pgbin); + let library_load_start_time = Utc::now(); - let remote_ext_metrics = self.prepare_preload_libraries(&compute_state)?; + let remote_ext_metrics = self.prepare_preload_libraries(&pspec.spec)?; let library_load_time = Utc::now() .signed_duration_since(library_load_start_time) @@ -761,7 +763,6 @@ impl ComputeNode { state.metrics.num_ext_downloaded = remote_ext_metrics.num_ext_downloaded; state.metrics.largest_ext_size = remote_ext_metrics.largest_ext_size; state.metrics.total_ext_download_size = remote_ext_metrics.total_ext_download_size; - state.metrics.prep_extensions_ms = remote_ext_metrics.prep_extensions_ms; info!( "Loading shared_preload_libraries took {:?}ms", library_load_time @@ -918,38 +919,11 @@ LIMIT 100", } } - // If remote extension storage is configured, - // download extension control files - pub async fn prepare_external_extensions(&self, compute_state: &ComputeState) -> Result<()> { - if let Some(ref ext_remote_storage) = self.ext_remote_storage { - let pspec = compute_state.pspec.as_ref().expect("spec must be set"); - let spec = &pspec.spec; - let custom_ext = spec.custom_extensions.clone().unwrap_or(Vec::new()); - info!("custom extensions: {:?}", &custom_ext); - - let (ext_remote_paths, library_index) = extension_server::get_available_extensions( - ext_remote_storage, - &self.pgbin, - &self.pgversion, - &custom_ext, - &self.build_tag, - ) - .await?; - self.ext_remote_paths - .set(ext_remote_paths) - .expect("this is the only time we set ext_remote_paths"); - self.library_index - .set(library_index) - .expect("this is the only time we set library_index"); - } - Ok(()) - } - // download an archive, unzip and place files in correct locations pub async fn download_extension( &self, - ext_name: &str, - is_library: bool, + real_ext_name: String, + ext_path: RemotePath, ) -> Result { let remote_storage = self .ext_remote_storage @@ -958,35 +932,6 @@ LIMIT 100", "Remote extensions storage is not configured", )))?; - let mut real_ext_name = ext_name; - if is_library { - // sometimes library names might have a suffix like - // library.so or library.so.3. We strip this off - // because library_index is based on the name without the file extension - let strip_lib_suffix = Regex::new(r"\.so.*").unwrap(); - let lib_raw_name = strip_lib_suffix.replace(real_ext_name, "").to_string(); - - real_ext_name = self - .library_index - .get() - .expect("must have already downloaded the library_index") - .get(&lib_raw_name) - .ok_or(DownloadError::BadInput(anyhow::anyhow!( - "library {} is not found", - lib_raw_name - )))?; - } - - let ext_path = &self - .ext_remote_paths - .get() - .expect("error accessing ext_remote_paths") - .get(real_ext_name) - .ok_or(DownloadError::BadInput(anyhow::anyhow!( - "real_ext_name {} is not found", - real_ext_name - )))?; - let ext_archive_name = ext_path.object_name().expect("bad path"); let mut first_try = false; @@ -1039,8 +984,8 @@ LIMIT 100", info!("downloading new extension {ext_archive_name}"); let download_size = extension_server::download_extension( - real_ext_name, - ext_path, + &real_ext_name, + &ext_path, remote_storage, &self.pgbin, ) @@ -1058,18 +1003,19 @@ LIMIT 100", #[tokio::main] pub async fn prepare_preload_libraries( &self, - compute_state: &ComputeState, + spec: &ComputeSpec, ) -> Result { if self.ext_remote_storage.is_none() { return Ok(RemoteExtensionMetrics { num_ext_downloaded: 0, largest_ext_size: 0, total_ext_download_size: 0, - prep_extensions_ms: 0, }); } - let pspec = compute_state.pspec.as_ref().expect("spec must be set"); - let spec = &pspec.spec; + let remote_extensions = spec + .remote_extensions + .as_ref() + .ok_or(anyhow::anyhow!("Remote extensions are not configured",))?; info!("parse shared_preload_libraries from spec.cluster.settings"); let mut libs_vec = Vec::new(); @@ -1081,6 +1027,7 @@ LIMIT 100", .collect(); } info!("parse shared_preload_libraries from provided postgresql.conf"); + // that is used in neon_local and python tests if let Some(conf) = &spec.cluster.postgresql_conf { let conf_lines = conf.split('\n').collect::>(); @@ -1101,30 +1048,16 @@ LIMIT 100", libs_vec.extend(preload_libs_vec); } - info!("Download ext_index.json, find the extension paths"); - let prep_ext_start_time = Utc::now(); - self.prepare_external_extensions(compute_state).await?; - let prep_ext_time_delta = Utc::now() - .signed_duration_since(prep_ext_start_time) - .to_std() - .unwrap() - .as_millis() as u64; - info!("Prepare extensions took {prep_ext_time_delta}ms"); - // Don't try to download libraries that are not in the index. // Assume that they are already present locally. - libs_vec.retain(|lib| { - self.library_index - .get() - .expect("error accessing ext_remote_paths") - .contains_key(lib) - }); + libs_vec.retain(|lib| remote_extensions.library_index.contains_key(lib)); info!("Downloading to shared preload libraries: {:?}", &libs_vec); let mut download_tasks = Vec::new(); for library in &libs_vec { - download_tasks.push(self.download_extension(library, true)); + let (ext_name, ext_path) = remote_extensions.get_ext(library, true)?; + download_tasks.push(self.download_extension(ext_name, ext_path)); } let results = join_all(download_tasks).await; @@ -1132,7 +1065,6 @@ LIMIT 100", num_ext_downloaded: 0, largest_ext_size: 0, total_ext_download_size: 0, - prep_extensions_ms: prep_ext_time_delta, }; for result in results { let download_size = match result { diff --git a/compute_tools/src/extension_server.rs b/compute_tools/src/extension_server.rs index 1b5e4cf974..cb54a603e0 100644 --- a/compute_tools/src/extension_server.rs +++ b/compute_tools/src/extension_server.rs @@ -73,10 +73,9 @@ More specifically, here is an example ext_index.json */ use anyhow::Context; use anyhow::{self, Result}; -use futures::future::join_all; +use compute_api::spec::RemoteExtSpec; use remote_storage::*; use serde_json; -use std::collections::HashMap; use std::io::Read; use std::num::{NonZeroU32, NonZeroUsize}; use std::path::Path; @@ -117,81 +116,6 @@ pub fn get_pg_version(pgbin: &str) -> String { panic!("Unsuported postgres version {human_version}"); } -// download control files for enabled_extensions -// return Hashmaps converting library names to extension names (library_index) -// and specifying the remote path to the archive for each extension name -pub async fn get_available_extensions( - remote_storage: &GenericRemoteStorage, - pgbin: &str, - pg_version: &str, - custom_extensions: &[String], - build_tag: &str, -) -> Result<(HashMap, HashMap)> { - let local_sharedir = Path::new(&get_pg_config("--sharedir", pgbin)).join("extension"); - let index_path = format!("{build_tag}/{pg_version}/ext_index.json"); - let index_path = RemotePath::new(Path::new(&index_path)).context("error forming path")?; - info!("download ext_index.json from: {:?}", &index_path); - - let mut download = remote_storage.download(&index_path).await?; - let mut ext_idx_buffer = Vec::new(); - download - .download_stream - .read_to_end(&mut ext_idx_buffer) - .await?; - info!("ext_index downloaded"); - - #[derive(Debug, serde::Deserialize)] - struct Index { - public_extensions: Vec, - library_index: HashMap, - extension_data: HashMap, - } - - #[derive(Debug, serde::Deserialize)] - struct ExtensionData { - control_data: HashMap, - archive_path: String, - } - - let ext_index_full = serde_json::from_slice::(&ext_idx_buffer)?; - let mut enabled_extensions = ext_index_full.public_extensions; - enabled_extensions.extend_from_slice(custom_extensions); - let mut library_index = ext_index_full.library_index; - let all_extension_data = ext_index_full.extension_data; - info!("library_index: {:?}", library_index); - - info!("enabled_extensions: {:?}", enabled_extensions); - let mut ext_remote_paths = HashMap::new(); - let mut file_create_tasks = Vec::new(); - for extension in enabled_extensions { - let ext_data = &all_extension_data[&extension]; - for (control_file, control_contents) in &ext_data.control_data { - let extension_name = control_file - .strip_suffix(".control") - .expect("control files must end in .control"); - let control_path = local_sharedir.join(control_file); - if !control_path.exists() { - ext_remote_paths.insert( - extension_name.to_string(), - RemotePath::from_string(&ext_data.archive_path)?, - ); - info!("writing file {:?}{:?}", control_path, control_contents); - file_create_tasks.push(tokio::fs::write(control_path, control_contents)); - } else { - warn!("control file {:?} exists both locally and remotely. ignoring the remote version.", control_file); - // also delete this from library index - library_index.retain(|_, value| value != extension_name); - } - } - } - let results = join_all(file_create_tasks).await; - for result in results { - result?; - } - info!("ext_remote_paths {:?}", ext_remote_paths); - Ok((ext_remote_paths, library_index)) -} - // download the archive for a given extension, // unzip it, and place files in the appropriate locations (share/lib) pub async fn download_extension( @@ -253,6 +177,22 @@ pub async fn download_extension( Ok(download_size) } +// Create extension control files from spec +pub fn create_control_files(remote_extensions: &RemoteExtSpec, pgbin: &str) { + let local_sharedir = Path::new(&get_pg_config("--sharedir", pgbin)).join("extension"); + for ext_data in remote_extensions.extension_data.values() { + for (control_name, control_content) in &ext_data.control_data { + let control_path = local_sharedir.join(control_name); + if !control_path.exists() { + info!("writing file {:?}{:?}", control_path, control_content); + std::fs::write(control_path, control_content).unwrap(); + } else { + warn!("control file {:?} exists both locally and remotely. ignoring the remote version.", control_path); + } + } + } +} + // This function initializes the necessary structs to use remote storage pub fn init_remote_storage(remote_ext_config: &str) -> anyhow::Result { #[derive(Debug, serde::Deserialize)] diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs index 7713d2bb51..841e533a3a 100644 --- a/compute_tools/src/http/api.rs +++ b/compute_tools/src/http/api.rs @@ -13,7 +13,7 @@ use hyper::{Body, Method, Request, Response, Server, StatusCode}; use num_cpus; use serde_json; use tokio::task; -use tracing::{error, info}; +use tracing::{error, info, warn}; use tracing_utils::http::OtelName; fn status_response_from_state(state: &ComputeState) -> ComputeStatusResponse { @@ -126,6 +126,15 @@ async fn routes(req: Request, compute: &Arc) -> Response, compute: &Arc) -> Response Response::new(Body::from("OK")), + // debug only + info!("spec: {:?}", spec); + + let remote_extensions = match spec.remote_extensions.as_ref() { + Some(r) => r, + None => { + info!("no remote extensions spec was provided"); + let mut resp = Response::new(Body::from("no remote storage configured")); + *resp.status_mut() = StatusCode::INTERNAL_SERVER_ERROR; + return resp; + } + }; + + remote_extensions.get_ext(&filename, is_library) + }; + + match ext { + Ok((ext_name, ext_path)) => { + match compute.download_extension(ext_name, ext_path).await { + Ok(_) => Response::new(Body::from("OK")), + Err(e) => { + error!("extension download failed: {}", e); + let mut resp = Response::new(Body::from(e.to_string())); + *resp.status_mut() = StatusCode::INTERNAL_SERVER_ERROR; + resp + } + } + } Err(e) => { - error!("extension download failed: {}", e); - let mut resp = Response::new(Body::from(e.to_string())); + warn!("extension download failed to find extension: {}", e); + let mut resp = Response::new(Body::from("failed to find file")); *resp.status_mut() = StatusCode::INTERNAL_SERVER_ERROR; resp } diff --git a/control_plane/src/endpoint.rs b/control_plane/src/endpoint.rs index 6060799458..7ebcf98ab0 100644 --- a/control_plane/src/endpoint.rs +++ b/control_plane/src/endpoint.rs @@ -493,7 +493,7 @@ impl Endpoint { pageserver_connstring: Some(pageserver_connstring), safekeeper_connstrings, storage_auth_token: auth_token.clone(), - custom_extensions: Some(vec![]), + remote_extensions: None, }; let spec_path = self.endpoint_path().join("spec.json"); std::fs::write(spec_path, serde_json::to_string_pretty(&spec)?)?; diff --git a/libs/compute_api/Cargo.toml b/libs/compute_api/Cargo.toml index 428d031a93..b377bd2cce 100644 --- a/libs/compute_api/Cargo.toml +++ b/libs/compute_api/Cargo.toml @@ -10,6 +10,9 @@ chrono.workspace = true serde.workspace = true serde_with.workspace = true serde_json.workspace = true +regex.workspace = true utils = { path = "../utils" } +remote_storage = { version = "0.1", path = "../remote_storage/" } + workspace_hack.workspace = true diff --git a/libs/compute_api/src/responses.rs b/libs/compute_api/src/responses.rs index b9aff7282b..92bbf79cd4 100644 --- a/libs/compute_api/src/responses.rs +++ b/libs/compute_api/src/responses.rs @@ -107,7 +107,6 @@ pub struct ComputeMetrics { pub num_ext_downloaded: u64, pub largest_ext_size: u64, // these are measured in bytes pub total_ext_download_size: u64, - pub prep_extensions_ms: u64, } /// Response of the `/computes/{compute_id}/spec` control-plane API. diff --git a/libs/compute_api/src/spec.rs b/libs/compute_api/src/spec.rs index 293f6dc294..3da922e517 100644 --- a/libs/compute_api/src/spec.rs +++ b/libs/compute_api/src/spec.rs @@ -3,11 +3,16 @@ //! The spec.json file is used to pass information to 'compute_ctl'. It contains //! all the information needed to start up the right version of PostgreSQL, //! and connect it to the storage nodes. +use std::collections::HashMap; + use serde::{Deserialize, Serialize}; use serde_with::{serde_as, DisplayFromStr}; use utils::id::{TenantId, TimelineId}; use utils::lsn::Lsn; +use regex::Regex; +use remote_storage::RemotePath; + /// String type alias representing Postgres identifier and /// intended to be used for DB / role names. pub type PgIdent = String; @@ -61,8 +66,55 @@ pub struct ComputeSpec { /// the pageserver and safekeepers. pub storage_auth_token: Option, - // list of prefixes to search for custom extensions in remote extension storage + // information about available remote extensions + pub remote_extensions: Option, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct RemoteExtSpec { + pub public_extensions: Option>, pub custom_extensions: Option>, + pub library_index: HashMap, + pub extension_data: HashMap, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ExtensionData { + pub control_data: HashMap, + pub archive_path: String, +} + +impl RemoteExtSpec { + pub fn get_ext( + &self, + ext_name: &str, + is_library: bool, + ) -> anyhow::Result<(String, RemotePath)> { + let mut real_ext_name = ext_name; + if is_library { + // sometimes library names might have a suffix like + // library.so or library.so.3. We strip this off + // because library_index is based on the name without the file extension + let strip_lib_suffix = Regex::new(r"\.so.*").unwrap(); + let lib_raw_name = strip_lib_suffix.replace(real_ext_name, "").to_string(); + + real_ext_name = self + .library_index + .get(&lib_raw_name) + .ok_or(anyhow::anyhow!("library {} is not found", lib_raw_name))?; + } + + match self.extension_data.get(real_ext_name) { + Some(ext_data) => Ok(( + real_ext_name.to_string(), + RemotePath::from_string(&ext_data.archive_path)?, + )), + None => Err(anyhow::anyhow!( + "real_ext_name {} is not found", + real_ext_name + )), + } + } } #[serde_as] diff --git a/libs/compute_api/tests/cluster_spec.json b/libs/compute_api/tests/cluster_spec.json index 8f81e7b3bd..96db13a5da 100644 --- a/libs/compute_api/tests/cluster_spec.json +++ b/libs/compute_api/tests/cluster_spec.json @@ -205,5 +205,43 @@ "name": "zenith new", "new_name": "zenith \"new\"" } - ] + ], + "remote_extensions": { + "library_index": { + "anon": "anon", + "postgis-3": "postgis", + "libpgrouting-3.4": "postgis", + "postgis_raster-3": "postgis", + "postgis_sfcgal-3": "postgis", + "postgis_topology-3": "postgis", + "address_standardizer-3": "postgis" + }, + "extension_data": { + "anon": { + "archive_path": "5834329303/v15/extensions/anon.tar.zst", + "control_data": { + "anon.control": "# PostgreSQL Anonymizer (anon) extension\ncomment = ''Data anonymization tools''\ndefault_version = ''1.1.0''\ndirectory=''extension/anon''\nrelocatable = false\nrequires = ''pgcrypto''\nsuperuser = false\nmodule_pathname = ''$libdir/anon''\ntrusted = true\n" + } + }, + "postgis": { + "archive_path": "5834329303/v15/extensions/postgis.tar.zst", + "control_data": { + "postgis.control": "# postgis extension\ncomment = ''PostGIS geometry and geography spatial types and functions''\ndefault_version = ''3.3.2''\nmodule_pathname = ''$libdir/postgis-3''\nrelocatable = false\ntrusted = true\n", + "pgrouting.control": "# pgRouting Extension\ncomment = ''pgRouting Extension''\ndefault_version = ''3.4.2''\nmodule_pathname = ''$libdir/libpgrouting-3.4''\nrelocatable = true\nrequires = ''plpgsql''\nrequires = ''postgis''\ntrusted = true\n", + "postgis_raster.control": "# postgis_raster extension\ncomment = ''PostGIS raster types and functions''\ndefault_version = ''3.3.2''\nmodule_pathname = ''$libdir/postgis_raster-3''\nrelocatable = false\nrequires = postgis\ntrusted = true\n", + "postgis_sfcgal.control": "# postgis topology extension\ncomment = ''PostGIS SFCGAL functions''\ndefault_version = ''3.3.2''\nrelocatable = true\nrequires = postgis\ntrusted = true\n", + "postgis_topology.control": "# postgis topology extension\ncomment = ''PostGIS topology spatial types and functions''\ndefault_version = ''3.3.2''\nrelocatable = false\nschema = topology\nrequires = postgis\ntrusted = true\n", + "address_standardizer.control": "# address_standardizer extension\ncomment = ''Used to parse an address into constituent elements. Generally used to support geocoding address normalization step.''\ndefault_version = ''3.3.2''\nrelocatable = true\ntrusted = true\n", + "postgis_tiger_geocoder.control": "# postgis tiger geocoder extension\ncomment = ''PostGIS tiger geocoder and reverse geocoder''\ndefault_version = ''3.3.2''\nrelocatable = false\nschema = tiger\nrequires = ''postgis,fuzzystrmatch''\nsuperuser= false\ntrusted = true\n", + "address_standardizer_data_us.control": "# address standardizer us dataset\ncomment = ''Address Standardizer US dataset example''\ndefault_version = ''3.3.2''\nrelocatable = true\ntrusted = true\n" + } + } + }, + "custom_extensions": [ + "anon" + ], + "public_extensions": [ + "postgis" + ] + } }