diff --git a/Cargo.lock b/Cargo.lock index 4c464c62b8..fe4cc35029 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1284,6 +1284,7 @@ name = "compute_tools" version = "0.1.0" dependencies = [ "anyhow", + "async-compression", "aws-config", "aws-sdk-kms", "aws-sdk-s3", @@ -1420,6 +1421,7 @@ dependencies = [ "clap", "comfy-table", "compute_api", + "endpoint_storage", "futures", "http-utils", "humantime", diff --git a/Cargo.toml b/Cargo.toml index 1c203af9e0..8d4cc4a75a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -243,6 +243,7 @@ azure_storage_blobs = { git = "https://github.com/neondatabase/azure-sdk-for-rus ## Local libraries compute_api = { version = "0.1", path = "./libs/compute_api/" } consumption_metrics = { version = "0.1", path = "./libs/consumption_metrics/" } +endpoint_storage = { version = "0.0.1", path = "./endpoint_storage/" } http-utils = { version = "0.1", path = "./libs/http-utils/" } metrics = { version = "0.1", path = "./libs/metrics/" } pageserver = { path = "./pageserver" } diff --git a/compute_tools/Cargo.toml b/compute_tools/Cargo.toml index 8c1e7ad149..8ee5dd0665 100644 --- a/compute_tools/Cargo.toml +++ b/compute_tools/Cargo.toml @@ -10,6 +10,7 @@ default = [] testing = ["fail/failpoints"] [dependencies] +async-compression.workspace = true base64.workspace = true aws-config.workspace = true aws-sdk-s3.workspace = true diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 8834f0d63d..08d915b331 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -1,17 +1,10 @@ -use std::collections::HashMap; -use std::os::unix::fs::{PermissionsExt, symlink}; -use std::path::Path; -use std::process::{Command, Stdio}; -use std::str::FromStr; -use std::sync::atomic::{AtomicU32, Ordering}; -use std::sync::{Arc, Condvar, Mutex, RwLock}; -use std::time::{Duration, Instant}; -use std::{env, fs}; - use anyhow::{Context, Result}; use chrono::{DateTime, Utc}; use compute_api::privilege::Privilege; -use compute_api::responses::{ComputeConfig, ComputeCtlConfig, ComputeMetrics, ComputeStatus}; +use compute_api::responses::{ + ComputeConfig, ComputeCtlConfig, ComputeMetrics, ComputeStatus, LfcOffloadState, + LfcPrewarmState, +}; use compute_api::spec::{ ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, ExtVersion, PgIdent, }; @@ -25,6 +18,16 @@ use postgres; use postgres::NoTls; use postgres::error::SqlState; use remote_storage::{DownloadError, RemotePath}; +use std::collections::HashMap; +use std::net::SocketAddr; +use std::os::unix::fs::{PermissionsExt, symlink}; +use std::path::Path; +use std::process::{Command, Stdio}; +use std::str::FromStr; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::sync::{Arc, Condvar, Mutex, RwLock}; +use std::time::{Duration, Instant}; +use std::{env, fs}; use tokio::spawn; use tracing::{Instrument, debug, error, info, instrument, warn}; use utils::id::{TenantId, TimelineId}; @@ -150,6 +153,9 @@ pub struct ComputeState { /// set up the span relationship ourselves. pub startup_span: Option, + pub lfc_prewarm_state: LfcPrewarmState, + pub lfc_offload_state: LfcOffloadState, + pub metrics: ComputeMetrics, } @@ -163,6 +169,8 @@ impl ComputeState { pspec: None, startup_span: None, metrics: ComputeMetrics::default(), + lfc_prewarm_state: LfcPrewarmState::default(), + lfc_offload_state: LfcOffloadState::default(), } } @@ -198,6 +206,8 @@ pub struct ParsedSpec { pub pageserver_connstr: String, pub safekeeper_connstrings: Vec, pub storage_auth_token: Option, + pub endpoint_storage_addr: Option, + pub endpoint_storage_token: Option, } impl TryFrom for ParsedSpec { @@ -251,6 +261,18 @@ impl TryFrom for ParsedSpec { .or(Err("invalid timeline id"))? }; + let endpoint_storage_addr: Option = spec + .endpoint_storage_addr + .clone() + .or_else(|| spec.cluster.settings.find("neon.endpoint_storage_addr")) + .unwrap_or_default() + .parse() + .ok(); + let endpoint_storage_token = spec + .endpoint_storage_token + .clone() + .or_else(|| spec.cluster.settings.find("neon.endpoint_storage_token")); + Ok(ParsedSpec { spec, pageserver_connstr, @@ -258,6 +280,8 @@ impl TryFrom for ParsedSpec { storage_auth_token, tenant_id, timeline_id, + endpoint_storage_addr, + endpoint_storage_token, }) } } @@ -736,6 +760,9 @@ impl ComputeNode { // Log metrics so that we can search for slow operations in logs info!(?metrics, postmaster_pid = %postmaster_pid, "compute start finished"); + if pspec.spec.prewarm_lfc_on_startup { + self.prewarm_lfc(); + } Ok(()) } diff --git a/compute_tools/src/compute_prewarm.rs b/compute_tools/src/compute_prewarm.rs new file mode 100644 index 0000000000..a6a84b3f1f --- /dev/null +++ b/compute_tools/src/compute_prewarm.rs @@ -0,0 +1,202 @@ +use crate::compute::ComputeNode; +use anyhow::{Context, Result, bail}; +use async_compression::tokio::bufread::{ZstdDecoder, ZstdEncoder}; +use compute_api::responses::LfcOffloadState; +use compute_api::responses::LfcPrewarmState; +use http::StatusCode; +use reqwest::Client; +use std::sync::Arc; +use tokio::{io::AsyncReadExt, spawn}; +use tracing::{error, info}; + +#[derive(serde::Serialize, Default)] +pub struct LfcPrewarmStateWithProgress { + #[serde(flatten)] + base: LfcPrewarmState, + total: i32, + prewarmed: i32, + skipped: i32, +} + +/// A pair of url and a token to query endpoint storage for LFC prewarm-related tasks +struct EndpointStoragePair { + url: String, + token: String, +} + +const KEY: &str = "lfc_state"; +impl TryFrom<&crate::compute::ParsedSpec> for EndpointStoragePair { + type Error = anyhow::Error; + fn try_from(pspec: &crate::compute::ParsedSpec) -> Result { + let Some(ref endpoint_id) = pspec.spec.endpoint_id else { + bail!("pspec.endpoint_id missing") + }; + let Some(ref base_uri) = pspec.endpoint_storage_addr else { + bail!("pspec.endpoint_storage_addr missing") + }; + let tenant_id = pspec.tenant_id; + let timeline_id = pspec.timeline_id; + + let url = format!("http://{base_uri}/{tenant_id}/{timeline_id}/{endpoint_id}/{KEY}"); + let Some(ref token) = pspec.endpoint_storage_token else { + bail!("pspec.endpoint_storage_token missing") + }; + let token = token.clone(); + Ok(EndpointStoragePair { url, token }) + } +} + +impl ComputeNode { + // If prewarm failed, we want to get overall number of segments as well as done ones. + // However, this function should be reliable even if querying postgres failed. + pub async fn lfc_prewarm_state(&self) -> LfcPrewarmStateWithProgress { + info!("requesting LFC prewarm state from postgres"); + let mut state = LfcPrewarmStateWithProgress::default(); + { + state.base = self.state.lock().unwrap().lfc_prewarm_state.clone(); + } + + let client = match ComputeNode::get_maintenance_client(&self.tokio_conn_conf).await { + Ok(client) => client, + Err(err) => { + error!(%err, "connecting to postgres"); + return state; + } + }; + let row = match client + .query_one("select * from get_prewarm_info()", &[]) + .await + { + Ok(row) => row, + Err(err) => { + error!(%err, "querying LFC prewarm status"); + return state; + } + }; + state.total = row.try_get(0).unwrap_or_default(); + state.prewarmed = row.try_get(1).unwrap_or_default(); + state.skipped = row.try_get(2).unwrap_or_default(); + state + } + + pub fn lfc_offload_state(&self) -> LfcOffloadState { + self.state.lock().unwrap().lfc_offload_state.clone() + } + + /// Returns false if there is a prewarm request ongoing, true otherwise + pub fn prewarm_lfc(self: &Arc) -> bool { + crate::metrics::LFC_PREWARM_REQUESTS.inc(); + { + let state = &mut self.state.lock().unwrap().lfc_prewarm_state; + if let LfcPrewarmState::Prewarming = + std::mem::replace(state, LfcPrewarmState::Prewarming) + { + return false; + } + } + + let cloned = self.clone(); + spawn(async move { + let Err(err) = cloned.prewarm_impl().await else { + cloned.state.lock().unwrap().lfc_prewarm_state = LfcPrewarmState::Completed; + return; + }; + error!(%err); + cloned.state.lock().unwrap().lfc_prewarm_state = LfcPrewarmState::Failed { + error: err.to_string(), + }; + }); + true + } + + fn endpoint_storage_pair(&self) -> Result { + let state = self.state.lock().unwrap(); + state.pspec.as_ref().unwrap().try_into() + } + + async fn prewarm_impl(&self) -> Result<()> { + let EndpointStoragePair { url, token } = self.endpoint_storage_pair()?; + info!(%url, "requesting LFC state from endpoint storage"); + + let request = Client::new().get(&url).bearer_auth(token); + let res = request.send().await.context("querying endpoint storage")?; + let status = res.status(); + if status != StatusCode::OK { + bail!("{status} querying endpoint storage") + } + + let mut uncompressed = Vec::new(); + let lfc_state = res + .bytes() + .await + .context("getting request body from endpoint storage")?; + ZstdDecoder::new(lfc_state.iter().as_slice()) + .read_to_end(&mut uncompressed) + .await + .context("decoding LFC state")?; + let uncompressed_len = uncompressed.len(); + info!(%url, "downloaded LFC state, uncompressed size {uncompressed_len}, loading into postgres"); + + ComputeNode::get_maintenance_client(&self.tokio_conn_conf) + .await + .context("connecting to postgres")? + .query_one("select prewarm_local_cache($1)", &[&uncompressed]) + .await + .context("loading LFC state into postgres") + .map(|_| ()) + } + + /// Returns false if there is an offload request ongoing, true otherwise + pub fn offload_lfc(self: &Arc) -> bool { + crate::metrics::LFC_OFFLOAD_REQUESTS.inc(); + { + let state = &mut self.state.lock().unwrap().lfc_offload_state; + if let LfcOffloadState::Offloading = + std::mem::replace(state, LfcOffloadState::Offloading) + { + return false; + } + } + + let cloned = self.clone(); + spawn(async move { + let Err(err) = cloned.offload_lfc_impl().await else { + cloned.state.lock().unwrap().lfc_offload_state = LfcOffloadState::Completed; + return; + }; + error!(%err); + cloned.state.lock().unwrap().lfc_offload_state = LfcOffloadState::Failed { + error: err.to_string(), + }; + }); + true + } + + async fn offload_lfc_impl(&self) -> Result<()> { + let EndpointStoragePair { url, token } = self.endpoint_storage_pair()?; + info!(%url, "requesting LFC state from postgres"); + + let mut compressed = Vec::new(); + ComputeNode::get_maintenance_client(&self.tokio_conn_conf) + .await + .context("connecting to postgres")? + .query_one("select get_local_cache_state()", &[]) + .await + .context("querying LFC state")? + .try_get::(0) + .context("deserializing LFC state") + .map(ZstdEncoder::new)? + .read_to_end(&mut compressed) + .await + .context("compressing LFC state")?; + let compressed_len = compressed.len(); + info!(%url, "downloaded LFC state, compressed size {compressed_len}, writing to endpoint storage"); + + let request = Client::new().put(url).bearer_auth(token).body(compressed); + match request.send().await { + Ok(res) if res.status() == StatusCode::OK => Ok(()), + Ok(res) => bail!("Error writing to endpoint storage: {}", res.status()), + Err(err) => Err(err).context("writing to endpoint storage"), + } + } +} diff --git a/compute_tools/src/http/routes/lfc.rs b/compute_tools/src/http/routes/lfc.rs new file mode 100644 index 0000000000..07bcc6bfb7 --- /dev/null +++ b/compute_tools/src/http/routes/lfc.rs @@ -0,0 +1,39 @@ +use crate::compute_prewarm::LfcPrewarmStateWithProgress; +use crate::http::JsonResponse; +use axum::response::{IntoResponse, Response}; +use axum::{Json, http::StatusCode}; +use compute_api::responses::LfcOffloadState; +type Compute = axum::extract::State>; + +pub(in crate::http) async fn prewarm_state(compute: Compute) -> Json { + Json(compute.lfc_prewarm_state().await) +} + +// Following functions are marked async for axum, as it's more convenient than wrapping these +// in async lambdas at call site + +pub(in crate::http) async fn offload_state(compute: Compute) -> Json { + Json(compute.lfc_offload_state()) +} + +pub(in crate::http) async fn prewarm(compute: Compute) -> Response { + if compute.prewarm_lfc() { + StatusCode::ACCEPTED.into_response() + } else { + JsonResponse::error( + StatusCode::TOO_MANY_REQUESTS, + "Multiple requests for prewarm are not allowed", + ) + } +} + +pub(in crate::http) async fn offload(compute: Compute) -> Response { + if compute.offload_lfc() { + StatusCode::ACCEPTED.into_response() + } else { + JsonResponse::error( + StatusCode::TOO_MANY_REQUESTS, + "Multiple requests for prewarm offload are not allowed", + ) + } +} diff --git a/compute_tools/src/http/routes/mod.rs b/compute_tools/src/http/routes/mod.rs index a67be7fd5a..432e66a830 100644 --- a/compute_tools/src/http/routes/mod.rs +++ b/compute_tools/src/http/routes/mod.rs @@ -11,6 +11,7 @@ pub(in crate::http) mod extensions; pub(in crate::http) mod failpoints; pub(in crate::http) mod grants; pub(in crate::http) mod insights; +pub(in crate::http) mod lfc; pub(in crate::http) mod metrics; pub(in crate::http) mod metrics_json; pub(in crate::http) mod status; diff --git a/compute_tools/src/http/server.rs b/compute_tools/src/http/server.rs index 10f767e97c..d5d2427971 100644 --- a/compute_tools/src/http/server.rs +++ b/compute_tools/src/http/server.rs @@ -23,7 +23,7 @@ use super::{ middleware::authorize::Authorize, routes::{ check_writability, configure, database_schema, dbs_and_roles, extension_server, extensions, - grants, insights, metrics, metrics_json, status, terminate, + grants, insights, lfc, metrics, metrics_json, status, terminate, }, }; use crate::compute::ComputeNode; @@ -85,6 +85,8 @@ impl From<&Server> for Router> { Router::>::new().route("/metrics", get(metrics::get_metrics)); let authenticated_router = Router::>::new() + .route("/lfc/prewarm", get(lfc::prewarm_state).post(lfc::prewarm)) + .route("/lfc/offload", get(lfc::offload_state).post(lfc::offload)) .route("/check_writability", post(check_writability::is_writable)) .route("/configure", post(configure::configure)) .route("/database_schema", get(database_schema::get_schema_dump)) diff --git a/compute_tools/src/lib.rs b/compute_tools/src/lib.rs index a681fad0b0..7218067a8a 100644 --- a/compute_tools/src/lib.rs +++ b/compute_tools/src/lib.rs @@ -11,6 +11,7 @@ pub mod http; pub mod logger; pub mod catalog; pub mod compute; +pub mod compute_prewarm; pub mod disk_quota; pub mod extension_server; pub mod installed_extensions; diff --git a/compute_tools/src/metrics.rs b/compute_tools/src/metrics.rs index e37d6120ac..90326b2074 100644 --- a/compute_tools/src/metrics.rs +++ b/compute_tools/src/metrics.rs @@ -1,7 +1,7 @@ use metrics::core::{AtomicF64, AtomicU64, Collector, GenericCounter, GenericGauge}; use metrics::proto::MetricFamily; use metrics::{ - IntCounterVec, IntGaugeVec, UIntGaugeVec, register_gauge, register_int_counter, + IntCounter, IntCounterVec, IntGaugeVec, UIntGaugeVec, register_gauge, register_int_counter, register_int_counter_vec, register_int_gauge_vec, register_uint_gauge_vec, }; use once_cell::sync::Lazy; @@ -97,6 +97,24 @@ pub(crate) static PG_TOTAL_DOWNTIME_MS: Lazy> = Lazy:: .expect("failed to define a metric") }); +/// Needed as neon.file_cache_prewarm_batch == 0 doesn't mean we never tried to prewarm. +/// On the other hand, LFC_PREWARMED_PAGES is excessive as we can GET /lfc/prewarm +pub(crate) static LFC_PREWARM_REQUESTS: Lazy = Lazy::new(|| { + register_int_counter!( + "compute_ctl_lfc_prewarm_requests_total", + "Total number of LFC prewarm requests made by compute_ctl", + ) + .expect("failed to define a metric") +}); + +pub(crate) static LFC_OFFLOAD_REQUESTS: Lazy = Lazy::new(|| { + register_int_counter!( + "compute_ctl_lfc_offload_requests_total", + "Total number of LFC offload requests made by compute_ctl", + ) + .expect("failed to define a metric") +}); + pub fn collect() -> Vec { let mut metrics = COMPUTE_CTL_UP.collect(); metrics.extend(INSTALLED_EXTENSIONS.collect()); @@ -106,5 +124,7 @@ pub fn collect() -> Vec { metrics.extend(AUDIT_LOG_DIR_SIZE.collect()); metrics.extend(PG_CURR_DOWNTIME_MS.collect()); metrics.extend(PG_TOTAL_DOWNTIME_MS.collect()); + metrics.extend(LFC_PREWARM_REQUESTS.collect()); + metrics.extend(LFC_OFFLOAD_REQUESTS.collect()); metrics } diff --git a/compute_tools/tests/pg_helpers_tests.rs b/compute_tools/tests/pg_helpers_tests.rs index b72c1293ee..53f2ddad84 100644 --- a/compute_tools/tests/pg_helpers_tests.rs +++ b/compute_tools/tests/pg_helpers_tests.rs @@ -30,6 +30,7 @@ mod pg_helpers_tests { r#"fsync = off wal_level = logical hot_standby = on +prewarm_lfc_on_startup = off neon.safekeepers = '127.0.0.1:6502,127.0.0.1:6503,127.0.0.1:6501' wal_log_hints = on log_connections = on diff --git a/control_plane/Cargo.toml b/control_plane/Cargo.toml index 92f0071bac..62c039047f 100644 --- a/control_plane/Cargo.toml +++ b/control_plane/Cargo.toml @@ -41,7 +41,7 @@ storage_broker.workspace = true http-utils.workspace = true utils.workspace = true whoami.workspace = true - +endpoint_storage.workspace = true compute_api.workspace = true workspace_hack.workspace = true tracing.workspace = true diff --git a/control_plane/src/bin/neon_local.rs b/control_plane/src/bin/neon_local.rs index 44698f7b23..fd625e9ed6 100644 --- a/control_plane/src/bin/neon_local.rs +++ b/control_plane/src/bin/neon_local.rs @@ -20,7 +20,7 @@ use compute_api::requests::ComputeClaimsScope; use compute_api::spec::ComputeMode; use control_plane::broker::StorageBroker; use control_plane::endpoint::ComputeControlPlane; -use control_plane::endpoint_storage::{ENDPOINT_STORAGE_DEFAULT_PORT, EndpointStorage}; +use control_plane::endpoint_storage::{ENDPOINT_STORAGE_DEFAULT_ADDR, EndpointStorage}; use control_plane::local_env; use control_plane::local_env::{ EndpointStorageConf, InitForceMode, LocalEnv, NeonBroker, NeonLocalInitConf, @@ -1022,7 +1022,7 @@ fn handle_init(args: &InitCmdArgs) -> anyhow::Result { }) .collect(), endpoint_storage: EndpointStorageConf { - port: ENDPOINT_STORAGE_DEFAULT_PORT, + listen_addr: ENDPOINT_STORAGE_DEFAULT_ADDR, }, pg_distrib_dir: None, neon_distrib_dir: None, @@ -1488,10 +1488,25 @@ async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Res None }; + let exp = (std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)? + + Duration::from_secs(86400)) + .as_secs(); + let claims = endpoint_storage::claims::EndpointStorageClaims { + tenant_id: endpoint.tenant_id, + timeline_id: endpoint.timeline_id, + endpoint_id: endpoint_id.to_string(), + exp, + }; + + let endpoint_storage_token = env.generate_auth_token(&claims)?; + let endpoint_storage_addr = env.endpoint_storage.listen_addr.to_string(); + println!("Starting existing endpoint {endpoint_id}..."); endpoint .start( &auth_token, + endpoint_storage_token, + endpoint_storage_addr, safekeepers_generation, safekeepers, pageservers, diff --git a/control_plane/src/endpoint.rs b/control_plane/src/endpoint.rs index 0b16339a6f..fe6a93eb5e 100644 --- a/control_plane/src/endpoint.rs +++ b/control_plane/src/endpoint.rs @@ -650,6 +650,8 @@ impl Endpoint { pub async fn start( &self, auth_token: &Option, + endpoint_storage_token: String, + endpoint_storage_addr: String, safekeepers_generation: Option, safekeepers: Vec, pageservers: Vec<(Host, u16)>, @@ -743,6 +745,9 @@ impl Endpoint { drop_subscriptions_before_start: self.drop_subscriptions_before_start, audit_log_level: ComputeAudit::Disabled, logs_export_host: None::, + endpoint_storage_addr: Some(endpoint_storage_addr), + endpoint_storage_token: Some(endpoint_storage_token), + prewarm_lfc_on_startup: false, }; // this strange code is needed to support respec() in tests diff --git a/control_plane/src/endpoint_storage.rs b/control_plane/src/endpoint_storage.rs index 102db91a22..171aaeddb4 100644 --- a/control_plane/src/endpoint_storage.rs +++ b/control_plane/src/endpoint_storage.rs @@ -3,17 +3,19 @@ use crate::local_env::LocalEnv; use anyhow::{Context, Result}; use camino::Utf8PathBuf; use std::io::Write; +use std::net::SocketAddr; use std::time::Duration; /// Directory within .neon which will be used by default for LocalFs remote storage. pub const ENDPOINT_STORAGE_REMOTE_STORAGE_DIR: &str = "local_fs_remote_storage/endpoint_storage"; -pub const ENDPOINT_STORAGE_DEFAULT_PORT: u16 = 9993; +pub const ENDPOINT_STORAGE_DEFAULT_ADDR: SocketAddr = + SocketAddr::new(std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST), 9993); pub struct EndpointStorage { pub bin: Utf8PathBuf, pub data_dir: Utf8PathBuf, pub pemfile: Utf8PathBuf, - pub port: u16, + pub addr: SocketAddr, } impl EndpointStorage { @@ -22,7 +24,7 @@ impl EndpointStorage { bin: Utf8PathBuf::from_path_buf(env.endpoint_storage_bin()).unwrap(), data_dir: Utf8PathBuf::from_path_buf(env.endpoint_storage_data_dir()).unwrap(), pemfile: Utf8PathBuf::from_path_buf(env.public_key_path.clone()).unwrap(), - port: env.endpoint_storage.port, + addr: env.endpoint_storage.listen_addr, } } @@ -31,7 +33,7 @@ impl EndpointStorage { } fn listen_addr(&self) -> Utf8PathBuf { - format!("127.0.0.1:{}", self.port).into() + format!("{}:{}", self.addr.ip(), self.addr.port()).into() } pub fn init(&self) -> Result<()> { diff --git a/control_plane/src/local_env.rs b/control_plane/src/local_env.rs index a18b34daa4..4a8892c6de 100644 --- a/control_plane/src/local_env.rs +++ b/control_plane/src/local_env.rs @@ -20,7 +20,9 @@ use utils::auth::encode_from_key_file; use utils::id::{NodeId, TenantId, TenantTimelineId, TimelineId}; use crate::broker::StorageBroker; -use crate::endpoint_storage::{ENDPOINT_STORAGE_REMOTE_STORAGE_DIR, EndpointStorage}; +use crate::endpoint_storage::{ + ENDPOINT_STORAGE_DEFAULT_ADDR, ENDPOINT_STORAGE_REMOTE_STORAGE_DIR, EndpointStorage, +}; use crate::pageserver::{PAGESERVER_REMOTE_STORAGE_DIR, PageServerNode}; use crate::safekeeper::SafekeeperNode; @@ -151,10 +153,10 @@ pub struct NeonLocalInitConf { pub generate_local_ssl_certs: bool, } -#[derive(Serialize, Default, Deserialize, PartialEq, Eq, Clone, Debug)] +#[derive(Serialize, Deserialize, PartialEq, Eq, Clone, Debug)] #[serde(default)] pub struct EndpointStorageConf { - pub port: u16, + pub listen_addr: SocketAddr, } /// Broker config for cluster internal communication. @@ -241,6 +243,14 @@ impl Default for NeonStorageControllerConf { } } +impl Default for EndpointStorageConf { + fn default() -> Self { + Self { + listen_addr: ENDPOINT_STORAGE_DEFAULT_ADDR, + } + } +} + impl NeonBroker { pub fn client_url(&self) -> Url { let url = if let Some(addr) = self.listen_https_addr { diff --git a/endpoint_storage/src/app.rs b/endpoint_storage/src/app.rs index f07ef06328..0bd7fe5f28 100644 --- a/endpoint_storage/src/app.rs +++ b/endpoint_storage/src/app.rs @@ -343,7 +343,7 @@ MC4CAQAwBQYDK2VwBCIEID/Drmc1AA6U/znNRWpF3zEGegOATQxfkdWxitcOMsIH TimelineId::from_array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 7]); const ENDPOINT_ID: &str = "ep-winter-frost-a662z3vg"; fn token() -> String { - let claims = endpoint_storage::Claims { + let claims = endpoint_storage::claims::EndpointStorageClaims { tenant_id: TENANT_ID, timeline_id: TIMELINE_ID, endpoint_id: ENDPOINT_ID.into(), @@ -489,16 +489,8 @@ MC4CAQAwBQYDK2VwBCIEID/Drmc1AA6U/znNRWpF3zEGegOATQxfkdWxitcOMsIH } fn delete_prefix_token(uri: &str) -> String { - use serde::Serialize; let parts = uri.split("/").collect::>(); - #[derive(Serialize)] - struct PrefixClaims { - tenant_id: TenantId, - timeline_id: Option, - endpoint_id: Option, - exp: u64, - } - let claims = PrefixClaims { + let claims = endpoint_storage::claims::DeletePrefixClaims { tenant_id: parts.get(1).map(|c| c.parse().unwrap()).unwrap(), timeline_id: parts.get(2).map(|c| c.parse().unwrap()), endpoint_id: parts.get(3).map(ToString::to_string), diff --git a/endpoint_storage/src/claims.rs b/endpoint_storage/src/claims.rs new file mode 100644 index 0000000000..ef0f0eb0b4 --- /dev/null +++ b/endpoint_storage/src/claims.rs @@ -0,0 +1,52 @@ +use serde::{Deserialize, Serialize}; +use std::fmt::Display; +use utils::id::{EndpointId, TenantId, TimelineId}; + +/// Claims to add, remove, or retrieve endpoint data. Used by compute_ctl +#[derive(Deserialize, Serialize, PartialEq)] +pub struct EndpointStorageClaims { + pub tenant_id: TenantId, + pub timeline_id: TimelineId, + pub endpoint_id: EndpointId, + pub exp: u64, +} + +/// Claims to remove tenant, timeline, or endpoint data. Used by control plane +#[derive(Deserialize, Serialize, PartialEq)] +pub struct DeletePrefixClaims { + pub tenant_id: TenantId, + /// None when tenant is deleted (endpoint_id is also None in this case) + pub timeline_id: Option, + /// None when timeline is deleted + pub endpoint_id: Option, + pub exp: u64, +} + +impl Display for EndpointStorageClaims { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "EndpointClaims(tenant_id={} timeline_id={} endpoint_id={} exp={})", + self.tenant_id, self.timeline_id, self.endpoint_id, self.exp + ) + } +} + +impl Display for DeletePrefixClaims { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "DeletePrefixClaims(tenant_id={} timeline_id={} endpoint_id={}, exp={})", + self.tenant_id, + self.timeline_id + .as_ref() + .map(ToString::to_string) + .unwrap_or("".to_string()), + self.endpoint_id + .as_ref() + .map(ToString::to_string) + .unwrap_or("".to_string()), + self.exp + ) + } +} diff --git a/endpoint_storage/src/lib.rs b/endpoint_storage/src/lib.rs index eb6b80c487..d1625dc843 100644 --- a/endpoint_storage/src/lib.rs +++ b/endpoint_storage/src/lib.rs @@ -1,3 +1,5 @@ +pub mod claims; +use crate::claims::{DeletePrefixClaims, EndpointStorageClaims}; use anyhow::Result; use axum::extract::{FromRequestParts, Path}; use axum::response::{IntoResponse, Response}; @@ -13,7 +15,7 @@ use std::result::Result as StdResult; use std::sync::Arc; use tokio_util::sync::CancellationToken; use tracing::{debug, error}; -use utils::id::{TenantId, TimelineId}; +use utils::id::{EndpointId, TenantId, TimelineId}; // simplified version of utils::auth::JwtAuth pub struct JwtAuth { @@ -79,26 +81,6 @@ pub struct Storage { pub max_upload_file_limit: usize, } -pub type EndpointId = String; // If needed, reuse small string from proxy/src/types.rc - -#[derive(Deserialize, Serialize, PartialEq)] -pub struct Claims { - pub tenant_id: TenantId, - pub timeline_id: TimelineId, - pub endpoint_id: EndpointId, - pub exp: u64, -} - -impl Display for Claims { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "Claims(tenant_id {} timeline_id {} endpoint_id {} exp {})", - self.tenant_id, self.timeline_id, self.endpoint_id, self.exp - ) - } -} - #[derive(Deserialize, Serialize)] struct KeyRequest { tenant_id: TenantId, @@ -107,6 +89,13 @@ struct KeyRequest { path: String, } +#[derive(Deserialize, Serialize, PartialEq)] +struct PrefixKeyRequest { + tenant_id: TenantId, + timeline_id: Option, + endpoint_id: Option, +} + #[derive(Debug, PartialEq)] pub struct S3Path { pub path: RemotePath, @@ -165,7 +154,7 @@ impl FromRequestParts> for S3Path { .extract::>>() .await .map_err(|e| bad_request(e, "invalid token"))?; - let claims: Claims = state + let claims: EndpointStorageClaims = state .auth .decode(bearer.token()) .map_err(|e| bad_request(e, "decoding token"))?; @@ -178,7 +167,7 @@ impl FromRequestParts> for S3Path { path.endpoint_id.clone() }; - let route = Claims { + let route = EndpointStorageClaims { tenant_id: path.tenant_id, timeline_id: path.timeline_id, endpoint_id, @@ -193,38 +182,13 @@ impl FromRequestParts> for S3Path { } } -#[derive(Deserialize, Serialize, PartialEq)] -pub struct PrefixKeyPath { - pub tenant_id: TenantId, - pub timeline_id: Option, - pub endpoint_id: Option, -} - -impl Display for PrefixKeyPath { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "PrefixKeyPath(tenant_id {} timeline_id {} endpoint_id {})", - self.tenant_id, - self.timeline_id - .as_ref() - .map(ToString::to_string) - .unwrap_or("".to_string()), - self.endpoint_id - .as_ref() - .map(ToString::to_string) - .unwrap_or("".to_string()) - ) - } -} - #[derive(Debug, PartialEq)] pub struct PrefixS3Path { pub path: RemotePath, } -impl From<&PrefixKeyPath> for PrefixS3Path { - fn from(path: &PrefixKeyPath) -> Self { +impl From<&DeletePrefixClaims> for PrefixS3Path { + fn from(path: &DeletePrefixClaims) -> Self { let timeline_id = path .timeline_id .as_ref() @@ -250,21 +214,27 @@ impl FromRequestParts> for PrefixS3Path { state: &Arc, ) -> Result { let Path(path) = parts - .extract::>() + .extract::>() .await .map_err(|e| bad_request(e, "invalid route"))?; let TypedHeader(Authorization(bearer)) = parts .extract::>>() .await .map_err(|e| bad_request(e, "invalid token"))?; - let claims: PrefixKeyPath = state + let claims: DeletePrefixClaims = state .auth .decode(bearer.token()) .map_err(|e| bad_request(e, "invalid token"))?; - if path != claims { - return Err(unauthorized(path, claims)); + let route = DeletePrefixClaims { + tenant_id: path.tenant_id, + timeline_id: path.timeline_id, + endpoint_id: path.endpoint_id, + exp: claims.exp, + }; + if route != claims { + return Err(unauthorized(route, claims)); } - Ok((&path).into()) + Ok((&route).into()) } } @@ -297,7 +267,7 @@ mod tests { #[test] fn s3_path() { - let auth = Claims { + let auth = EndpointStorageClaims { tenant_id: TENANT_ID, timeline_id: TIMELINE_ID, endpoint_id: ENDPOINT_ID.into(), @@ -327,10 +297,11 @@ mod tests { #[test] fn prefix_s3_path() { - let mut path = PrefixKeyPath { + let mut path = DeletePrefixClaims { tenant_id: TENANT_ID, timeline_id: None, endpoint_id: None, + exp: 0, }; let prefix_path = |s: String| RemotePath::from_string(&s).unwrap(); assert_eq!( diff --git a/libs/compute_api/src/responses.rs b/libs/compute_api/src/responses.rs index b7d6b7ca34..24d371c6eb 100644 --- a/libs/compute_api/src/responses.rs +++ b/libs/compute_api/src/responses.rs @@ -46,6 +46,30 @@ pub struct ExtensionInstallResponse { pub version: ExtVersion, } +#[derive(Serialize, Default, Debug, Clone)] +#[serde(tag = "status", rename_all = "snake_case")] +pub enum LfcPrewarmState { + #[default] + NotPrewarmed, + Prewarming, + Completed, + Failed { + error: String, + }, +} + +#[derive(Serialize, Default, Debug, Clone)] +#[serde(tag = "status", rename_all = "snake_case")] +pub enum LfcOffloadState { + #[default] + NotOffloaded, + Offloading, + Completed, + Failed { + error: String, + }, +} + /// Response of the /status API #[derive(Serialize, Debug, Deserialize)] #[serde(rename_all = "snake_case")] diff --git a/libs/compute_api/src/spec.rs b/libs/compute_api/src/spec.rs index ad246c48ec..09b550b96c 100644 --- a/libs/compute_api/src/spec.rs +++ b/libs/compute_api/src/spec.rs @@ -172,6 +172,15 @@ pub struct ComputeSpec { /// Hostname and the port of the otel collector. Leave empty to disable Postgres logs forwarding. /// Example: config-shy-breeze-123-collector-monitoring.neon-telemetry.svc.cluster.local:10514 pub logs_export_host: Option, + + /// Address of endpoint storage service + pub endpoint_storage_addr: Option, + /// JWT for authorizing requests to endpoint storage service + pub endpoint_storage_token: Option, + + /// If true, download LFC state from endpoint_storage and pass it to Postgres on startup + #[serde(default)] + pub prewarm_lfc_on_startup: bool, } /// Feature flag to signal `compute_ctl` to enable certain experimental functionality. diff --git a/libs/compute_api/tests/cluster_spec.json b/libs/compute_api/tests/cluster_spec.json index 37de24be5b..30e788a601 100644 --- a/libs/compute_api/tests/cluster_spec.json +++ b/libs/compute_api/tests/cluster_spec.json @@ -84,6 +84,11 @@ "value": "on", "vartype": "bool" }, + { + "name": "prewarm_lfc_on_startup", + "value": "off", + "vartype": "bool" + }, { "name": "neon.safekeepers", "value": "127.0.0.1:6502,127.0.0.1:6503,127.0.0.1:6501", diff --git a/libs/utils/src/id.rs b/libs/utils/src/id.rs index 6016c23a01..68cb1f0209 100644 --- a/libs/utils/src/id.rs +++ b/libs/utils/src/id.rs @@ -295,6 +295,9 @@ pub struct TenantId(Id); id_newtype!(TenantId); +/// If needed, reuse small string from proxy/src/types.rc +pub type EndpointId = String; + // A pair uniquely identifying Neon instance. #[derive(Debug, Clone, Copy, PartialOrd, Ord, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct TenantTimelineId { diff --git a/test_runner/fixtures/endpoint/http.py b/test_runner/fixtures/endpoint/http.py index beed1dcd93..4b4b98aa6c 100644 --- a/test_runner/fixtures/endpoint/http.py +++ b/test_runner/fixtures/endpoint/http.py @@ -10,6 +10,7 @@ from requests.auth import AuthBase from typing_extensions import override from fixtures.log_helper import log +from fixtures.utils import wait_until if TYPE_CHECKING: from requests import PreparedRequest @@ -62,6 +63,35 @@ class EndpointHttpClient(requests.Session): res.raise_for_status() return res.json() + def prewarm_lfc_status(self) -> dict[str, str]: + res = self.get(f"http://localhost:{self.external_port}/lfc/prewarm") + res.raise_for_status() + json: dict[str, str] = res.json() + return json + + def prewarm_lfc(self): + self.post(f"http://localhost:{self.external_port}/lfc/prewarm").raise_for_status() + + def prewarmed(): + json = self.prewarm_lfc_status() + status, err = json["status"], json.get("error") + assert status == "completed", f"{status}, error {err}" + + wait_until(prewarmed) + + def offload_lfc(self): + url = f"http://localhost:{self.external_port}/lfc/offload" + self.post(url).raise_for_status() + + def offloaded(): + res = self.get(url) + res.raise_for_status() + json = res.json() + status, err = json["status"], json.get("error") + assert status == "completed", f"{status}, error {err}" + + wait_until(offloaded) + def database_schema(self, database: str): res = self.get( f"http://localhost:{self.external_port}/database_schema?database={urllib.parse.quote(database, safe='')}", diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index 133be5c045..d4a750ad3b 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -1185,7 +1185,9 @@ class NeonEnv: "broker": {}, "safekeepers": [], "pageservers": [], - "endpoint_storage": {"port": self.port_distributor.get_port()}, + "endpoint_storage": { + "listen_addr": f"127.0.0.1:{self.port_distributor.get_port()}", + }, "generate_local_ssl_certs": self.generate_local_ssl_certs, } diff --git a/test_runner/regress/test_endpoint_storage.py b/test_runner/regress/test_endpoint_storage.py index 04029114ec..1e27ef4b14 100644 --- a/test_runner/regress/test_endpoint_storage.py +++ b/test_runner/regress/test_endpoint_storage.py @@ -4,10 +4,12 @@ import pytest from aiohttp import ClientSession from fixtures.log_helper import log from fixtures.neon_fixtures import NeonEnv +from fixtures.utils import run_only_on_default_postgres from jwcrypto import jwk, jwt @pytest.mark.asyncio +@run_only_on_default_postgres("test doesn't use postgres") async def test_endpoint_storage_insert_retrieve_delete(neon_simple_env: NeonEnv): """ Inserts, retrieves, and deletes test file using a JWT token @@ -35,7 +37,6 @@ async def test_endpoint_storage_insert_retrieve_delete(neon_simple_env: NeonEnv) key = f"http://{base_url}/{tenant_id}/{timeline_id}/{endpoint_id}/key" headers = {"Authorization": f"Bearer {token}"} log.info(f"cache key url {key}") - log.info(f"token {token}") async with ClientSession(headers=headers) as session: async with session.get(key) as res: diff --git a/test_runner/regress/test_lfc_prewarm.py b/test_runner/regress/test_lfc_prewarm.py index dd0ae1921d..82e1e9fcba 100644 --- a/test_runner/regress/test_lfc_prewarm.py +++ b/test_runner/regress/test_lfc_prewarm.py @@ -1,11 +1,24 @@ import random import threading import time +from enum import Enum import pytest +from fixtures.endpoint.http import EndpointHttpClient from fixtures.log_helper import log from fixtures.neon_fixtures import NeonEnv from fixtures.utils import USE_LFC +from prometheus_client.parser import text_string_to_metric_families as prom_parse_impl + + +class LfcQueryMethod(Enum): + COMPUTE_CTL = False + POSTGRES = True + + +PREWARM_LABEL = "compute_ctl_lfc_prewarm_requests_total" +OFFLOAD_LABEL = "compute_ctl_lfc_offload_requests_total" +QUERY_OPTIONS = LfcQueryMethod.POSTGRES, LfcQueryMethod.COMPUTE_CTL def check_pinned_entries(cur): @@ -19,11 +32,20 @@ def check_pinned_entries(cur): assert n_pinned == 0 +def prom_parse(client: EndpointHttpClient) -> dict[str, float]: + return { + sample.name: sample.value + for family in prom_parse_impl(client.metrics()) + for sample in family.samples + if sample.name in (PREWARM_LABEL, OFFLOAD_LABEL) + } + + @pytest.mark.skipif(not USE_LFC, reason="LFC is disabled, skipping") -def test_lfc_prewarm(neon_simple_env: NeonEnv): +@pytest.mark.parametrize("query", QUERY_OPTIONS, ids=["postgres", "compute-ctl"]) +def test_lfc_prewarm(neon_simple_env: NeonEnv, query: LfcQueryMethod): env = neon_simple_env n_records = 1000000 - endpoint = env.endpoints.create_start( branch_name="main", config_lines=[ @@ -34,30 +56,57 @@ def test_lfc_prewarm(neon_simple_env: NeonEnv): "neon.file_cache_prewarm_limit=1000", ], ) - conn = endpoint.connect() - cur = conn.cursor() - cur.execute("create extension neon version '1.6'") - cur.execute("create table t(pk integer primary key, payload text default repeat('?', 128))") - cur.execute(f"insert into t (pk) values (generate_series(1,{n_records}))") - cur.execute("select get_local_cache_state()") - lfc_state = cur.fetchall()[0][0] + + pg_conn = endpoint.connect() + pg_cur = pg_conn.cursor() + pg_cur.execute("create extension neon version '1.6'") + pg_cur.execute("create database lfc") + + lfc_conn = endpoint.connect(dbname="lfc") + lfc_cur = lfc_conn.cursor() + log.info(f"Inserting {n_records} rows") + lfc_cur.execute("create table t(pk integer primary key, payload text default repeat('?', 128))") + lfc_cur.execute(f"insert into t (pk) values (generate_series(1,{n_records}))") + log.info(f"Inserted {n_records} rows") + + http_client = endpoint.http_client() + if query is LfcQueryMethod.COMPUTE_CTL: + status = http_client.prewarm_lfc_status() + assert status["status"] == "not_prewarmed" + assert "error" not in status + http_client.offload_lfc() + assert http_client.prewarm_lfc_status()["status"] == "not_prewarmed" + assert prom_parse(http_client) == {OFFLOAD_LABEL: 1, PREWARM_LABEL: 0} + else: + pg_cur.execute("select get_local_cache_state()") + lfc_state = pg_cur.fetchall()[0][0] endpoint.stop() endpoint.start() - conn = endpoint.connect() - cur = conn.cursor() - time.sleep(1) # wait until compute_ctl complete downgrade of extension to default version - cur.execute("alter extension neon update to '1.6'") - cur.execute("select prewarm_local_cache(%s)", (lfc_state,)) + # wait until compute_ctl completes downgrade of extension to default version + time.sleep(1) + pg_conn = endpoint.connect() + pg_cur = pg_conn.cursor() + pg_cur.execute("alter extension neon update to '1.6'") - cur.execute("select lfc_value from neon_lfc_stats where lfc_key='file_cache_used_pages'") - lfc_used_pages = cur.fetchall()[0][0] + lfc_conn = endpoint.connect(dbname="lfc") + lfc_cur = lfc_conn.cursor() + + if query is LfcQueryMethod.COMPUTE_CTL: + http_client.prewarm_lfc() + else: + pg_cur.execute("select prewarm_local_cache(%s)", (lfc_state,)) + + pg_cur.execute("select lfc_value from neon_lfc_stats where lfc_key='file_cache_used_pages'") + lfc_used_pages = pg_cur.fetchall()[0][0] log.info(f"Used LFC size: {lfc_used_pages}") - cur.execute("select * from get_prewarm_info()") - prewarm_info = cur.fetchall()[0] + pg_cur.execute("select * from get_prewarm_info()") + prewarm_info = pg_cur.fetchall()[0] log.info(f"Prewarm info: {prewarm_info}") - log.info(f"Prewarm progress: {(prewarm_info[1] + prewarm_info[2]) * 100 // prewarm_info[0]}%") + total, prewarmed, skipped, _ = prewarm_info + progress = (prewarmed + skipped) * 100 // total + log.info(f"Prewarm progress: {progress}%") assert lfc_used_pages > 10000 assert ( @@ -66,18 +115,23 @@ def test_lfc_prewarm(neon_simple_env: NeonEnv): and prewarm_info[0] == prewarm_info[1] + prewarm_info[2] ) - cur.execute("select sum(pk) from t") - assert cur.fetchall()[0][0] == n_records * (n_records + 1) / 2 + lfc_cur.execute("select sum(pk) from t") + assert lfc_cur.fetchall()[0][0] == n_records * (n_records + 1) / 2 - check_pinned_entries(cur) + check_pinned_entries(pg_cur) + + desired = {"status": "completed", "total": total, "prewarmed": prewarmed, "skipped": skipped} + if query is LfcQueryMethod.COMPUTE_CTL: + assert http_client.prewarm_lfc_status() == desired + assert prom_parse(http_client) == {OFFLOAD_LABEL: 0, PREWARM_LABEL: 1} @pytest.mark.skipif(not USE_LFC, reason="LFC is disabled, skipping") -def test_lfc_prewarm_under_workload(neon_simple_env: NeonEnv): +@pytest.mark.parametrize("query", QUERY_OPTIONS, ids=["postgres", "compute-ctl"]) +def test_lfc_prewarm_under_workload(neon_simple_env: NeonEnv, query: LfcQueryMethod): env = neon_simple_env n_records = 10000 n_threads = 4 - endpoint = env.endpoints.create_start( branch_name="main", config_lines=[ @@ -87,40 +141,58 @@ def test_lfc_prewarm_under_workload(neon_simple_env: NeonEnv): "neon.file_cache_prewarm_limit=1000000", ], ) - conn = endpoint.connect() - cur = conn.cursor() - cur.execute("create extension neon version '1.6'") - cur.execute( + + pg_conn = endpoint.connect() + pg_cur = pg_conn.cursor() + pg_cur.execute("create extension neon version '1.6'") + pg_cur.execute("CREATE DATABASE lfc") + + lfc_conn = endpoint.connect(dbname="lfc") + lfc_cur = lfc_conn.cursor() + lfc_cur.execute( "create table accounts(id integer primary key, balance bigint default 0, payload text default repeat('?', 1000)) with (fillfactor=10)" ) - cur.execute(f"insert into accounts(id) values (generate_series(1,{n_records}))") - cur.execute("select get_local_cache_state()") - lfc_state = cur.fetchall()[0][0] + log.info(f"Inserting {n_records} rows") + lfc_cur.execute(f"insert into accounts(id) values (generate_series(1,{n_records}))") + log.info(f"Inserted {n_records} rows") + + http_client = endpoint.http_client() + if query is LfcQueryMethod.COMPUTE_CTL: + http_client.offload_lfc() + else: + pg_cur.execute("select get_local_cache_state()") + lfc_state = pg_cur.fetchall()[0][0] running = True + n_prewarms = 0 def workload(): - conn = endpoint.connect() - cur = conn.cursor() + lfc_conn = endpoint.connect(dbname="lfc") + lfc_cur = lfc_conn.cursor() n_transfers = 0 while running: src = random.randint(1, n_records) dst = random.randint(1, n_records) - cur.execute("update accounts set balance=balance-100 where id=%s", (src,)) - cur.execute("update accounts set balance=balance+100 where id=%s", (dst,)) + lfc_cur.execute("update accounts set balance=balance-100 where id=%s", (src,)) + lfc_cur.execute("update accounts set balance=balance+100 where id=%s", (dst,)) n_transfers += 1 log.info(f"Number of transfers: {n_transfers}") def prewarm(): - conn = endpoint.connect() - cur = conn.cursor() - n_prewarms = 0 + pg_conn = endpoint.connect() + pg_cur = pg_conn.cursor() while running: - cur.execute("alter system set neon.file_cache_size_limit='1MB'") - cur.execute("select pg_reload_conf()") - cur.execute("alter system set neon.file_cache_size_limit='1GB'") - cur.execute("select pg_reload_conf()") - cur.execute("select prewarm_local_cache(%s)", (lfc_state,)) + pg_cur.execute("alter system set neon.file_cache_size_limit='1MB'") + pg_cur.execute("select pg_reload_conf()") + pg_cur.execute("alter system set neon.file_cache_size_limit='1GB'") + pg_cur.execute("select pg_reload_conf()") + + if query is LfcQueryMethod.COMPUTE_CTL: + http_client.prewarm_lfc() + else: + pg_cur.execute("select prewarm_local_cache(%s)", (lfc_state,)) + + nonlocal n_prewarms n_prewarms += 1 log.info(f"Number of prewarms: {n_prewarms}") @@ -140,8 +212,10 @@ def test_lfc_prewarm_under_workload(neon_simple_env: NeonEnv): t.join() prewarm_thread.join() - cur.execute("select sum(balance) from accounts") - total_balance = cur.fetchall()[0][0] + lfc_cur.execute("select sum(balance) from accounts") + total_balance = lfc_cur.fetchall()[0][0] assert total_balance == 0 - check_pinned_entries(cur) + check_pinned_entries(pg_cur) + if query is LfcQueryMethod.COMPUTE_CTL: + assert prom_parse(http_client) == {OFFLOAD_LABEL: 1, PREWARM_LABEL: n_prewarms}