diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 880d6044f2..2b88f09b3d 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -1131,7 +1131,7 @@ jobs: # TODO: move deployPreprodRegion to release (`"$GITHUB_REF_NAME" == "release"` block), once Staging support different compute tag prefixes for different regions gh workflow --repo neondatabase/aws run deploy-dev.yml --ref main -f branch=main -f dockerTag=${{needs.tag.outputs.build-tag}} -f deployPreprodRegion=true elif [[ "$GITHUB_REF_NAME" == "release" ]]; then - gh workflow --repo neondatabase/aws run deploy-prod.yml --ref main -f branch=main -f dockerTag=${{needs.tag.outputs.build-tag}} -f disclamerAcknowledged=true + gh workflow --repo neondatabase/aws run deploy-prod.yml --ref main -f branch=main -f dockerTag=${{needs.tag.outputs.build-tag}} else echo "GITHUB_REF_NAME (value '$GITHUB_REF_NAME') is not set to either 'main' or 'release'" exit 1 diff --git a/Cargo.lock b/Cargo.lock index c5c2523e29..0053c3b08c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -270,6 +270,32 @@ dependencies = [ "critical-section", ] +[[package]] +name = "attachment_service" +version = "0.1.0" +dependencies = [ + "anyhow", + "camino", + "clap", + "control_plane", + "futures", + "git-version", + "hyper", + "metrics", + "pageserver_api", + "pageserver_client", + "postgres_backend", + "postgres_connection", + "serde", + "serde_json", + "thiserror", + "tokio", + "tokio-util", + "tracing", + "utils", + "workspace_hack", +] + [[package]] name = "autocfg" version = "1.1.0" @@ -1748,6 +1774,12 @@ dependencies = [ "termcolor", ] +[[package]] +name = "equivalent" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" + [[package]] name = "errno" version = "0.3.1" @@ -2106,9 +2138,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.3.19" +version = "0.3.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d357c7ae988e7d2182f7d7871d0b963962420b0678b0997ce7de72001aeab782" +checksum = "bb2c4422095b67ee78da96fbb51a4cc413b3b25883c7717ff7ca1ab31022c9c9" dependencies = [ "bytes", "fnv", @@ -2116,7 +2148,7 @@ dependencies = [ "futures-sink", "futures-util", "http", - "indexmap", + "indexmap 2.0.1", "slab", "tokio", "tokio-util", @@ -2452,6 +2484,16 @@ dependencies = [ "serde", ] +[[package]] +name = "indexmap" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad227c3af19d4914570ad36d30409928b75967c298feb9ea1969db3a610bb14e" +dependencies = [ + "equivalent", + "hashbrown 0.14.0", +] + [[package]] name = "infer" version = "0.2.3" @@ -3130,7 +3172,7 @@ dependencies = [ "fnv", "futures-channel", "futures-util", - "indexmap", + "indexmap 1.9.3", "once_cell", "pin-project-lite", "thiserror", @@ -3340,6 +3382,7 @@ dependencies = [ "const_format", "enum-map", "hex", + "humantime-serde", "postgres_ffi", "rand 0.8.5", "serde", @@ -3524,7 +3567,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4dd7d28ee937e54fe3080c91faa1c3a46c06de6252988a7f4592ba2310ef22a4" dependencies = [ "fixedbitset", - "indexmap", + "indexmap 1.9.3", ] [[package]] @@ -4929,7 +4972,7 @@ dependencies = [ "base64 0.13.1", "chrono", "hex", - "indexmap", + "indexmap 1.9.3", "serde", "serde_json", "serde_with_macros", @@ -5630,7 +5673,7 @@ version = "0.19.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2380d56e8670370eee6566b0bfd4265f65b3f432e8c6d85623f728d4fa31f739" dependencies = [ - "indexmap", + "indexmap 1.9.3", "serde", "serde_spanned", "toml_datetime", @@ -5722,7 +5765,7 @@ checksum = "b8fa9be0de6cf49e536ce1851f987bd21a43b771b09473c3549a6c853db37c1c" dependencies = [ "futures-core", "futures-util", - "indexmap", + "indexmap 1.9.3", "pin-project", "pin-project-lite", "rand 0.8.5", @@ -6593,9 +6636,11 @@ dependencies = [ "futures-sink", "futures-util", "getrandom 0.2.11", + "hashbrown 0.14.0", "hex", "hmac", "hyper", + "indexmap 1.9.3", "itertools", "libc", "log", diff --git a/Cargo.toml b/Cargo.toml index 2d8fbaffa8..5d5d2f4a55 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,7 @@ resolver = "2" members = [ "compute_tools", "control_plane", + "control_plane/attachment_service", "pageserver", "pageserver/ctl", "pageserver/client", diff --git a/compute_tools/src/monitor.rs b/compute_tools/src/monitor.rs index 1a28aab479..f09bd02664 100644 --- a/compute_tools/src/monitor.rs +++ b/compute_tools/src/monitor.rs @@ -137,6 +137,28 @@ fn watch_compute_activity(compute: &ComputeNode) { continue; } } + // + // Do not suspend compute if autovacuum is running + // + let autovacuum_count_query = "select count(*) from pg_stat_activity where backend_type = 'autovacuum worker'"; + match cli.query_one(autovacuum_count_query, &[]) { + Ok(r) => match r.try_get::<&str, i64>("count") { + Ok(num_workers) => { + if num_workers > 0 { + compute.update_last_active(Some(Utc::now())); + continue; + } + } + Err(e) => { + warn!("failed to parse autovacuum workers count: {:?}", e); + continue; + } + }, + Err(e) => { + warn!("failed to get list of autovacuum workers: {:?}", e); + continue; + } + } } Err(e) => { debug!("could not connect to Postgres: {}, retrying", e); diff --git a/control_plane/attachment_service/Cargo.toml b/control_plane/attachment_service/Cargo.toml new file mode 100644 index 0000000000..2e2286dbab --- /dev/null +++ b/control_plane/attachment_service/Cargo.toml @@ -0,0 +1,32 @@ +[package] +name = "attachment_service" +version = "0.1.0" +edition.workspace = true +license.workspace = true + +[dependencies] +anyhow.workspace = true +camino.workspace = true +clap.workspace = true +futures.workspace = true +git-version.workspace = true +hyper.workspace = true +pageserver_api.workspace = true +pageserver_client.workspace = true +postgres_connection.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokio.workspace = true +tokio-util.workspace = true +tracing.workspace = true + +# TODO: remove this after DB persistence is added, it is only used for +# a parsing function when loading pageservers from neon_local LocalEnv +postgres_backend.workspace = true + +utils = { path = "../../libs/utils/" } +metrics = { path = "../../libs/metrics/" } +control_plane = { path = ".." } +workspace_hack = { version = "0.1", path = "../../workspace_hack" } + diff --git a/control_plane/attachment_service/src/compute_hook.rs b/control_plane/attachment_service/src/compute_hook.rs new file mode 100644 index 0000000000..02617cd065 --- /dev/null +++ b/control_plane/attachment_service/src/compute_hook.rs @@ -0,0 +1,116 @@ +use std::collections::HashMap; + +use control_plane::endpoint::ComputeControlPlane; +use control_plane::local_env::LocalEnv; +use pageserver_api::shard::{ShardCount, ShardIndex, TenantShardId}; +use postgres_connection::parse_host_port; +use utils::id::{NodeId, TenantId}; + +pub(super) struct ComputeHookTenant { + shards: Vec<(ShardIndex, NodeId)>, +} + +impl ComputeHookTenant { + pub(super) async fn maybe_reconfigure(&mut self, tenant_id: TenantId) -> anyhow::Result<()> { + // Find the highest shard count and drop any shards that aren't + // for that shard count. + let shard_count = self.shards.iter().map(|(k, _v)| k.shard_count).max(); + let Some(shard_count) = shard_count else { + // No shards, nothing to do. + tracing::info!("ComputeHookTenant::maybe_reconfigure: no shards"); + return Ok(()); + }; + + self.shards.retain(|(k, _v)| k.shard_count == shard_count); + self.shards + .sort_by_key(|(shard, _node_id)| shard.shard_number); + + if self.shards.len() == shard_count.0 as usize || shard_count == ShardCount(0) { + // We have pageservers for all the shards: proceed to reconfigure compute + let env = match LocalEnv::load_config() { + Ok(e) => e, + Err(e) => { + tracing::warn!( + "Couldn't load neon_local config, skipping compute update ({e})" + ); + return Ok(()); + } + }; + let cplane = ComputeControlPlane::load(env.clone()) + .expect("Error loading compute control plane"); + + let compute_pageservers = self + .shards + .iter() + .map(|(_shard, node_id)| { + let ps_conf = env + .get_pageserver_conf(*node_id) + .expect("Unknown pageserver"); + let (pg_host, pg_port) = parse_host_port(&ps_conf.listen_pg_addr) + .expect("Unable to parse listen_pg_addr"); + (pg_host, pg_port.unwrap_or(5432)) + }) + .collect::>(); + + for (endpoint_name, endpoint) in &cplane.endpoints { + if endpoint.tenant_id == tenant_id && endpoint.status() == "running" { + tracing::info!("🔁 Reconfiguring endpoint {}", endpoint_name,); + endpoint.reconfigure(compute_pageservers.clone()).await?; + } + } + } else { + tracing::info!( + "ComputeHookTenant::maybe_reconfigure: not enough shards ({}/{})", + self.shards.len(), + shard_count.0 + ); + } + + Ok(()) + } +} + +/// The compute hook is a destination for notifications about changes to tenant:pageserver +/// mapping. It aggregates updates for the shards in a tenant, and when appropriate reconfigures +/// the compute connection string. +pub(super) struct ComputeHook { + state: tokio::sync::Mutex>, +} + +impl ComputeHook { + pub(super) fn new() -> Self { + Self { + state: Default::default(), + } + } + + pub(super) async fn notify( + &self, + tenant_shard_id: TenantShardId, + node_id: NodeId, + ) -> anyhow::Result<()> { + tracing::info!("ComputeHook::notify: {}->{}", tenant_shard_id, node_id); + let mut locked = self.state.lock().await; + let entry = locked + .entry(tenant_shard_id.tenant_id) + .or_insert_with(|| ComputeHookTenant { shards: Vec::new() }); + + let shard_index = ShardIndex { + shard_count: tenant_shard_id.shard_count, + shard_number: tenant_shard_id.shard_number, + }; + + let mut set = false; + for (existing_shard, existing_node) in &mut entry.shards { + if *existing_shard == shard_index { + *existing_node = node_id; + set = true; + } + } + if !set { + entry.shards.push((shard_index, node_id)); + } + + entry.maybe_reconfigure(tenant_shard_id.tenant_id).await + } +} diff --git a/control_plane/attachment_service/src/http.rs b/control_plane/attachment_service/src/http.rs new file mode 100644 index 0000000000..30f6dd66ee --- /dev/null +++ b/control_plane/attachment_service/src/http.rs @@ -0,0 +1,218 @@ +use crate::reconciler::ReconcileError; +use crate::service::Service; +use hyper::{Body, Request, Response}; +use hyper::{StatusCode, Uri}; +use pageserver_api::models::{TenantCreateRequest, TimelineCreateRequest}; +use pageserver_api::shard::TenantShardId; +use std::sync::Arc; +use utils::auth::SwappableJwtAuth; +use utils::http::endpoint::{auth_middleware, request_span}; +use utils::http::request::parse_request_param; +use utils::id::TenantId; + +use utils::{ + http::{ + endpoint::{self}, + error::ApiError, + json::{json_request, json_response}, + RequestExt, RouterBuilder, + }, + id::NodeId, +}; + +use pageserver_api::control_api::{ReAttachRequest, ValidateRequest}; + +use control_plane::attachment_service::{ + AttachHookRequest, InspectRequest, NodeConfigureRequest, NodeRegisterRequest, + TenantShardMigrateRequest, +}; + +/// State available to HTTP request handlers +#[derive(Clone)] +pub struct HttpState { + service: Arc, + auth: Option>, + allowlist_routes: Vec, +} + +impl HttpState { + pub fn new(service: Arc, auth: Option>) -> Self { + let allowlist_routes = ["/status"] + .iter() + .map(|v| v.parse().unwrap()) + .collect::>(); + Self { + service, + auth, + allowlist_routes, + } + } +} + +#[inline(always)] +fn get_state(request: &Request) -> &HttpState { + request + .data::>() + .expect("unknown state type") + .as_ref() +} + +/// Pageserver calls into this on startup, to learn which tenants it should attach +async fn handle_re_attach(mut req: Request) -> Result, ApiError> { + let reattach_req = json_request::(&mut req).await?; + let state = get_state(&req); + json_response( + StatusCode::OK, + state + .service + .re_attach(reattach_req) + .await + .map_err(ApiError::InternalServerError)?, + ) +} + +/// Pageserver calls into this before doing deletions, to confirm that it still +/// holds the latest generation for the tenants with deletions enqueued +async fn handle_validate(mut req: Request) -> Result, ApiError> { + let validate_req = json_request::(&mut req).await?; + let state = get_state(&req); + json_response(StatusCode::OK, state.service.validate(validate_req)) +} + +/// Call into this before attaching a tenant to a pageserver, to acquire a generation number +/// (in the real control plane this is unnecessary, because the same program is managing +/// generation numbers and doing attachments). +async fn handle_attach_hook(mut req: Request) -> Result, ApiError> { + let attach_req = json_request::(&mut req).await?; + let state = get_state(&req); + + json_response( + StatusCode::OK, + state + .service + .attach_hook(attach_req) + .await + .map_err(ApiError::InternalServerError)?, + ) +} + +async fn handle_inspect(mut req: Request) -> Result, ApiError> { + let inspect_req = json_request::(&mut req).await?; + + let state = get_state(&req); + + json_response(StatusCode::OK, state.service.inspect(inspect_req)) +} + +async fn handle_tenant_create(mut req: Request) -> Result, ApiError> { + let create_req = json_request::(&mut req).await?; + let state = get_state(&req); + json_response( + StatusCode::OK, + state.service.tenant_create(create_req).await?, + ) +} + +async fn handle_tenant_timeline_create(mut req: Request) -> Result, ApiError> { + let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?; + let create_req = json_request::(&mut req).await?; + + let state = get_state(&req); + json_response( + StatusCode::OK, + state + .service + .tenant_timeline_create(tenant_id, create_req) + .await?, + ) +} + +async fn handle_tenant_locate(req: Request) -> Result, ApiError> { + let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?; + let state = get_state(&req); + + json_response(StatusCode::OK, state.service.tenant_locate(tenant_id)?) +} + +async fn handle_node_register(mut req: Request) -> Result, ApiError> { + let register_req = json_request::(&mut req).await?; + let state = get_state(&req); + state.service.node_register(register_req).await?; + json_response(StatusCode::OK, ()) +} + +async fn handle_node_configure(mut req: Request) -> Result, ApiError> { + let node_id: NodeId = parse_request_param(&req, "node_id")?; + let config_req = json_request::(&mut req).await?; + if node_id != config_req.node_id { + return Err(ApiError::BadRequest(anyhow::anyhow!( + "Path and body node_id differ" + ))); + } + let state = get_state(&req); + + json_response(StatusCode::OK, state.service.node_configure(config_req)?) +} + +async fn handle_tenant_shard_migrate(mut req: Request) -> Result, ApiError> { + let tenant_shard_id: TenantShardId = parse_request_param(&req, "tenant_shard_id")?; + let migrate_req = json_request::(&mut req).await?; + let state = get_state(&req); + json_response( + StatusCode::OK, + state + .service + .tenant_shard_migrate(tenant_shard_id, migrate_req) + .await?, + ) +} + +/// Status endpoint is just used for checking that our HTTP listener is up +async fn handle_status(_req: Request) -> Result, ApiError> { + json_response(StatusCode::OK, ()) +} + +impl From for ApiError { + fn from(value: ReconcileError) -> Self { + ApiError::Conflict(format!("Reconciliation error: {}", value)) + } +} + +pub fn make_router( + service: Arc, + auth: Option>, +) -> RouterBuilder { + let mut router = endpoint::make_router(); + if auth.is_some() { + router = router.middleware(auth_middleware(|request| { + let state = get_state(request); + if state.allowlist_routes.contains(request.uri()) { + None + } else { + state.auth.as_deref() + } + })) + } + + router + .data(Arc::new(HttpState::new(service, auth))) + .get("/status", |r| request_span(r, handle_status)) + .post("/re-attach", |r| request_span(r, handle_re_attach)) + .post("/validate", |r| request_span(r, handle_validate)) + .post("/attach-hook", |r| request_span(r, handle_attach_hook)) + .post("/inspect", |r| request_span(r, handle_inspect)) + .post("/node", |r| request_span(r, handle_node_register)) + .put("/node/:node_id/config", |r| { + request_span(r, handle_node_configure) + }) + .post("/tenant", |r| request_span(r, handle_tenant_create)) + .post("/tenant/:tenant_id/timeline", |r| { + request_span(r, handle_tenant_timeline_create) + }) + .get("/tenant/:tenant_id/locate", |r| { + request_span(r, handle_tenant_locate) + }) + .put("/tenant/:tenant_shard_id/migrate", |r| { + request_span(r, handle_tenant_shard_migrate) + }) +} diff --git a/control_plane/attachment_service/src/lib.rs b/control_plane/attachment_service/src/lib.rs new file mode 100644 index 0000000000..d8f996952a --- /dev/null +++ b/control_plane/attachment_service/src/lib.rs @@ -0,0 +1,57 @@ +use serde::{Deserialize, Serialize}; +use utils::seqwait::MonotonicCounter; + +mod compute_hook; +pub mod http; +mod node; +pub mod persistence; +mod reconciler; +mod scheduler; +pub mod service; +mod tenant_state; + +#[derive(Clone, Serialize, Deserialize)] +enum PlacementPolicy { + /// Cheapest way to attach a tenant: just one pageserver, no secondary + Single, + /// Production-ready way to attach a tenant: one attached pageserver and + /// some number of secondaries. + Double(usize), +} + +#[derive(Ord, PartialOrd, Eq, PartialEq, Copy, Clone)] +struct Sequence(u64); + +impl Sequence { + fn initial() -> Self { + Self(0) + } +} + +impl std::fmt::Display for Sequence { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +impl MonotonicCounter for Sequence { + fn cnt_advance(&mut self, v: Sequence) { + assert!(*self <= v); + *self = v; + } + fn cnt_value(&self) -> Sequence { + *self + } +} + +impl Sequence { + fn next(&self) -> Sequence { + Sequence(self.0 + 1) + } +} + +impl Default for PlacementPolicy { + fn default() -> Self { + PlacementPolicy::Double(1) + } +} diff --git a/control_plane/attachment_service/src/main.rs b/control_plane/attachment_service/src/main.rs new file mode 100644 index 0000000000..ee2a22ee53 --- /dev/null +++ b/control_plane/attachment_service/src/main.rs @@ -0,0 +1,100 @@ +/// The attachment service mimics the aspects of the control plane API +/// that are required for a pageserver to operate. +/// +/// This enables running & testing pageservers without a full-blown +/// deployment of the Neon cloud platform. +/// +use anyhow::anyhow; +use attachment_service::http::make_router; +use attachment_service::persistence::Persistence; +use attachment_service::service::{Config, Service}; +use camino::Utf8PathBuf; +use clap::Parser; +use metrics::launch_timestamp::LaunchTimestamp; +use std::sync::Arc; +use utils::auth::{JwtAuth, SwappableJwtAuth}; +use utils::logging::{self, LogFormat}; +use utils::signals::{ShutdownSignals, Signal}; + +use utils::{project_build_tag, project_git_version, tcp_listener}; + +project_git_version!(GIT_VERSION); +project_build_tag!(BUILD_TAG); + +#[derive(Parser)] +#[command(author, version, about, long_about = None)] +#[command(arg_required_else_help(true))] +struct Cli { + /// Host and port to listen on, like `127.0.0.1:1234` + #[arg(short, long)] + listen: std::net::SocketAddr, + + /// Path to public key for JWT authentication of clients + #[arg(long)] + public_key: Option, + + /// Token for authenticating this service with the pageservers it controls + #[arg(short, long)] + jwt_token: Option, + + /// Path to the .json file to store state (will be created if it doesn't exist) + #[arg(short, long)] + path: Utf8PathBuf, +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let launch_ts = Box::leak(Box::new(LaunchTimestamp::generate())); + + logging::init( + LogFormat::Plain, + logging::TracingErrorLayerEnablement::Disabled, + logging::Output::Stdout, + )?; + + let args = Cli::parse(); + tracing::info!( + "version: {}, launch_timestamp: {}, build_tag {}, state at {}, listening on {}", + GIT_VERSION, + launch_ts.to_string(), + BUILD_TAG, + args.path, + args.listen + ); + + let config = Config { + jwt_token: args.jwt_token, + }; + + let persistence = Arc::new(Persistence::new(&args.path).await); + + let service = Service::spawn(config, persistence).await?; + + let http_listener = tcp_listener::bind(args.listen)?; + + let auth = if let Some(public_key_path) = &args.public_key { + let jwt_auth = JwtAuth::from_key_path(public_key_path)?; + Some(Arc::new(SwappableJwtAuth::new(jwt_auth))) + } else { + None + }; + let router = make_router(service, auth) + .build() + .map_err(|err| anyhow!(err))?; + let service = utils::http::RouterService::new(router).unwrap(); + let server = hyper::Server::from_tcp(http_listener)?.serve(service); + + tracing::info!("Serving on {0}", args.listen); + + tokio::task::spawn(server); + + ShutdownSignals::handle(|signal| match signal { + Signal::Interrupt | Signal::Terminate | Signal::Quit => { + tracing::info!("Got {}. Terminating", signal.name()); + // We're just a test helper: no graceful shutdown. + std::process::exit(0); + } + })?; + + Ok(()) +} diff --git a/control_plane/attachment_service/src/node.rs b/control_plane/attachment_service/src/node.rs new file mode 100644 index 0000000000..efd3f8f49b --- /dev/null +++ b/control_plane/attachment_service/src/node.rs @@ -0,0 +1,37 @@ +use control_plane::attachment_service::{NodeAvailability, NodeSchedulingPolicy}; +use utils::id::NodeId; + +#[derive(Clone)] +pub(crate) struct Node { + pub(crate) id: NodeId, + + pub(crate) availability: NodeAvailability, + pub(crate) scheduling: NodeSchedulingPolicy, + + pub(crate) listen_http_addr: String, + pub(crate) listen_http_port: u16, + + pub(crate) listen_pg_addr: String, + pub(crate) listen_pg_port: u16, +} + +impl Node { + pub(crate) fn base_url(&self) -> String { + format!("http://{}:{}", self.listen_http_addr, self.listen_http_port) + } + + /// Is this node elegible to have work scheduled onto it? + pub(crate) fn may_schedule(&self) -> bool { + match self.availability { + NodeAvailability::Active => {} + NodeAvailability::Offline => return false, + } + + match self.scheduling { + NodeSchedulingPolicy::Active => true, + NodeSchedulingPolicy::Draining => false, + NodeSchedulingPolicy::Filling => true, + NodeSchedulingPolicy::Pause => false, + } + } +} diff --git a/control_plane/attachment_service/src/persistence.rs b/control_plane/attachment_service/src/persistence.rs new file mode 100644 index 0000000000..58708be140 --- /dev/null +++ b/control_plane/attachment_service/src/persistence.rs @@ -0,0 +1,272 @@ +use std::{collections::HashMap, str::FromStr}; + +use camino::{Utf8Path, Utf8PathBuf}; +use control_plane::{ + attachment_service::{NodeAvailability, NodeSchedulingPolicy}, + local_env::LocalEnv, +}; +use pageserver_api::{ + models::TenantConfig, + shard::{ShardCount, ShardNumber, TenantShardId}, +}; +use postgres_connection::parse_host_port; +use serde::{Deserialize, Serialize}; +use utils::{ + generation::Generation, + id::{NodeId, TenantId}, +}; + +use crate::{node::Node, PlacementPolicy}; + +/// Placeholder for storage. This will be replaced with a database client. +pub struct Persistence { + state: std::sync::Mutex, +} + +// Top level state available to all HTTP handlers +#[derive(Serialize, Deserialize)] +struct PersistentState { + tenants: HashMap, + + #[serde(skip)] + path: Utf8PathBuf, +} + +/// A convenience for serializing the state inside a sync lock, and then +/// writing it to disk outside of the lock. This will go away when switching +/// to a database backend. +struct PendingWrite { + bytes: Vec, + path: Utf8PathBuf, +} + +impl PendingWrite { + async fn commit(&self) -> anyhow::Result<()> { + tokio::fs::write(&self.path, &self.bytes).await?; + + Ok(()) + } +} + +impl PersistentState { + fn save(&self) -> PendingWrite { + PendingWrite { + bytes: serde_json::to_vec(self).expect("Serialization error"), + path: self.path.clone(), + } + } + + async fn load(path: &Utf8Path) -> anyhow::Result { + let bytes = tokio::fs::read(path).await?; + let mut decoded = serde_json::from_slice::(&bytes)?; + decoded.path = path.to_owned(); + + for (tenant_id, tenant) in &mut decoded.tenants { + // Backward compat: an old attachments.json from before PR #6251, replace + // empty strings with proper defaults. + if tenant.tenant_id.is_empty() { + tenant.tenant_id = format!("{}", tenant_id); + tenant.config = serde_json::to_string(&TenantConfig::default())?; + tenant.placement_policy = serde_json::to_string(&PlacementPolicy::default())?; + } + } + + Ok(decoded) + } + + async fn load_or_new(path: &Utf8Path) -> Self { + match Self::load(path).await { + Ok(s) => { + tracing::info!("Loaded state file at {}", path); + s + } + Err(e) + if e.downcast_ref::() + .map(|e| e.kind() == std::io::ErrorKind::NotFound) + .unwrap_or(false) => + { + tracing::info!("Will create state file at {}", path); + Self { + tenants: HashMap::new(), + path: path.to_owned(), + } + } + Err(e) => { + panic!("Failed to load state from '{}': {e:#} (maybe your .neon/ dir was written by an older version?)", path) + } + } + } +} + +impl Persistence { + pub async fn new(path: &Utf8Path) -> Self { + let state = PersistentState::load_or_new(path).await; + Self { + state: std::sync::Mutex::new(state), + } + } + + /// When registering a node, persist it so that on next start we will be able to + /// iterate over known nodes to synchronize their tenant shard states with our observed state. + pub(crate) async fn insert_node(&self, _node: &Node) -> anyhow::Result<()> { + // TODO: node persitence will come with database backend + Ok(()) + } + + /// At startup, we populate the service's list of nodes, and use this list to call into + /// each node to do an initial reconciliation of the state of the world with our in-memory + /// observed state. + pub(crate) async fn list_nodes(&self) -> anyhow::Result> { + let env = LocalEnv::load_config()?; + // TODO: node persitence will come with database backend + + // XXX hack: enable test_backward_compatibility to work by populating our list of + // nodes from LocalEnv when it is not present in persistent storage. Otherwise at + // first startup in the compat test, we may have shards but no nodes. + let mut result = Vec::new(); + tracing::info!( + "Loaded {} pageserver nodes from LocalEnv", + env.pageservers.len() + ); + for ps_conf in env.pageservers { + let (pg_host, pg_port) = + parse_host_port(&ps_conf.listen_pg_addr).expect("Unable to parse listen_pg_addr"); + let (http_host, http_port) = parse_host_port(&ps_conf.listen_http_addr) + .expect("Unable to parse listen_http_addr"); + result.push(Node { + id: ps_conf.id, + listen_pg_addr: pg_host.to_string(), + listen_pg_port: pg_port.unwrap_or(5432), + listen_http_addr: http_host.to_string(), + listen_http_port: http_port.unwrap_or(80), + availability: NodeAvailability::Active, + scheduling: NodeSchedulingPolicy::Active, + }); + } + + Ok(result) + } + + /// At startup, we populate our map of tenant shards from persistent storage. + pub(crate) async fn list_tenant_shards(&self) -> anyhow::Result> { + let locked = self.state.lock().unwrap(); + Ok(locked.tenants.values().cloned().collect()) + } + + /// Tenants must be persisted before we schedule them for the first time. This enables us + /// to correctly retain generation monotonicity, and the externally provided placement policy & config. + pub(crate) async fn insert_tenant_shards( + &self, + shards: Vec, + ) -> anyhow::Result<()> { + let write = { + let mut locked = self.state.lock().unwrap(); + for shard in shards { + let tenant_shard_id = TenantShardId { + tenant_id: TenantId::from_str(shard.tenant_id.as_str())?, + shard_number: ShardNumber(shard.shard_number as u8), + shard_count: ShardCount(shard.shard_count as u8), + }; + + locked.tenants.insert(tenant_shard_id, shard); + } + locked.save() + }; + + write.commit().await?; + + Ok(()) + } + + /// Reconciler calls this immediately before attaching to a new pageserver, to acquire a unique, monotonically + /// advancing generation number. We also store the NodeId for which the generation was issued, so that in + /// [`Self::re_attach`] we can do a bulk UPDATE on the generations for that node. + pub(crate) async fn increment_generation( + &self, + tenant_shard_id: TenantShardId, + node_id: Option, + ) -> anyhow::Result { + let (write, gen) = { + let mut locked = self.state.lock().unwrap(); + let Some(shard) = locked.tenants.get_mut(&tenant_shard_id) else { + anyhow::bail!("Tried to increment generation of unknown shard"); + }; + + // If we're called with a None pageserver, we need only update the generation + // record to disassociate it with this pageserver, not actually increment the number, as + // the increment is guaranteed to happen the next time this tenant is attached. + if node_id.is_some() { + shard.generation += 1; + } + + shard.generation_pageserver = node_id; + let gen = Generation::new(shard.generation); + (locked.save(), gen) + }; + + write.commit().await?; + Ok(gen) + } + + pub(crate) async fn re_attach( + &self, + node_id: NodeId, + ) -> anyhow::Result> { + let (write, result) = { + let mut result = HashMap::new(); + let mut locked = self.state.lock().unwrap(); + for (tenant_shard_id, shard) in locked.tenants.iter_mut() { + if shard.generation_pageserver == Some(node_id) { + shard.generation += 1; + result.insert(*tenant_shard_id, Generation::new(shard.generation)); + } + } + + (locked.save(), result) + }; + + write.commit().await?; + Ok(result) + } + + // TODO: when we start shard splitting, we must durably mark the tenant so that + // on restart, we know that we must go through recovery (list shards that exist + // and pick up where we left off and/or revert to parent shards). + #[allow(dead_code)] + pub(crate) async fn begin_shard_split(&self, _tenant_id: TenantId) -> anyhow::Result<()> { + todo!(); + } + + // TODO: when we finish shard splitting, we must atomically clean up the old shards + // and insert the new shards, and clear the splitting marker. + #[allow(dead_code)] + pub(crate) async fn complete_shard_split(&self, _tenant_id: TenantId) -> anyhow::Result<()> { + todo!(); + } +} + +/// Parts of [`crate::tenant_state::TenantState`] that are stored durably +#[derive(Serialize, Deserialize, Clone)] +pub(crate) struct TenantShardPersistence { + #[serde(default)] + pub(crate) tenant_id: String, + #[serde(default)] + pub(crate) shard_number: i32, + #[serde(default)] + pub(crate) shard_count: i32, + #[serde(default)] + pub(crate) shard_stripe_size: i32, + + // Currently attached pageserver + #[serde(rename = "pageserver")] + pub(crate) generation_pageserver: Option, + + // Latest generation number: next time we attach, increment this + // and use the incremented number when attaching + pub(crate) generation: u32, + + #[serde(default)] + pub(crate) placement_policy: String, + #[serde(default)] + pub(crate) config: String, +} diff --git a/control_plane/attachment_service/src/reconciler.rs b/control_plane/attachment_service/src/reconciler.rs new file mode 100644 index 0000000000..b08339b3b4 --- /dev/null +++ b/control_plane/attachment_service/src/reconciler.rs @@ -0,0 +1,495 @@ +use crate::persistence::Persistence; +use crate::service; +use control_plane::attachment_service::NodeAvailability; +use pageserver_api::models::{ + LocationConfig, LocationConfigMode, LocationConfigSecondary, TenantConfig, +}; +use pageserver_api::shard::{ShardIdentity, TenantShardId}; +use pageserver_client::mgmt_api; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; +use tokio_util::sync::CancellationToken; +use utils::generation::Generation; +use utils::id::{NodeId, TimelineId}; +use utils::lsn::Lsn; + +use crate::compute_hook::ComputeHook; +use crate::node::Node; +use crate::tenant_state::{IntentState, ObservedState, ObservedStateLocation}; + +/// Object with the lifetime of the background reconcile task that is created +/// for tenants which have a difference between their intent and observed states. +pub(super) struct Reconciler { + /// See [`crate::tenant_state::TenantState`] for the meanings of these fields: they are a snapshot + /// of a tenant's state from when we spawned a reconcile task. + pub(super) tenant_shard_id: TenantShardId, + pub(crate) shard: ShardIdentity, + pub(crate) generation: Generation, + pub(crate) intent: IntentState, + pub(crate) config: TenantConfig, + pub(crate) observed: ObservedState, + + pub(crate) service_config: service::Config, + + /// A snapshot of the pageservers as they were when we were asked + /// to reconcile. + pub(crate) pageservers: Arc>, + + /// A hook to notify the running postgres instances when we change the location + /// of a tenant + pub(crate) compute_hook: Arc, + + /// A means to abort background reconciliation: it is essential to + /// call this when something changes in the original TenantState that + /// will make this reconciliation impossible or unnecessary, for + /// example when a pageserver node goes offline, or the PlacementPolicy for + /// the tenant is changed. + pub(crate) cancel: CancellationToken, + + /// Access to persistent storage for updating generation numbers + pub(crate) persistence: Arc, +} + +#[derive(thiserror::Error, Debug)] +pub enum ReconcileError { + #[error(transparent)] + Other(#[from] anyhow::Error), +} + +impl Reconciler { + async fn location_config( + &mut self, + node_id: NodeId, + config: LocationConfig, + flush_ms: Option, + ) -> anyhow::Result<()> { + let node = self + .pageservers + .get(&node_id) + .expect("Pageserver may not be removed while referenced"); + + self.observed + .locations + .insert(node.id, ObservedStateLocation { conf: None }); + + tracing::info!("location_config({}) calling: {:?}", node_id, config); + let client = + mgmt_api::Client::new(node.base_url(), self.service_config.jwt_token.as_deref()); + client + .location_config(self.tenant_shard_id, config.clone(), flush_ms) + .await?; + tracing::info!("location_config({}) complete: {:?}", node_id, config); + + self.observed + .locations + .insert(node.id, ObservedStateLocation { conf: Some(config) }); + + Ok(()) + } + + async fn maybe_live_migrate(&mut self) -> Result<(), ReconcileError> { + let destination = if let Some(node_id) = self.intent.attached { + match self.observed.locations.get(&node_id) { + Some(conf) => { + // We will do a live migration only if the intended destination is not + // currently in an attached state. + match &conf.conf { + Some(conf) if conf.mode == LocationConfigMode::Secondary => { + // Fall through to do a live migration + node_id + } + None | Some(_) => { + // Attached or uncertain: don't do a live migration, proceed + // with a general-case reconciliation + tracing::info!("maybe_live_migrate: destination is None or attached"); + return Ok(()); + } + } + } + None => { + // Our destination is not attached: maybe live migrate if some other + // node is currently attached. Fall through. + node_id + } + } + } else { + // No intent to be attached + tracing::info!("maybe_live_migrate: no attached intent"); + return Ok(()); + }; + + let mut origin = None; + for (node_id, state) in &self.observed.locations { + if let Some(observed_conf) = &state.conf { + if observed_conf.mode == LocationConfigMode::AttachedSingle { + let node = self + .pageservers + .get(node_id) + .expect("Nodes may not be removed while referenced"); + // We will only attempt live migration if the origin is not offline: this + // avoids trying to do it while reconciling after responding to an HA failover. + if !matches!(node.availability, NodeAvailability::Offline) { + origin = Some(*node_id); + break; + } + } + } + } + + let Some(origin) = origin else { + tracing::info!("maybe_live_migrate: no origin found"); + return Ok(()); + }; + + // We have an origin and a destination: proceed to do the live migration + tracing::info!("Live migrating {}->{}", origin, destination); + self.live_migrate(origin, destination).await?; + + Ok(()) + } + + async fn get_lsns( + &self, + tenant_shard_id: TenantShardId, + node_id: &NodeId, + ) -> anyhow::Result> { + let node = self + .pageservers + .get(node_id) + .expect("Pageserver may not be removed while referenced"); + + let client = + mgmt_api::Client::new(node.base_url(), self.service_config.jwt_token.as_deref()); + + let timelines = client.timeline_list(&tenant_shard_id).await?; + Ok(timelines + .into_iter() + .map(|t| (t.timeline_id, t.last_record_lsn)) + .collect()) + } + + async fn secondary_download(&self, tenant_shard_id: TenantShardId, node_id: &NodeId) { + let node = self + .pageservers + .get(node_id) + .expect("Pageserver may not be removed while referenced"); + + let client = + mgmt_api::Client::new(node.base_url(), self.service_config.jwt_token.as_deref()); + + match client.tenant_secondary_download(tenant_shard_id).await { + Ok(()) => {} + Err(_) => { + tracing::info!(" (skipping, destination wasn't in secondary mode)") + } + } + } + + async fn await_lsn( + &self, + tenant_shard_id: TenantShardId, + pageserver_id: &NodeId, + baseline: HashMap, + ) -> anyhow::Result<()> { + loop { + let latest = match self.get_lsns(tenant_shard_id, pageserver_id).await { + Ok(l) => l, + Err(e) => { + println!( + "🕑 Can't get LSNs on pageserver {} yet, waiting ({e})", + pageserver_id + ); + std::thread::sleep(Duration::from_millis(500)); + continue; + } + }; + + let mut any_behind: bool = false; + for (timeline_id, baseline_lsn) in &baseline { + match latest.get(timeline_id) { + Some(latest_lsn) => { + println!("🕑 LSN origin {baseline_lsn} vs destination {latest_lsn}"); + if latest_lsn < baseline_lsn { + any_behind = true; + } + } + None => { + // Expected timeline isn't yet visible on migration destination. + // (IRL we would have to account for timeline deletion, but this + // is just test helper) + any_behind = true; + } + } + } + + if !any_behind { + println!("✅ LSN caught up. Proceeding..."); + break; + } else { + std::thread::sleep(Duration::from_millis(500)); + } + } + + Ok(()) + } + + pub async fn live_migrate( + &mut self, + origin_ps_id: NodeId, + dest_ps_id: NodeId, + ) -> anyhow::Result<()> { + // `maybe_live_migrate` is responsibble for sanity of inputs + assert!(origin_ps_id != dest_ps_id); + + fn build_location_config( + shard: &ShardIdentity, + config: &TenantConfig, + mode: LocationConfigMode, + generation: Option, + secondary_conf: Option, + ) -> LocationConfig { + LocationConfig { + mode, + generation: generation.map(|g| g.into().unwrap()), + secondary_conf, + tenant_conf: config.clone(), + shard_number: shard.number.0, + shard_count: shard.count.0, + shard_stripe_size: shard.stripe_size.0, + } + } + + tracing::info!( + "🔁 Switching origin pageserver {} to stale mode", + origin_ps_id + ); + + // FIXME: it is incorrect to use self.generation here, we should use the generation + // from the ObservedState of the origin pageserver (it might be older than self.generation) + let stale_conf = build_location_config( + &self.shard, + &self.config, + LocationConfigMode::AttachedStale, + Some(self.generation), + None, + ); + self.location_config(origin_ps_id, stale_conf, Some(Duration::from_secs(10))) + .await?; + + let baseline_lsns = Some(self.get_lsns(self.tenant_shard_id, &origin_ps_id).await?); + + // If we are migrating to a destination that has a secondary location, warm it up first + if let Some(destination_conf) = self.observed.locations.get(&dest_ps_id) { + if let Some(destination_conf) = &destination_conf.conf { + if destination_conf.mode == LocationConfigMode::Secondary { + tracing::info!( + "🔁 Downloading latest layers to destination pageserver {}", + dest_ps_id, + ); + self.secondary_download(self.tenant_shard_id, &dest_ps_id) + .await; + } + } + } + + // Increment generation before attaching to new pageserver + self.generation = self + .persistence + .increment_generation(self.tenant_shard_id, Some(dest_ps_id)) + .await?; + + let dest_conf = build_location_config( + &self.shard, + &self.config, + LocationConfigMode::AttachedMulti, + Some(self.generation), + None, + ); + + tracing::info!("🔁 Attaching to pageserver {}", dest_ps_id); + self.location_config(dest_ps_id, dest_conf, None).await?; + + if let Some(baseline) = baseline_lsns { + tracing::info!("🕑 Waiting for LSN to catch up..."); + self.await_lsn(self.tenant_shard_id, &dest_ps_id, baseline) + .await?; + } + + tracing::info!("🔁 Notifying compute to use pageserver {}", dest_ps_id); + self.compute_hook + .notify(self.tenant_shard_id, dest_ps_id) + .await?; + + // Downgrade the origin to secondary. If the tenant's policy is PlacementPolicy::Single, then + // this location will be deleted in the general case reconciliation that runs after this. + let origin_secondary_conf = build_location_config( + &self.shard, + &self.config, + LocationConfigMode::Secondary, + None, + Some(LocationConfigSecondary { warm: true }), + ); + self.location_config(origin_ps_id, origin_secondary_conf.clone(), None) + .await?; + // TODO: we should also be setting the ObservedState on earlier API calls, in case we fail + // partway through. In fact, all location conf API calls should be in a wrapper that sets + // the observed state to None, then runs, then sets it to what we wrote. + self.observed.locations.insert( + origin_ps_id, + ObservedStateLocation { + conf: Some(origin_secondary_conf), + }, + ); + + println!( + "🔁 Switching to AttachedSingle mode on pageserver {}", + dest_ps_id + ); + let dest_final_conf = build_location_config( + &self.shard, + &self.config, + LocationConfigMode::AttachedSingle, + Some(self.generation), + None, + ); + self.location_config(dest_ps_id, dest_final_conf.clone(), None) + .await?; + self.observed.locations.insert( + dest_ps_id, + ObservedStateLocation { + conf: Some(dest_final_conf), + }, + ); + + println!("✅ Migration complete"); + + Ok(()) + } + + /// Reconciling a tenant makes API calls to pageservers until the observed state + /// matches the intended state. + /// + /// First we apply special case handling (e.g. for live migrations), and then a + /// general case reconciliation where we walk through the intent by pageserver + /// and call out to the pageserver to apply the desired state. + pub(crate) async fn reconcile(&mut self) -> Result<(), ReconcileError> { + // TODO: if any of self.observed is None, call to remote pageservers + // to learn correct state. + + // Special case: live migration + self.maybe_live_migrate().await?; + + // If the attached pageserver is not attached, do so now. + if let Some(node_id) = self.intent.attached { + let mut wanted_conf = + attached_location_conf(self.generation, &self.shard, &self.config); + match self.observed.locations.get(&node_id) { + Some(conf) if conf.conf.as_ref() == Some(&wanted_conf) => { + // Nothing to do + tracing::info!("Observed configuration already correct.") + } + _ => { + // In all cases other than a matching observed configuration, we will + // reconcile this location. This includes locations with different configurations, as well + // as locations with unknown (None) observed state. + self.generation = self + .persistence + .increment_generation(self.tenant_shard_id, Some(node_id)) + .await?; + wanted_conf.generation = self.generation.into(); + tracing::info!("Observed configuration requires update."); + self.location_config(node_id, wanted_conf, None).await?; + if let Err(e) = self + .compute_hook + .notify(self.tenant_shard_id, node_id) + .await + { + tracing::warn!( + "Failed to notify compute of newly attached pageserver {node_id}: {e}" + ); + } + } + } + } + + // Configure secondary locations: if these were previously attached this + // implicitly downgrades them from attached to secondary. + let mut changes = Vec::new(); + for node_id in &self.intent.secondary { + let wanted_conf = secondary_location_conf(&self.shard, &self.config); + match self.observed.locations.get(node_id) { + Some(conf) if conf.conf.as_ref() == Some(&wanted_conf) => { + // Nothing to do + tracing::info!(%node_id, "Observed configuration already correct.") + } + _ => { + // In all cases other than a matching observed configuration, we will + // reconcile this location. + tracing::info!(%node_id, "Observed configuration requires update."); + changes.push((*node_id, wanted_conf)) + } + } + } + + // Detach any extraneous pageservers that are no longer referenced + // by our intent. + let all_pageservers = self.intent.all_pageservers(); + for node_id in self.observed.locations.keys() { + if all_pageservers.contains(node_id) { + // We are only detaching pageservers that aren't used at all. + continue; + } + + changes.push(( + *node_id, + LocationConfig { + mode: LocationConfigMode::Detached, + generation: None, + secondary_conf: None, + shard_number: self.shard.number.0, + shard_count: self.shard.count.0, + shard_stripe_size: self.shard.stripe_size.0, + tenant_conf: self.config.clone(), + }, + )); + } + + for (node_id, conf) in changes { + self.location_config(node_id, conf, None).await?; + } + + Ok(()) + } +} + +pub(crate) fn attached_location_conf( + generation: Generation, + shard: &ShardIdentity, + config: &TenantConfig, +) -> LocationConfig { + LocationConfig { + mode: LocationConfigMode::AttachedSingle, + generation: generation.into(), + secondary_conf: None, + shard_number: shard.number.0, + shard_count: shard.count.0, + shard_stripe_size: shard.stripe_size.0, + tenant_conf: config.clone(), + } +} + +pub(crate) fn secondary_location_conf( + shard: &ShardIdentity, + config: &TenantConfig, +) -> LocationConfig { + LocationConfig { + mode: LocationConfigMode::Secondary, + generation: None, + secondary_conf: Some(LocationConfigSecondary { warm: true }), + shard_number: shard.number.0, + shard_count: shard.count.0, + shard_stripe_size: shard.stripe_size.0, + tenant_conf: config.clone(), + } +} diff --git a/control_plane/attachment_service/src/scheduler.rs b/control_plane/attachment_service/src/scheduler.rs new file mode 100644 index 0000000000..1966a7ea2a --- /dev/null +++ b/control_plane/attachment_service/src/scheduler.rs @@ -0,0 +1,89 @@ +use pageserver_api::shard::TenantShardId; +use std::collections::{BTreeMap, HashMap}; +use utils::{http::error::ApiError, id::NodeId}; + +use crate::{node::Node, tenant_state::TenantState}; + +/// Scenarios in which we cannot find a suitable location for a tenant shard +#[derive(thiserror::Error, Debug)] +pub enum ScheduleError { + #[error("No pageservers found")] + NoPageservers, + #[error("No pageserver found matching constraint")] + ImpossibleConstraint, +} + +impl From for ApiError { + fn from(value: ScheduleError) -> Self { + ApiError::Conflict(format!("Scheduling error: {}", value)) + } +} + +pub(crate) struct Scheduler { + tenant_counts: HashMap, +} + +impl Scheduler { + pub(crate) fn new( + tenants: &BTreeMap, + nodes: &HashMap, + ) -> Self { + let mut tenant_counts = HashMap::new(); + for node_id in nodes.keys() { + tenant_counts.insert(*node_id, 0); + } + + for tenant in tenants.values() { + if let Some(ps) = tenant.intent.attached { + let entry = tenant_counts.entry(ps).or_insert(0); + *entry += 1; + } + } + + for (node_id, node) in nodes { + if !node.may_schedule() { + tenant_counts.remove(node_id); + } + } + + Self { tenant_counts } + } + + pub(crate) fn schedule_shard( + &mut self, + hard_exclude: &[NodeId], + ) -> Result { + if self.tenant_counts.is_empty() { + return Err(ScheduleError::NoPageservers); + } + + let mut tenant_counts: Vec<(NodeId, usize)> = self + .tenant_counts + .iter() + .filter_map(|(k, v)| { + if hard_exclude.contains(k) { + None + } else { + Some((*k, *v)) + } + }) + .collect(); + + // Sort by tenant count. Nodes with the same tenant count are sorted by ID. + tenant_counts.sort_by_key(|i| (i.1, i.0)); + + if tenant_counts.is_empty() { + // After applying constraints, no pageservers were left + return Err(ScheduleError::ImpossibleConstraint); + } + + for (node_id, count) in &tenant_counts { + tracing::info!("tenant_counts[{node_id}]={count}"); + } + + let node_id = tenant_counts.first().unwrap().0; + tracing::info!("scheduler selected node {node_id}"); + *self.tenant_counts.get_mut(&node_id).unwrap() += 1; + Ok(node_id) + } +} diff --git a/control_plane/attachment_service/src/service.rs b/control_plane/attachment_service/src/service.rs new file mode 100644 index 0000000000..5999d48fd9 --- /dev/null +++ b/control_plane/attachment_service/src/service.rs @@ -0,0 +1,1137 @@ +use std::{ + collections::{BTreeMap, HashMap}, + str::FromStr, + sync::Arc, + time::{Duration, Instant}, +}; + +use control_plane::attachment_service::{ + AttachHookRequest, AttachHookResponse, InspectRequest, InspectResponse, NodeAvailability, + NodeConfigureRequest, NodeRegisterRequest, NodeSchedulingPolicy, TenantCreateResponse, + TenantCreateResponseShard, TenantLocateResponse, TenantLocateResponseShard, + TenantShardMigrateRequest, TenantShardMigrateResponse, +}; +use hyper::StatusCode; +use pageserver_api::{ + control_api::{ + ReAttachRequest, ReAttachResponse, ReAttachResponseTenant, ValidateRequest, + ValidateResponse, ValidateResponseTenant, + }, + models, + models::{ + LocationConfig, LocationConfigMode, ShardParameters, TenantConfig, TenantCreateRequest, + TimelineCreateRequest, TimelineInfo, + }, + shard::{ShardCount, ShardIdentity, ShardNumber, ShardStripeSize, TenantShardId}, +}; +use pageserver_client::mgmt_api; +use utils::{ + generation::Generation, + http::error::ApiError, + id::{NodeId, TenantId}, + seqwait::SeqWait, +}; + +use crate::{ + compute_hook::ComputeHook, + node::Node, + persistence::{Persistence, TenantShardPersistence}, + scheduler::Scheduler, + tenant_state::{ + IntentState, ObservedState, ObservedStateLocation, ReconcileResult, ReconcileWaitError, + ReconcilerWaiter, TenantState, + }, + PlacementPolicy, Sequence, +}; + +const RECONCILE_TIMEOUT: Duration = Duration::from_secs(30); + +// Top level state available to all HTTP handlers +struct ServiceState { + tenants: BTreeMap, + + nodes: Arc>, + + compute_hook: Arc, + + result_tx: tokio::sync::mpsc::UnboundedSender, +} + +impl ServiceState { + fn new( + result_tx: tokio::sync::mpsc::UnboundedSender, + nodes: HashMap, + tenants: BTreeMap, + ) -> Self { + Self { + tenants, + nodes: Arc::new(nodes), + compute_hook: Arc::new(ComputeHook::new()), + result_tx, + } + } +} + +#[derive(Clone)] +pub struct Config { + // All pageservers managed by one instance of this service must have + // the same public key. + pub jwt_token: Option, +} + +pub struct Service { + inner: Arc>, + config: Config, + persistence: Arc, +} + +impl From for ApiError { + fn from(value: ReconcileWaitError) -> Self { + match value { + ReconcileWaitError::Shutdown => ApiError::ShuttingDown, + e @ ReconcileWaitError::Timeout(_) => ApiError::Timeout(format!("{e}").into()), + e @ ReconcileWaitError::Failed(..) => ApiError::InternalServerError(anyhow::anyhow!(e)), + } + } +} + +impl Service { + pub async fn spawn(config: Config, persistence: Arc) -> anyhow::Result> { + let (result_tx, mut result_rx) = tokio::sync::mpsc::unbounded_channel(); + + tracing::info!("Loading nodes from database..."); + let mut nodes = persistence.list_nodes().await?; + tracing::info!("Loaded {} nodes from database.", nodes.len()); + + tracing::info!("Loading shards from database..."); + let tenant_shard_persistence = persistence.list_tenant_shards().await?; + tracing::info!( + "Loaded {} shards from database.", + tenant_shard_persistence.len() + ); + + let mut tenants = BTreeMap::new(); + + for tsp in tenant_shard_persistence { + let tenant_shard_id = TenantShardId { + tenant_id: TenantId::from_str(tsp.tenant_id.as_str())?, + shard_number: ShardNumber(tsp.shard_number as u8), + shard_count: ShardCount(tsp.shard_count as u8), + }; + let shard_identity = if tsp.shard_count == 0 { + ShardIdentity::unsharded() + } else { + ShardIdentity::new( + ShardNumber(tsp.shard_number as u8), + ShardCount(tsp.shard_count as u8), + ShardStripeSize(tsp.shard_stripe_size as u32), + )? + }; + let new_tenant = TenantState { + tenant_shard_id, + shard: shard_identity, + sequence: Sequence::initial(), + // Note that we load generation, but don't care about generation_pageserver. We will either end up finding + // our existing attached location and it will match generation_pageserver, or we will attach somewhere new + // and update generation_pageserver in the process. + generation: Generation::new(tsp.generation), + policy: serde_json::from_str(&tsp.placement_policy).unwrap(), + intent: IntentState::new(), + observed: ObservedState::new(), + config: serde_json::from_str(&tsp.config).unwrap(), + reconciler: None, + waiter: Arc::new(SeqWait::new(Sequence::initial())), + error_waiter: Arc::new(SeqWait::new(Sequence::initial())), + last_error: Arc::default(), + }; + + tenants.insert(tenant_shard_id, new_tenant); + } + + // For all tenant shards, a vector of observed states on nodes (where None means + // indeterminate, same as in [`ObservedStateLocation`]) + let mut observed = HashMap::new(); + + // TODO: issue these requests concurrently + for node in &mut nodes { + let client = mgmt_api::Client::new(node.base_url(), config.jwt_token.as_deref()); + + tracing::info!("Scanning shards on node {}...", node.id); + match client.list_location_config().await { + Err(e) => { + tracing::warn!("Could not contact pageserver {} ({e})", node.id); + // TODO: be more tolerant, apply a generous 5-10 second timeout + // TODO: setting a node to Offline is a dramatic thing to do, and can + // prevent neon_local from starting up (it starts this service before + // any pageservers are running). It may make sense to give nodes + // a Pending state to accomodate this situation, and allow (but deprioritize) + // scheduling on Pending nodes. + //node.availability = NodeAvailability::Offline; + } + Ok(listing) => { + tracing::info!( + "Received {} shard statuses from pageserver {}, setting it to Active", + listing.tenant_shards.len(), + node.id + ); + node.availability = NodeAvailability::Active; + + for (tenant_shard_id, conf_opt) in listing.tenant_shards { + observed.insert(tenant_shard_id, (node.id, conf_opt)); + } + } + } + } + + let mut cleanup = Vec::new(); + + // Populate intent and observed states for all tenants, based on reported state on pageservers + for (tenant_shard_id, (node_id, observed_loc)) in observed { + let Some(tenant_state) = tenants.get_mut(&tenant_shard_id) else { + cleanup.push((tenant_shard_id, node_id)); + continue; + }; + + tenant_state + .observed + .locations + .insert(node_id, ObservedStateLocation { conf: observed_loc }); + } + + // State of nodes is now frozen, transform to a HashMap. + let mut nodes: HashMap = nodes.into_iter().map(|n| (n.id, n)).collect(); + + // Populate each tenant's intent state + let mut scheduler = Scheduler::new(&tenants, &nodes); + for (tenant_shard_id, tenant_state) in tenants.iter_mut() { + tenant_state.intent_from_observed(); + if let Err(e) = tenant_state.schedule(&mut scheduler) { + // Non-fatal error: we are unable to properly schedule the tenant, perhaps because + // not enough pageservers are available. The tenant may well still be available + // to clients. + tracing::error!("Failed to schedule tenant {tenant_shard_id} at startup: {e}"); + } + } + + // Clean up any tenants that were found on pageservers but are not known to us. + for (tenant_shard_id, node_id) in cleanup { + // A node reported a tenant_shard_id which is unknown to us: detach it. + let node = nodes + .get_mut(&node_id) + .expect("Always exists: only known nodes are scanned"); + + let client = mgmt_api::Client::new(node.base_url(), config.jwt_token.as_deref()); + match client + .location_config( + tenant_shard_id, + LocationConfig { + mode: LocationConfigMode::Detached, + generation: None, + secondary_conf: None, + shard_number: tenant_shard_id.shard_number.0, + shard_count: tenant_shard_id.shard_count.0, + shard_stripe_size: 0, + tenant_conf: models::TenantConfig::default(), + }, + None, + ) + .await + { + Ok(()) => { + tracing::info!( + "Detached unknown shard {tenant_shard_id} on pageserver {node_id}" + ); + } + Err(e) => { + // Non-fatal error: leaving a tenant shard behind that we are not managing shouldn't + // break anything. + tracing::error!( + "Failed to detach unknkown shard {tenant_shard_id} on pageserver {node_id}: {e}" + ); + } + } + } + + let shard_count = tenants.len(); + let this = Arc::new(Self { + inner: Arc::new(std::sync::RwLock::new(ServiceState::new( + result_tx, nodes, tenants, + ))), + config, + persistence, + }); + + let result_task_this = this.clone(); + tokio::task::spawn(async move { + while let Some(result) = result_rx.recv().await { + tracing::info!( + "Reconcile result for sequence {}, ok={}", + result.sequence, + result.result.is_ok() + ); + let mut locked = result_task_this.inner.write().unwrap(); + let Some(tenant) = locked.tenants.get_mut(&result.tenant_shard_id) else { + // A reconciliation result might race with removing a tenant: drop results for + // tenants that aren't in our map. + continue; + }; + + // Usually generation should only be updated via this path, so the max() isn't + // needed, but it is used to handle out-of-band updates via. e.g. test hook. + tenant.generation = std::cmp::max(tenant.generation, result.generation); + + match result.result { + Ok(()) => { + for (node_id, loc) in &result.observed.locations { + if let Some(conf) = &loc.conf { + tracing::info!( + "Updating observed location {}: {:?}", + node_id, + conf + ); + } else { + tracing::info!("Setting observed location {} to None", node_id,) + } + } + tenant.observed = result.observed; + tenant.waiter.advance(result.sequence); + } + Err(e) => { + tracing::warn!( + "Reconcile error on tenant {}: {}", + tenant.tenant_shard_id, + e + ); + + // Ordering: populate last_error before advancing error_seq, + // so that waiters will see the correct error after waiting. + *(tenant.last_error.lock().unwrap()) = format!("{e}"); + tenant.error_waiter.advance(result.sequence); + + for (node_id, o) in result.observed.locations { + tenant.observed.locations.insert(node_id, o); + } + } + } + } + }); + + // Finally, now that the service is up and running, launch reconcile operations for any tenants + // which require it: under normal circumstances this should only include tenants that were in some + // transient state before we restarted. + let reconcile_tasks = this.reconcile_all(); + tracing::info!("Startup complete, spawned {reconcile_tasks} reconciliation tasks ({shard_count} shards total)"); + + Ok(this) + } + + pub(crate) async fn attach_hook( + &self, + attach_req: AttachHookRequest, + ) -> anyhow::Result { + // This is a test hook. To enable using it on tenants that were created directly with + // the pageserver API (not via this service), we will auto-create any missing tenant + // shards with default state. + let insert = { + let locked = self.inner.write().unwrap(); + !locked.tenants.contains_key(&attach_req.tenant_shard_id) + }; + + if insert { + let tsp = TenantShardPersistence { + tenant_id: attach_req.tenant_shard_id.tenant_id.to_string(), + shard_number: attach_req.tenant_shard_id.shard_number.0 as i32, + shard_count: attach_req.tenant_shard_id.shard_count.0 as i32, + shard_stripe_size: 0, + generation: 0, + generation_pageserver: None, + placement_policy: serde_json::to_string(&PlacementPolicy::default()).unwrap(), + config: serde_json::to_string(&TenantConfig::default()).unwrap(), + }; + + self.persistence.insert_tenant_shards(vec![tsp]).await?; + + let mut locked = self.inner.write().unwrap(); + locked.tenants.insert( + attach_req.tenant_shard_id, + TenantState::new( + attach_req.tenant_shard_id, + ShardIdentity::unsharded(), + PlacementPolicy::Single, + ), + ); + } + + let new_generation = if attach_req.node_id.is_some() { + Some( + self.persistence + .increment_generation(attach_req.tenant_shard_id, attach_req.node_id) + .await?, + ) + } else { + None + }; + + let mut locked = self.inner.write().unwrap(); + let tenant_state = locked + .tenants + .get_mut(&attach_req.tenant_shard_id) + .expect("Checked for existence above"); + + if let Some(new_generation) = new_generation { + tenant_state.generation = new_generation; + } + + if let Some(attaching_pageserver) = attach_req.node_id.as_ref() { + tracing::info!( + tenant_id = %attach_req.tenant_shard_id, + ps_id = %attaching_pageserver, + generation = ?tenant_state.generation, + "issuing", + ); + } else if let Some(ps_id) = tenant_state.intent.attached { + tracing::info!( + tenant_id = %attach_req.tenant_shard_id, + %ps_id, + generation = ?tenant_state.generation, + "dropping", + ); + } else { + tracing::info!( + tenant_id = %attach_req.tenant_shard_id, + "no-op: tenant already has no pageserver"); + } + tenant_state.intent.attached = attach_req.node_id; + + tracing::info!( + "attach_hook: tenant {} set generation {:?}, pageserver {}", + attach_req.tenant_shard_id, + tenant_state.generation, + attach_req.node_id.unwrap_or(utils::id::NodeId(0xfffffff)) + ); + + Ok(AttachHookResponse { + gen: attach_req + .node_id + .map(|_| tenant_state.generation.into().unwrap()), + }) + } + + pub(crate) fn inspect(&self, inspect_req: InspectRequest) -> InspectResponse { + let locked = self.inner.read().unwrap(); + + let tenant_state = locked.tenants.get(&inspect_req.tenant_shard_id); + + InspectResponse { + attachment: tenant_state.and_then(|s| { + s.intent + .attached + .map(|ps| (s.generation.into().unwrap(), ps)) + }), + } + } + + pub(crate) async fn re_attach( + &self, + reattach_req: ReAttachRequest, + ) -> anyhow::Result { + // Ordering: we must persist generation number updates before making them visible in the in-memory state + let incremented_generations = self.persistence.re_attach(reattach_req.node_id).await?; + + // Apply the updated generation to our in-memory state + let mut locked = self.inner.write().unwrap(); + + let mut response = ReAttachResponse { + tenants: Vec::new(), + }; + + for (tenant_shard_id, new_gen) in incremented_generations { + response.tenants.push(ReAttachResponseTenant { + id: tenant_shard_id, + gen: new_gen.into().unwrap(), + }); + + // Apply the new generation number to our in-memory state + let shard_state = locked.tenants.get_mut(&tenant_shard_id); + let Some(shard_state) = shard_state else { + // Not fatal. This edge case requires a re-attach to happen + // between inserting a new tenant shard in to the database, and updating our in-memory + // state to know about the shard, _and_ that the state inserted to the database referenced + // a pageserver. Should never happen, but handle it rather than panicking, since it should + // be harmless. + tracing::error!( + "Shard {} is in database for node {} but not in-memory state", + tenant_shard_id, + reattach_req.node_id + ); + continue; + }; + + shard_state.generation = std::cmp::max(shard_state.generation, new_gen); + + // TODO: cancel/restart any running reconciliation for this tenant, it might be trying + // to call location_conf API with an old generation. Wait for cancellation to complete + // before responding to this request. Requires well implemented CancellationToken logic + // all the way to where we call location_conf. Even then, there can still be a location_conf + // request in flight over the network: TODO handle that by making location_conf API refuse + // to go backward in generations. + } + Ok(response) + } + + pub(crate) fn validate(&self, validate_req: ValidateRequest) -> ValidateResponse { + let locked = self.inner.read().unwrap(); + + let mut response = ValidateResponse { + tenants: Vec::new(), + }; + + for req_tenant in validate_req.tenants { + if let Some(tenant_state) = locked.tenants.get(&req_tenant.id) { + let valid = tenant_state.generation == Generation::new(req_tenant.gen); + tracing::info!( + "handle_validate: {}(gen {}): valid={valid} (latest {:?})", + req_tenant.id, + req_tenant.gen, + tenant_state.generation + ); + response.tenants.push(ValidateResponseTenant { + id: req_tenant.id, + valid, + }); + } + } + response + } + + pub(crate) async fn tenant_create( + &self, + create_req: TenantCreateRequest, + ) -> Result { + // Shard count 0 is valid: it means create a single shard (ShardCount(0) means "unsharded") + let literal_shard_count = if create_req.shard_parameters.is_unsharded() { + 1 + } else { + create_req.shard_parameters.count.0 + }; + + // This service expects to handle sharding itself: it is an error to try and directly create + // a particular shard here. + let tenant_id = if create_req.new_tenant_id.shard_count > ShardCount(1) { + return Err(ApiError::BadRequest(anyhow::anyhow!( + "Attempted to create a specific shard, this API is for creating the whole tenant" + ))); + } else { + create_req.new_tenant_id.tenant_id + }; + + tracing::info!( + "Creating tenant {}, shard_count={:?}", + create_req.new_tenant_id, + create_req.shard_parameters.count, + ); + + let create_ids = (0..literal_shard_count) + .map(|i| TenantShardId { + tenant_id, + shard_number: ShardNumber(i), + shard_count: create_req.shard_parameters.count, + }) + .collect::>(); + + // TODO: enable specifying this. Using Single as a default helps legacy tests to work (they + // have no expectation of HA). + let placement_policy: PlacementPolicy = PlacementPolicy::Single; + + // Ordering: we persist tenant shards before creating them on the pageserver. This enables a caller + // to clean up after themselves by issuing a tenant deletion if something goes wrong and we restart + // during the creation, rather than risking leaving orphan objects in S3. + let persist_tenant_shards = create_ids + .iter() + .map(|tenant_shard_id| TenantShardPersistence { + tenant_id: tenant_shard_id.tenant_id.to_string(), + shard_number: tenant_shard_id.shard_number.0 as i32, + shard_count: tenant_shard_id.shard_count.0 as i32, + shard_stripe_size: create_req.shard_parameters.stripe_size.0 as i32, + generation: 0, + generation_pageserver: None, + placement_policy: serde_json::to_string(&placement_policy).unwrap(), + config: serde_json::to_string(&create_req.config).unwrap(), + }) + .collect(); + self.persistence + .insert_tenant_shards(persist_tenant_shards) + .await + .map_err(|e| { + // TODO: distinguish primary key constraint (idempotent, OK), from other errors + ApiError::InternalServerError(anyhow::anyhow!(e)) + })?; + + let (waiters, response_shards) = { + let mut locked = self.inner.write().unwrap(); + + let mut response_shards = Vec::new(); + + let mut scheduler = Scheduler::new(&locked.tenants, &locked.nodes); + + for tenant_shard_id in create_ids { + tracing::info!("Creating shard {tenant_shard_id}..."); + + use std::collections::btree_map::Entry; + match locked.tenants.entry(tenant_shard_id) { + Entry::Occupied(mut entry) => { + tracing::info!( + "Tenant shard {tenant_shard_id} already exists while creating" + ); + + // TODO: schedule() should take an anti-affinity expression that pushes + // attached and secondary locations (independently) away frorm those + // pageservers also holding a shard for this tenant. + + entry.get_mut().schedule(&mut scheduler).map_err(|e| { + ApiError::Conflict(format!( + "Failed to schedule shard {tenant_shard_id}: {e}" + )) + })?; + + response_shards.push(TenantCreateResponseShard { + node_id: entry + .get() + .intent + .attached + .expect("We just set pageserver if it was None"), + generation: entry.get().generation.into().unwrap(), + }); + + continue; + } + Entry::Vacant(entry) => { + let mut state = TenantState::new( + tenant_shard_id, + ShardIdentity::from_params( + tenant_shard_id.shard_number, + &create_req.shard_parameters, + ), + placement_policy.clone(), + ); + + if let Some(create_gen) = create_req.generation { + state.generation = Generation::new(create_gen); + } + state.config = create_req.config.clone(); + + state.schedule(&mut scheduler).map_err(|e| { + ApiError::Conflict(format!( + "Failed to schedule shard {tenant_shard_id}: {e}" + )) + })?; + + response_shards.push(TenantCreateResponseShard { + node_id: state + .intent + .attached + .expect("We just set pageserver if it was None"), + generation: state.generation.into().unwrap(), + }); + entry.insert(state) + } + }; + } + + // Take a snapshot of pageservers + let pageservers = locked.nodes.clone(); + + let result_tx = locked.result_tx.clone(); + let compute_hook = locked.compute_hook.clone(); + + let waiters = locked + .tenants + .range_mut(TenantShardId::tenant_range(tenant_id)) + .filter_map(|(_shard_id, shard)| { + shard.maybe_reconcile( + result_tx.clone(), + &pageservers, + &compute_hook, + &self.config, + &self.persistence, + ) + }) + .collect::>(); + (waiters, response_shards) + }; + + let deadline = Instant::now().checked_add(Duration::from_secs(5)).unwrap(); + for waiter in waiters { + let timeout = deadline.duration_since(Instant::now()); + waiter.wait_timeout(timeout).await?; + } + Ok(TenantCreateResponse { + shards: response_shards, + }) + } + + pub(crate) async fn tenant_timeline_create( + &self, + tenant_id: TenantId, + mut create_req: TimelineCreateRequest, + ) -> Result { + let mut timeline_info = None; + + let ensure_waiters = { + let locked = self.inner.write().unwrap(); + tracing::info!( + "Creating timeline {}/{}, have {} pageservers", + tenant_id, + create_req.new_timeline_id, + locked.nodes.len() + ); + + self.ensure_attached(locked, tenant_id) + .map_err(ApiError::InternalServerError)? + }; + + let deadline = Instant::now().checked_add(Duration::from_secs(5)).unwrap(); + for waiter in ensure_waiters { + let timeout = deadline.duration_since(Instant::now()); + waiter.wait_timeout(timeout).await?; + } + + let targets = { + let locked = self.inner.read().unwrap(); + let mut targets = Vec::new(); + + for (tenant_shard_id, shard) in + locked.tenants.range(TenantShardId::tenant_range(tenant_id)) + { + let node_id = shard.intent.attached.ok_or_else(|| { + ApiError::InternalServerError(anyhow::anyhow!("Shard not scheduled")) + })?; + let node = locked + .nodes + .get(&node_id) + .expect("Pageservers may not be deleted while referenced"); + + targets.push((*tenant_shard_id, node.clone())); + } + targets + }; + + if targets.is_empty() { + return Err(ApiError::NotFound( + anyhow::anyhow!("Tenant not found").into(), + )); + } + + for (tenant_shard_id, node) in targets { + // TODO: issue shard timeline creates in parallel, once the 0th is done. + + let client = mgmt_api::Client::new(node.base_url(), self.config.jwt_token.as_deref()); + + tracing::info!( + "Creating timeline on shard {}/{}, attached to node {}", + tenant_shard_id, + create_req.new_timeline_id, + node.id + ); + + let shard_timeline_info = client + .timeline_create(tenant_shard_id, &create_req) + .await + .map_err(|e| match e { + mgmt_api::Error::ApiError(status, msg) + if status == StatusCode::INTERNAL_SERVER_ERROR + || status == StatusCode::NOT_ACCEPTABLE => + { + // TODO: handle more error codes, e.g. 503 should be passed through. Make a general wrapper + // for pass-through API calls. + ApiError::InternalServerError(anyhow::anyhow!(msg)) + } + _ => ApiError::Conflict(format!("Failed to create timeline: {e}")), + })?; + + if timeline_info.is_none() { + // If the caller specified an ancestor but no ancestor LSN, we are responsible for + // propagating the LSN chosen by the first shard to the other shards: it is important + // that all shards end up with the same ancestor_start_lsn. + if create_req.ancestor_timeline_id.is_some() + && create_req.ancestor_start_lsn.is_none() + { + create_req.ancestor_start_lsn = shard_timeline_info.ancestor_lsn; + } + + // We will return the TimelineInfo from the first shard + timeline_info = Some(shard_timeline_info); + } + } + Ok(timeline_info.expect("targets cannot be empty")) + } + + pub(crate) fn tenant_locate( + &self, + tenant_id: TenantId, + ) -> Result { + let locked = self.inner.read().unwrap(); + tracing::info!("Locating shards for tenant {tenant_id}"); + + // Take a snapshot of pageservers + let pageservers = locked.nodes.clone(); + + let mut result = Vec::new(); + let mut shard_params: Option = None; + + for (tenant_shard_id, shard) in locked.tenants.range(TenantShardId::tenant_range(tenant_id)) + { + let node_id = shard + .intent + .attached + .ok_or(ApiError::BadRequest(anyhow::anyhow!( + "Cannot locate a tenant that is not attached" + )))?; + + let node = pageservers + .get(&node_id) + .expect("Pageservers may not be deleted while referenced"); + + result.push(TenantLocateResponseShard { + shard_id: *tenant_shard_id, + node_id, + listen_http_addr: node.listen_http_addr.clone(), + listen_http_port: node.listen_http_port, + listen_pg_addr: node.listen_pg_addr.clone(), + listen_pg_port: node.listen_pg_port, + }); + + match &shard_params { + None => { + shard_params = Some(ShardParameters { + stripe_size: shard.shard.stripe_size, + count: shard.shard.count, + }); + } + Some(params) => { + if params.stripe_size != shard.shard.stripe_size { + // This should never happen. We enforce at runtime because it's simpler than + // adding an extra per-tenant data structure to store the things that should be the same + return Err(ApiError::InternalServerError(anyhow::anyhow!( + "Inconsistent shard stripe size parameters!" + ))); + } + } + } + } + + if result.is_empty() { + return Err(ApiError::NotFound( + anyhow::anyhow!("No shards for this tenant ID found").into(), + )); + } + let shard_params = shard_params.expect("result is non-empty, therefore this is set"); + tracing::info!( + "Located tenant {} with params {:?} on shards {}", + tenant_id, + shard_params, + result + .iter() + .map(|s| format!("{:?}", s)) + .collect::>() + .join(",") + ); + + Ok(TenantLocateResponse { + shards: result, + shard_params, + }) + } + + pub(crate) async fn tenant_shard_migrate( + &self, + tenant_shard_id: TenantShardId, + migrate_req: TenantShardMigrateRequest, + ) -> Result { + let waiter = { + let mut locked = self.inner.write().unwrap(); + + let result_tx = locked.result_tx.clone(); + let pageservers = locked.nodes.clone(); + let compute_hook = locked.compute_hook.clone(); + + let Some(shard) = locked.tenants.get_mut(&tenant_shard_id) else { + return Err(ApiError::NotFound( + anyhow::anyhow!("Tenant shard not found").into(), + )); + }; + + if shard.intent.attached == Some(migrate_req.node_id) { + // No-op case: we will still proceed to wait for reconciliation in case it is + // incomplete from an earlier update to the intent. + tracing::info!("Migrating: intent is unchanged {:?}", shard.intent); + } else { + let old_attached = shard.intent.attached; + + shard.intent.attached = Some(migrate_req.node_id); + match shard.policy { + PlacementPolicy::Single => { + shard.intent.secondary.clear(); + } + PlacementPolicy::Double(_n) => { + // If our new attached node was a secondary, it no longer should be. + shard.intent.secondary.retain(|s| s != &migrate_req.node_id); + + // If we were already attached to something, demote that to a secondary + if let Some(old_attached) = old_attached { + shard.intent.secondary.push(old_attached); + } + } + } + + tracing::info!("Migrating: new intent {:?}", shard.intent); + shard.sequence = shard.sequence.next(); + } + + shard.maybe_reconcile( + result_tx, + &pageservers, + &compute_hook, + &self.config, + &self.persistence, + ) + }; + + if let Some(waiter) = waiter { + waiter.wait_timeout(RECONCILE_TIMEOUT).await?; + } else { + tracing::warn!("Migration is a no-op"); + } + + Ok(TenantShardMigrateResponse {}) + } + + pub(crate) async fn node_register( + &self, + register_req: NodeRegisterRequest, + ) -> Result<(), ApiError> { + // Pre-check for an already-existing node + { + let locked = self.inner.read().unwrap(); + if let Some(node) = locked.nodes.get(®ister_req.node_id) { + // Note that we do not do a total equality of the struct, because we don't require + // the availability/scheduling states to agree for a POST to be idempotent. + if node.listen_http_addr == register_req.listen_http_addr + && node.listen_http_port == register_req.listen_http_port + && node.listen_pg_addr == register_req.listen_pg_addr + && node.listen_pg_port == register_req.listen_pg_port + { + tracing::info!( + "Node {} re-registered with matching address", + register_req.node_id + ); + return Ok(()); + } else { + // TODO: decide if we want to allow modifying node addresses without removing and re-adding + // the node. Safest/simplest thing is to refuse it, and usually we deploy with + // a fixed address through the lifetime of a node. + tracing::warn!( + "Node {} tried to register with different address", + register_req.node_id + ); + return Err(ApiError::Conflict( + "Node is already registered with different address".to_string(), + )); + } + } + } + + // Ordering: we must persist the new node _before_ adding it to in-memory state. + // This ensures that before we use it for anything or expose it via any external + // API, it is guaranteed to be available after a restart. + let new_node = Node { + id: register_req.node_id, + listen_http_addr: register_req.listen_http_addr, + listen_http_port: register_req.listen_http_port, + listen_pg_addr: register_req.listen_pg_addr, + listen_pg_port: register_req.listen_pg_port, + scheduling: NodeSchedulingPolicy::Filling, + // TODO: we shouldn't really call this Active until we've heartbeated it. + availability: NodeAvailability::Active, + }; + // TODO: idempotency if the node already exists in the database + self.persistence + .insert_node(&new_node) + .await + .map_err(ApiError::InternalServerError)?; + + let mut locked = self.inner.write().unwrap(); + let mut new_nodes = (*locked.nodes).clone(); + + new_nodes.insert(register_req.node_id, new_node); + + locked.nodes = Arc::new(new_nodes); + + tracing::info!( + "Registered pageserver {}, now have {} pageservers", + register_req.node_id, + locked.nodes.len() + ); + Ok(()) + } + + pub(crate) fn node_configure(&self, config_req: NodeConfigureRequest) -> Result<(), ApiError> { + let mut locked = self.inner.write().unwrap(); + let result_tx = locked.result_tx.clone(); + let compute_hook = locked.compute_hook.clone(); + + let mut new_nodes = (*locked.nodes).clone(); + + let Some(node) = new_nodes.get_mut(&config_req.node_id) else { + return Err(ApiError::NotFound( + anyhow::anyhow!("Node not registered").into(), + )); + }; + + let mut offline_transition = false; + let mut active_transition = false; + + if let Some(availability) = &config_req.availability { + match (availability, &node.availability) { + (NodeAvailability::Offline, NodeAvailability::Active) => { + tracing::info!("Node {} transition to offline", config_req.node_id); + offline_transition = true; + } + (NodeAvailability::Active, NodeAvailability::Offline) => { + tracing::info!("Node {} transition to active", config_req.node_id); + active_transition = true; + } + _ => { + tracing::info!("Node {} no change during config", config_req.node_id); + // No change + } + }; + node.availability = *availability; + } + + if let Some(scheduling) = config_req.scheduling { + node.scheduling = scheduling; + + // TODO: once we have a background scheduling ticker for fill/drain, kick it + // to wake up and start working. + } + + let new_nodes = Arc::new(new_nodes); + + let mut scheduler = Scheduler::new(&locked.tenants, &new_nodes); + if offline_transition { + for (tenant_shard_id, tenant_state) in &mut locked.tenants { + if let Some(observed_loc) = + tenant_state.observed.locations.get_mut(&config_req.node_id) + { + // When a node goes offline, we set its observed configuration to None, indicating unknown: we will + // not assume our knowledge of the node's configuration is accurate until it comes back online + observed_loc.conf = None; + } + + if tenant_state.intent.notify_offline(config_req.node_id) { + tenant_state.sequence = tenant_state.sequence.next(); + match tenant_state.schedule(&mut scheduler) { + Err(e) => { + // It is possible that some tenants will become unschedulable when too many pageservers + // go offline: in this case there isn't much we can do other than make the issue observable. + // TODO: give TenantState a scheduling error attribute to be queried later. + tracing::warn!(%tenant_shard_id, "Scheduling error when marking pageserver {} offline: {e}", config_req.node_id); + } + Ok(()) => { + tenant_state.maybe_reconcile( + result_tx.clone(), + &new_nodes, + &compute_hook, + &self.config, + &self.persistence, + ); + } + } + } + } + } + + if active_transition { + // When a node comes back online, we must reconcile any tenant that has a None observed + // location on the node. + for tenant_state in locked.tenants.values_mut() { + if let Some(observed_loc) = + tenant_state.observed.locations.get_mut(&config_req.node_id) + { + if observed_loc.conf.is_none() { + tenant_state.maybe_reconcile( + result_tx.clone(), + &new_nodes, + &compute_hook, + &self.config, + &self.persistence, + ); + } + } + } + + // TODO: in the background, we should balance work back onto this pageserver + } + + locked.nodes = new_nodes; + + Ok(()) + } + + /// Helper for methods that will try and call pageserver APIs for + /// a tenant, such as timeline CRUD: they cannot proceed unless the tenant + /// is attached somewhere. + fn ensure_attached( + &self, + mut locked: std::sync::RwLockWriteGuard<'_, ServiceState>, + tenant_id: TenantId, + ) -> Result, anyhow::Error> { + let mut waiters = Vec::new(); + let result_tx = locked.result_tx.clone(); + let compute_hook = locked.compute_hook.clone(); + let mut scheduler = Scheduler::new(&locked.tenants, &locked.nodes); + let pageservers = locked.nodes.clone(); + + for (_tenant_shard_id, shard) in locked + .tenants + .range_mut(TenantShardId::tenant_range(tenant_id)) + { + shard.schedule(&mut scheduler)?; + + if let Some(waiter) = shard.maybe_reconcile( + result_tx.clone(), + &pageservers, + &compute_hook, + &self.config, + &self.persistence, + ) { + waiters.push(waiter); + } + } + Ok(waiters) + } + + /// Check all tenants for pending reconciliation work, and reconcile those in need + /// + /// Returns how many reconciliation tasks were started + fn reconcile_all(&self) -> usize { + let mut locked = self.inner.write().unwrap(); + let result_tx = locked.result_tx.clone(); + let compute_hook = locked.compute_hook.clone(); + let pageservers = locked.nodes.clone(); + locked + .tenants + .iter_mut() + .filter_map(|(_tenant_shard_id, shard)| { + shard.maybe_reconcile( + result_tx.clone(), + &pageservers, + &compute_hook, + &self.config, + &self.persistence, + ) + }) + .count() + } +} diff --git a/control_plane/attachment_service/src/tenant_state.rs b/control_plane/attachment_service/src/tenant_state.rs new file mode 100644 index 0000000000..a907628eff --- /dev/null +++ b/control_plane/attachment_service/src/tenant_state.rs @@ -0,0 +1,455 @@ +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use control_plane::attachment_service::NodeAvailability; +use pageserver_api::{ + models::{LocationConfig, LocationConfigMode, TenantConfig}, + shard::{ShardIdentity, TenantShardId}, +}; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; +use utils::{ + generation::Generation, + id::NodeId, + seqwait::{SeqWait, SeqWaitError}, +}; + +use crate::{ + compute_hook::ComputeHook, + node::Node, + persistence::Persistence, + reconciler::{attached_location_conf, secondary_location_conf, ReconcileError, Reconciler}, + scheduler::{ScheduleError, Scheduler}, + service, PlacementPolicy, Sequence, +}; + +pub(crate) struct TenantState { + pub(crate) tenant_shard_id: TenantShardId, + + pub(crate) shard: ShardIdentity, + + // Runtime only: sequence used to coordinate when updating this object while + // with background reconcilers may be running. A reconciler runs to a particular + // sequence. + pub(crate) sequence: Sequence, + + // Latest generation number: next time we attach, increment this + // and use the incremented number when attaching + pub(crate) generation: Generation, + + // High level description of how the tenant should be set up. Provided + // externally. + pub(crate) policy: PlacementPolicy, + + // Low level description of exactly which pageservers should fulfil + // which role. Generated by `Self::schedule`. + pub(crate) intent: IntentState, + + // Low level description of how the tenant is configured on pageservers: + // if this does not match `Self::intent` then the tenant needs reconciliation + // with `Self::reconcile`. + pub(crate) observed: ObservedState, + + // Tenant configuration, passed through opaquely to the pageserver. Identical + // for all shards in a tenant. + pub(crate) config: TenantConfig, + + /// If a reconcile task is currently in flight, it may be joined here (it is + /// only safe to join if either the result has been received or the reconciler's + /// cancellation token has been fired) + pub(crate) reconciler: Option, + + /// Optionally wait for reconciliation to complete up to a particular + /// sequence number. + pub(crate) waiter: std::sync::Arc>, + + /// Indicates sequence number for which we have encountered an error reconciling. If + /// this advances ahead of [`Self::waiter`] then a reconciliation error has occurred, + /// and callers should stop waiting for `waiter` and propagate the error. + pub(crate) error_waiter: std::sync::Arc>, + + /// The most recent error from a reconcile on this tenant + /// TODO: generalize to an array of recent events + /// TOOD: use a ArcSwap instead of mutex for faster reads? + pub(crate) last_error: std::sync::Arc>, +} + +#[derive(Default, Clone, Debug)] +pub(crate) struct IntentState { + pub(crate) attached: Option, + pub(crate) secondary: Vec, +} + +#[derive(Default, Clone)] +pub(crate) struct ObservedState { + pub(crate) locations: HashMap, +} + +/// Our latest knowledge of how this tenant is configured in the outside world. +/// +/// Meaning: +/// * No instance of this type exists for a node: we are certain that we have nothing configured on that +/// node for this shard. +/// * Instance exists with conf==None: we *might* have some state on that node, but we don't know +/// what it is (e.g. we failed partway through configuring it) +/// * Instance exists with conf==Some: this tells us what we last successfully configured on this node, +/// and that configuration will still be present unless something external interfered. +#[derive(Clone)] +pub(crate) struct ObservedStateLocation { + /// If None, it means we do not know the status of this shard's location on this node, but + /// we know that we might have some state on this node. + pub(crate) conf: Option, +} +pub(crate) struct ReconcilerWaiter { + // For observability purposes, remember the ID of the shard we're + // waiting for. + pub(crate) tenant_shard_id: TenantShardId, + + seq_wait: std::sync::Arc>, + error_seq_wait: std::sync::Arc>, + error: std::sync::Arc>, + seq: Sequence, +} + +#[derive(thiserror::Error, Debug)] +pub enum ReconcileWaitError { + #[error("Timeout waiting for shard {0}")] + Timeout(TenantShardId), + #[error("shutting down")] + Shutdown, + #[error("Reconcile error on shard {0}: {1}")] + Failed(TenantShardId, String), +} + +impl ReconcilerWaiter { + pub(crate) async fn wait_timeout(&self, timeout: Duration) -> Result<(), ReconcileWaitError> { + tokio::select! { + result = self.seq_wait.wait_for_timeout(self.seq, timeout)=> { + result.map_err(|e| match e { + SeqWaitError::Timeout => ReconcileWaitError::Timeout(self.tenant_shard_id), + SeqWaitError::Shutdown => ReconcileWaitError::Shutdown + })?; + }, + result = self.error_seq_wait.wait_for(self.seq) => { + result.map_err(|e| match e { + SeqWaitError::Shutdown => ReconcileWaitError::Shutdown, + SeqWaitError::Timeout => unreachable!() + })?; + + return Err(ReconcileWaitError::Failed(self.tenant_shard_id, self.error.lock().unwrap().clone())) + } + } + + Ok(()) + } +} + +/// Having spawned a reconciler task, the tenant shard's state will carry enough +/// information to optionally cancel & await it later. +pub(crate) struct ReconcilerHandle { + sequence: Sequence, + handle: JoinHandle<()>, + cancel: CancellationToken, +} + +/// When a reconcile task completes, it sends this result object +/// to be applied to the primary TenantState. +pub(crate) struct ReconcileResult { + pub(crate) sequence: Sequence, + /// On errors, `observed` should be treated as an incompleted description + /// of state (i.e. any nodes present in the result should override nodes + /// present in the parent tenant state, but any unmentioned nodes should + /// not be removed from parent tenant state) + pub(crate) result: Result<(), ReconcileError>, + + pub(crate) tenant_shard_id: TenantShardId, + pub(crate) generation: Generation, + pub(crate) observed: ObservedState, +} + +impl IntentState { + pub(crate) fn new() -> Self { + Self { + attached: None, + secondary: vec![], + } + } + pub(crate) fn all_pageservers(&self) -> Vec { + let mut result = Vec::new(); + if let Some(p) = self.attached { + result.push(p) + } + + result.extend(self.secondary.iter().copied()); + + result + } + + /// When a node goes offline, we update intents to avoid using it + /// as their attached pageserver. + /// + /// Returns true if a change was made + pub(crate) fn notify_offline(&mut self, node_id: NodeId) -> bool { + if self.attached == Some(node_id) { + self.attached = None; + self.secondary.push(node_id); + true + } else { + false + } + } +} + +impl ObservedState { + pub(crate) fn new() -> Self { + Self { + locations: HashMap::new(), + } + } +} + +impl TenantState { + pub(crate) fn new( + tenant_shard_id: TenantShardId, + shard: ShardIdentity, + policy: PlacementPolicy, + ) -> Self { + Self { + tenant_shard_id, + policy, + intent: IntentState::default(), + generation: Generation::new(0), + shard, + observed: ObservedState::default(), + config: TenantConfig::default(), + reconciler: None, + sequence: Sequence(1), + waiter: Arc::new(SeqWait::new(Sequence(0))), + error_waiter: Arc::new(SeqWait::new(Sequence(0))), + last_error: Arc::default(), + } + } + + /// For use on startup when learning state from pageservers: generate my [`IntentState`] from my + /// [`ObservedState`], even if it violates my [`PlacementPolicy`]. Call [`Self::schedule`] next, + /// to get an intent state that complies with placement policy. The overall goal is to do scheduling + /// in a way that makes use of any configured locations that already exist in the outside world. + pub(crate) fn intent_from_observed(&mut self) { + // Choose an attached location by filtering observed locations, and then sorting to get the highest + // generation + let mut attached_locs = self + .observed + .locations + .iter() + .filter_map(|(node_id, l)| { + if let Some(conf) = &l.conf { + if conf.mode == LocationConfigMode::AttachedMulti + || conf.mode == LocationConfigMode::AttachedSingle + || conf.mode == LocationConfigMode::AttachedStale + { + Some((node_id, conf.generation)) + } else { + None + } + } else { + None + } + }) + .collect::>(); + + attached_locs.sort_by_key(|i| i.1); + if let Some((node_id, _gen)) = attached_locs.into_iter().last() { + self.intent.attached = Some(*node_id); + } + + // All remaining observed locations generate secondary intents. This includes None + // observations, as these may well have some local content on disk that is usable (this + // is an edge case that might occur if we restarted during a migration or other change) + self.observed.locations.keys().for_each(|node_id| { + if Some(*node_id) != self.intent.attached { + self.intent.secondary.push(*node_id); + } + }); + } + + pub(crate) fn schedule(&mut self, scheduler: &mut Scheduler) -> Result<(), ScheduleError> { + // TODO: before scheduling new nodes, check if any existing content in + // self.intent refers to pageservers that are offline, and pick other + // pageservers if so. + + // Build the set of pageservers already in use by this tenant, to avoid scheduling + // more work on the same pageservers we're already using. + let mut used_pageservers = self.intent.all_pageservers(); + let mut modified = false; + + use PlacementPolicy::*; + match self.policy { + Single => { + // Should have exactly one attached, and zero secondaries + if self.intent.attached.is_none() { + let node_id = scheduler.schedule_shard(&used_pageservers)?; + self.intent.attached = Some(node_id); + used_pageservers.push(node_id); + modified = true; + } + if !self.intent.secondary.is_empty() { + self.intent.secondary.clear(); + modified = true; + } + } + Double(secondary_count) => { + // Should have exactly one attached, and N secondaries + if self.intent.attached.is_none() { + let node_id = scheduler.schedule_shard(&used_pageservers)?; + self.intent.attached = Some(node_id); + used_pageservers.push(node_id); + modified = true; + } + + while self.intent.secondary.len() < secondary_count { + let node_id = scheduler.schedule_shard(&used_pageservers)?; + self.intent.secondary.push(node_id); + used_pageservers.push(node_id); + modified = true; + } + } + } + + if modified { + self.sequence.0 += 1; + } + + Ok(()) + } + + fn dirty(&self) -> bool { + if let Some(node_id) = self.intent.attached { + let wanted_conf = attached_location_conf(self.generation, &self.shard, &self.config); + match self.observed.locations.get(&node_id) { + Some(conf) if conf.conf.as_ref() == Some(&wanted_conf) => {} + Some(_) | None => { + return true; + } + } + } + + for node_id in &self.intent.secondary { + let wanted_conf = secondary_location_conf(&self.shard, &self.config); + match self.observed.locations.get(node_id) { + Some(conf) if conf.conf.as_ref() == Some(&wanted_conf) => {} + Some(_) | None => { + return true; + } + } + } + + false + } + + pub(crate) fn maybe_reconcile( + &mut self, + result_tx: tokio::sync::mpsc::UnboundedSender, + pageservers: &Arc>, + compute_hook: &Arc, + service_config: &service::Config, + persistence: &Arc, + ) -> Option { + // If there are any ambiguous observed states, and the nodes they refer to are available, + // we should reconcile to clean them up. + let mut dirty_observed = false; + for (node_id, observed_loc) in &self.observed.locations { + let node = pageservers + .get(node_id) + .expect("Nodes may not be removed while referenced"); + if observed_loc.conf.is_none() + && !matches!(node.availability, NodeAvailability::Offline) + { + dirty_observed = true; + break; + } + } + + if !self.dirty() && !dirty_observed { + tracing::info!("Not dirty, no reconciliation needed."); + return None; + } + + // Reconcile already in flight for the current sequence? + if let Some(handle) = &self.reconciler { + if handle.sequence == self.sequence { + return Some(ReconcilerWaiter { + tenant_shard_id: self.tenant_shard_id, + seq_wait: self.waiter.clone(), + error_seq_wait: self.error_waiter.clone(), + error: self.last_error.clone(), + seq: self.sequence, + }); + } + } + + // Reconcile in flight for a stale sequence? Our sequence's task will wait for it before + // doing our sequence's work. + let old_handle = self.reconciler.take(); + + let cancel = CancellationToken::new(); + let mut reconciler = Reconciler { + tenant_shard_id: self.tenant_shard_id, + shard: self.shard, + generation: self.generation, + intent: self.intent.clone(), + config: self.config.clone(), + observed: self.observed.clone(), + pageservers: pageservers.clone(), + compute_hook: compute_hook.clone(), + service_config: service_config.clone(), + cancel: cancel.clone(), + persistence: persistence.clone(), + }; + + let reconcile_seq = self.sequence; + + tracing::info!("Spawning Reconciler for sequence {}", self.sequence); + let join_handle = tokio::task::spawn(async move { + // Wait for any previous reconcile task to complete before we start + if let Some(old_handle) = old_handle { + old_handle.cancel.cancel(); + if let Err(e) = old_handle.handle.await { + // We can't do much with this other than log it: the task is done, so + // we may proceed with our work. + tracing::error!("Unexpected join error waiting for reconcile task: {e}"); + } + } + + // Early check for cancellation before doing any work + // TODO: wrap all remote API operations in cancellation check + // as well. + if reconciler.cancel.is_cancelled() { + return; + } + + let result = reconciler.reconcile().await; + result_tx + .send(ReconcileResult { + sequence: reconcile_seq, + result, + tenant_shard_id: reconciler.tenant_shard_id, + generation: reconciler.generation, + observed: reconciler.observed, + }) + .ok(); + }); + + self.reconciler = Some(ReconcilerHandle { + sequence: self.sequence, + handle: join_handle, + cancel, + }); + + Some(ReconcilerWaiter { + tenant_shard_id: self.tenant_shard_id, + seq_wait: self.waiter.clone(), + error_seq_wait: self.error_waiter.clone(), + error: self.last_error.clone(), + seq: self.sequence, + }) + } +} diff --git a/control_plane/src/attachment_service.rs b/control_plane/src/attachment_service.rs index 731c05809e..0a353d8b12 100644 --- a/control_plane/src/attachment_service.rs +++ b/control_plane/src/attachment_service.rs @@ -1,14 +1,27 @@ use crate::{background_process, local_env::LocalEnv}; -use anyhow::anyhow; use camino::Utf8PathBuf; -use serde::{Deserialize, Serialize}; -use std::{path::PathBuf, process::Child}; -use utils::id::{NodeId, TenantId}; +use hyper::Method; +use pageserver_api::{ + models::{ShardParameters, TenantCreateRequest, TimelineCreateRequest, TimelineInfo}, + shard::TenantShardId, +}; +use pageserver_client::mgmt_api::ResponseErrorMessageExt; +use postgres_backend::AuthType; +use postgres_connection::parse_host_port; +use serde::{de::DeserializeOwned, Deserialize, Serialize}; +use std::{path::PathBuf, process::Child, str::FromStr}; +use tracing::instrument; +use utils::{ + auth::{Claims, Scope}, + id::{NodeId, TenantId}, +}; pub struct AttachmentService { env: LocalEnv, listen: String, path: PathBuf, + jwt_token: Option, + public_key_path: Option, client: reqwest::Client, } @@ -16,7 +29,7 @@ const COMMAND: &str = "attachment_service"; #[derive(Serialize, Deserialize)] pub struct AttachHookRequest { - pub tenant_id: TenantId, + pub tenant_shard_id: TenantShardId, pub node_id: Option, } @@ -27,7 +40,7 @@ pub struct AttachHookResponse { #[derive(Serialize, Deserialize)] pub struct InspectRequest { - pub tenant_id: TenantId, + pub tenant_shard_id: TenantShardId, } #[derive(Serialize, Deserialize)] @@ -35,6 +48,125 @@ pub struct InspectResponse { pub attachment: Option<(u32, NodeId)>, } +#[derive(Serialize, Deserialize)] +pub struct TenantCreateResponseShard { + pub node_id: NodeId, + pub generation: u32, +} + +#[derive(Serialize, Deserialize)] +pub struct TenantCreateResponse { + pub shards: Vec, +} + +#[derive(Serialize, Deserialize)] +pub struct NodeRegisterRequest { + pub node_id: NodeId, + + pub listen_pg_addr: String, + pub listen_pg_port: u16, + + pub listen_http_addr: String, + pub listen_http_port: u16, +} + +#[derive(Serialize, Deserialize)] +pub struct NodeConfigureRequest { + pub node_id: NodeId, + + pub availability: Option, + pub scheduling: Option, +} + +#[derive(Serialize, Deserialize, Debug)] +pub struct TenantLocateResponseShard { + pub shard_id: TenantShardId, + pub node_id: NodeId, + + pub listen_pg_addr: String, + pub listen_pg_port: u16, + + pub listen_http_addr: String, + pub listen_http_port: u16, +} + +#[derive(Serialize, Deserialize)] +pub struct TenantLocateResponse { + pub shards: Vec, + pub shard_params: ShardParameters, +} + +/// Explicitly migrating a particular shard is a low level operation +/// TODO: higher level "Reschedule tenant" operation where the request +/// specifies some constraints, e.g. asking it to get off particular node(s) +#[derive(Serialize, Deserialize, Debug)] +pub struct TenantShardMigrateRequest { + pub tenant_shard_id: TenantShardId, + pub node_id: NodeId, +} + +#[derive(Serialize, Deserialize, Clone, Copy)] +pub enum NodeAvailability { + // Normal, happy state + Active, + // Offline: Tenants shouldn't try to attach here, but they may assume that their + // secondary locations on this node still exist. Newly added nodes are in this + // state until we successfully contact them. + Offline, +} + +impl FromStr for NodeAvailability { + type Err = anyhow::Error; + + fn from_str(s: &str) -> Result { + match s { + "active" => Ok(Self::Active), + "offline" => Ok(Self::Offline), + _ => Err(anyhow::anyhow!("Unknown availability state '{s}'")), + } + } +} + +/// FIXME: this is a duplicate of the type in the attachment_service crate, because the +/// type needs to be defined with diesel traits in there. +#[derive(Serialize, Deserialize, Clone, Copy)] +pub enum NodeSchedulingPolicy { + Active, + Filling, + Pause, + Draining, +} + +impl FromStr for NodeSchedulingPolicy { + type Err = anyhow::Error; + + fn from_str(s: &str) -> Result { + match s { + "active" => Ok(Self::Active), + "filling" => Ok(Self::Filling), + "pause" => Ok(Self::Pause), + "draining" => Ok(Self::Draining), + _ => Err(anyhow::anyhow!("Unknown scheduling state '{s}'")), + } + } +} + +impl From for String { + fn from(value: NodeSchedulingPolicy) -> String { + use NodeSchedulingPolicy::*; + match value { + Active => "active", + Filling => "filling", + Pause => "pause", + Draining => "draining", + } + .to_string() + } +} + +#[derive(Serialize, Deserialize, Debug)] +pub struct TenantShardMigrateResponse {} + impl AttachmentService { pub fn from_env(env: &LocalEnv) -> Self { let path = env.base_data_dir.join("attachments.json"); @@ -49,10 +181,34 @@ impl AttachmentService { listen_url.port().unwrap() ); + // Assume all pageservers have symmetric auth configuration: this service + // expects to use one JWT token to talk to all of them. + let ps_conf = env + .pageservers + .first() + .expect("Config is validated to contain at least one pageserver"); + let (jwt_token, public_key_path) = match ps_conf.http_auth_type { + AuthType::Trust => (None, None), + AuthType::NeonJWT => { + let jwt_token = env + .generate_auth_token(&Claims::new(None, Scope::PageServerApi)) + .unwrap(); + + // If pageserver auth is enabled, this implicitly enables auth for this service, + // using the same credentials. + let public_key_path = + camino::Utf8PathBuf::try_from(env.base_data_dir.join("auth_public_key.pem")) + .unwrap(); + (Some(jwt_token), Some(public_key_path)) + } + }; + Self { env: env.clone(), path, listen, + jwt_token, + public_key_path, client: reqwest::ClientBuilder::new() .build() .expect("Failed to construct http client"), @@ -67,72 +223,199 @@ impl AttachmentService { pub async fn start(&self) -> anyhow::Result { let path_str = self.path.to_string_lossy(); - background_process::start_process( + let mut args = vec!["-l", &self.listen, "-p", &path_str] + .into_iter() + .map(|s| s.to_string()) + .collect::>(); + if let Some(jwt_token) = &self.jwt_token { + args.push(format!("--jwt-token={jwt_token}")); + } + + if let Some(public_key_path) = &self.public_key_path { + args.push(format!("--public-key={public_key_path}")); + } + + let result = background_process::start_process( COMMAND, &self.env.base_data_dir, &self.env.attachment_service_bin(), - ["-l", &self.listen, "-p", &path_str], - [], + args, + [( + "NEON_REPO_DIR".to_string(), + self.env.base_data_dir.to_string_lossy().to_string(), + )], background_process::InitialPidFile::Create(self.pid_file()), - // TODO: a real status check - || async move { anyhow::Ok(true) }, + || async { + match self.status().await { + Ok(_) => Ok(true), + Err(_) => Ok(false), + } + }, ) - .await + .await; + + for ps_conf in &self.env.pageservers { + let (pg_host, pg_port) = + parse_host_port(&ps_conf.listen_pg_addr).expect("Unable to parse listen_pg_addr"); + let (http_host, http_port) = parse_host_port(&ps_conf.listen_http_addr) + .expect("Unable to parse listen_http_addr"); + self.node_register(NodeRegisterRequest { + node_id: ps_conf.id, + listen_pg_addr: pg_host.to_string(), + listen_pg_port: pg_port.unwrap_or(5432), + listen_http_addr: http_host.to_string(), + listen_http_port: http_port.unwrap_or(80), + }) + .await?; + } + + result } pub fn stop(&self, immediate: bool) -> anyhow::Result<()> { background_process::stop_process(immediate, COMMAND, &self.pid_file()) } - - /// Call into the attach_hook API, for use before handing out attachments to pageservers - pub async fn attach_hook( + /// Simple HTTP request wrapper for calling into attachment service + async fn dispatch( &self, - tenant_id: TenantId, - pageserver_id: NodeId, - ) -> anyhow::Result> { - use hyper::StatusCode; - + method: hyper::Method, + path: String, + body: Option, + ) -> anyhow::Result + where + RQ: Serialize + Sized, + RS: DeserializeOwned + Sized, + { let url = self .env .control_plane_api .clone() .unwrap() - .join("attach-hook") + .join(&path) .unwrap(); + let mut builder = self.client.request(method, url); + if let Some(body) = body { + builder = builder.json(&body) + } + if let Some(jwt_token) = &self.jwt_token { + builder = builder.header( + reqwest::header::AUTHORIZATION, + format!("Bearer {jwt_token}"), + ); + } + + let response = builder.send().await?; + let response = response.error_from_body().await?; + + Ok(response + .json() + .await + .map_err(pageserver_client::mgmt_api::Error::ReceiveBody)?) + } + + /// Call into the attach_hook API, for use before handing out attachments to pageservers + #[instrument(skip(self))] + pub async fn attach_hook( + &self, + tenant_shard_id: TenantShardId, + pageserver_id: NodeId, + ) -> anyhow::Result> { let request = AttachHookRequest { - tenant_id, + tenant_shard_id, node_id: Some(pageserver_id), }; - let response = self.client.post(url).json(&request).send().await?; - if response.status() != StatusCode::OK { - return Err(anyhow!("Unexpected status {}", response.status())); - } + let response = self + .dispatch::<_, AttachHookResponse>( + Method::POST, + "attach-hook".to_string(), + Some(request), + ) + .await?; - let response = response.json::().await?; Ok(response.gen) } - pub async fn inspect(&self, tenant_id: TenantId) -> anyhow::Result> { - use hyper::StatusCode; + #[instrument(skip(self))] + pub async fn inspect( + &self, + tenant_shard_id: TenantShardId, + ) -> anyhow::Result> { + let request = InspectRequest { tenant_shard_id }; - let url = self - .env - .control_plane_api - .clone() - .unwrap() - .join("inspect") - .unwrap(); + let response = self + .dispatch::<_, InspectResponse>(Method::POST, "inspect".to_string(), Some(request)) + .await?; - let request = InspectRequest { tenant_id }; - - let response = self.client.post(url).json(&request).send().await?; - if response.status() != StatusCode::OK { - return Err(anyhow!("Unexpected status {}", response.status())); - } - - let response = response.json::().await?; Ok(response.attachment) } + + #[instrument(skip(self))] + pub async fn tenant_create( + &self, + req: TenantCreateRequest, + ) -> anyhow::Result { + self.dispatch(Method::POST, "tenant".to_string(), Some(req)) + .await + } + + #[instrument(skip(self))] + pub async fn tenant_locate(&self, tenant_id: TenantId) -> anyhow::Result { + self.dispatch::<(), _>(Method::GET, format!("tenant/{tenant_id}/locate"), None) + .await + } + + #[instrument(skip(self))] + pub async fn tenant_migrate( + &self, + tenant_shard_id: TenantShardId, + node_id: NodeId, + ) -> anyhow::Result { + self.dispatch( + Method::PUT, + format!("tenant/{tenant_shard_id}/migrate"), + Some(TenantShardMigrateRequest { + tenant_shard_id, + node_id, + }), + ) + .await + } + + #[instrument(skip_all, fields(node_id=%req.node_id))] + pub async fn node_register(&self, req: NodeRegisterRequest) -> anyhow::Result<()> { + self.dispatch::<_, ()>(Method::POST, "node".to_string(), Some(req)) + .await + } + + #[instrument(skip_all, fields(node_id=%req.node_id))] + pub async fn node_configure(&self, req: NodeConfigureRequest) -> anyhow::Result<()> { + self.dispatch::<_, ()>( + Method::PUT, + format!("node/{}/config", req.node_id), + Some(req), + ) + .await + } + + #[instrument(skip(self))] + pub async fn status(&self) -> anyhow::Result<()> { + self.dispatch::<(), ()>(Method::GET, "status".to_string(), None) + .await + } + + #[instrument(skip_all, fields(%tenant_id, timeline_id=%req.new_timeline_id))] + pub async fn tenant_timeline_create( + &self, + tenant_id: TenantId, + req: TimelineCreateRequest, + ) -> anyhow::Result { + self.dispatch( + Method::POST, + format!("tenant/{tenant_id}/timeline"), + Some(req), + ) + .await + } } diff --git a/control_plane/src/bin/attachment_service.rs b/control_plane/src/bin/attachment_service.rs deleted file mode 100644 index e50c8fbba0..0000000000 --- a/control_plane/src/bin/attachment_service.rs +++ /dev/null @@ -1,337 +0,0 @@ -/// The attachment service mimics the aspects of the control plane API -/// that are required for a pageserver to operate. -/// -/// This enables running & testing pageservers without a full-blown -/// deployment of the Neon cloud platform. -/// -use anyhow::anyhow; -use clap::Parser; -use hex::FromHex; -use hyper::StatusCode; -use hyper::{Body, Request, Response}; -use pageserver_api::shard::TenantShardId; -use serde::{Deserialize, Serialize}; -use std::path::{Path, PathBuf}; -use std::{collections::HashMap, sync::Arc}; -use utils::http::endpoint::request_span; -use utils::logging::{self, LogFormat}; -use utils::signals::{ShutdownSignals, Signal}; - -use utils::{ - http::{ - endpoint::{self}, - error::ApiError, - json::{json_request, json_response}, - RequestExt, RouterBuilder, - }, - id::{NodeId, TenantId}, - tcp_listener, -}; - -use pageserver_api::control_api::{ - ReAttachRequest, ReAttachResponse, ReAttachResponseTenant, ValidateRequest, ValidateResponse, - ValidateResponseTenant, -}; - -use control_plane::attachment_service::{ - AttachHookRequest, AttachHookResponse, InspectRequest, InspectResponse, -}; - -#[derive(Parser)] -#[command(author, version, about, long_about = None)] -#[command(arg_required_else_help(true))] -struct Cli { - /// Host and port to listen on, like `127.0.0.1:1234` - #[arg(short, long)] - listen: std::net::SocketAddr, - - /// Path to the .json file to store state (will be created if it doesn't exist) - #[arg(short, long)] - path: PathBuf, -} - -// The persistent state of each Tenant -#[derive(Serialize, Deserialize, Clone)] -struct TenantState { - // Currently attached pageserver - pageserver: Option, - - // Latest generation number: next time we attach, increment this - // and use the incremented number when attaching - generation: u32, -} - -fn to_hex_map(input: &HashMap, serializer: S) -> Result -where - S: serde::Serializer, - V: Clone + Serialize, -{ - let transformed = input.iter().map(|(k, v)| (hex::encode(k), v.clone())); - - transformed - .collect::>() - .serialize(serializer) -} - -fn from_hex_map<'de, D, V>(deserializer: D) -> Result, D::Error> -where - D: serde::de::Deserializer<'de>, - V: Deserialize<'de>, -{ - let hex_map = HashMap::::deserialize(deserializer)?; - hex_map - .into_iter() - .map(|(k, v)| { - TenantId::from_hex(k) - .map(|k| (k, v)) - .map_err(serde::de::Error::custom) - }) - .collect() -} - -// Top level state available to all HTTP handlers -#[derive(Serialize, Deserialize)] -struct PersistentState { - #[serde(serialize_with = "to_hex_map", deserialize_with = "from_hex_map")] - tenants: HashMap, - - #[serde(skip)] - path: PathBuf, -} - -impl PersistentState { - async fn save(&self) -> anyhow::Result<()> { - let bytes = serde_json::to_vec(self)?; - tokio::fs::write(&self.path, &bytes).await?; - - Ok(()) - } - - async fn load(path: &Path) -> anyhow::Result { - let bytes = tokio::fs::read(path).await?; - let mut decoded = serde_json::from_slice::(&bytes)?; - decoded.path = path.to_owned(); - Ok(decoded) - } - - async fn load_or_new(path: &Path) -> Self { - match Self::load(path).await { - Ok(s) => { - tracing::info!("Loaded state file at {}", path.display()); - s - } - Err(e) - if e.downcast_ref::() - .map(|e| e.kind() == std::io::ErrorKind::NotFound) - .unwrap_or(false) => - { - tracing::info!("Will create state file at {}", path.display()); - Self { - tenants: HashMap::new(), - path: path.to_owned(), - } - } - Err(e) => { - panic!("Failed to load state from '{}': {e:#} (maybe your .neon/ dir was written by an older version?)", path.display()) - } - } - } -} - -/// State available to HTTP request handlers -#[derive(Clone)] -struct State { - inner: Arc>, -} - -impl State { - fn new(persistent_state: PersistentState) -> State { - Self { - inner: Arc::new(tokio::sync::RwLock::new(persistent_state)), - } - } -} - -#[inline(always)] -fn get_state(request: &Request) -> &State { - request - .data::>() - .expect("unknown state type") - .as_ref() -} - -/// Pageserver calls into this on startup, to learn which tenants it should attach -async fn handle_re_attach(mut req: Request) -> Result, ApiError> { - let reattach_req = json_request::(&mut req).await?; - - let state = get_state(&req).inner.clone(); - let mut locked = state.write().await; - - let mut response = ReAttachResponse { - tenants: Vec::new(), - }; - for (t, state) in &mut locked.tenants { - if state.pageserver == Some(reattach_req.node_id) { - state.generation += 1; - response.tenants.push(ReAttachResponseTenant { - // TODO(sharding): make this shard-aware - id: TenantShardId::unsharded(*t), - gen: state.generation, - }); - } - } - - locked.save().await.map_err(ApiError::InternalServerError)?; - - json_response(StatusCode::OK, response) -} - -/// Pageserver calls into this before doing deletions, to confirm that it still -/// holds the latest generation for the tenants with deletions enqueued -async fn handle_validate(mut req: Request) -> Result, ApiError> { - let validate_req = json_request::(&mut req).await?; - - let locked = get_state(&req).inner.read().await; - - let mut response = ValidateResponse { - tenants: Vec::new(), - }; - - for req_tenant in validate_req.tenants { - // TODO(sharding): make this shard-aware - if let Some(tenant_state) = locked.tenants.get(&req_tenant.id.tenant_id) { - let valid = tenant_state.generation == req_tenant.gen; - tracing::info!( - "handle_validate: {}(gen {}): valid={valid} (latest {})", - req_tenant.id, - req_tenant.gen, - tenant_state.generation - ); - response.tenants.push(ValidateResponseTenant { - id: req_tenant.id, - valid, - }); - } - } - - json_response(StatusCode::OK, response) -} -/// Call into this before attaching a tenant to a pageserver, to acquire a generation number -/// (in the real control plane this is unnecessary, because the same program is managing -/// generation numbers and doing attachments). -async fn handle_attach_hook(mut req: Request) -> Result, ApiError> { - let attach_req = json_request::(&mut req).await?; - - let state = get_state(&req).inner.clone(); - let mut locked = state.write().await; - - let tenant_state = locked - .tenants - .entry(attach_req.tenant_id) - .or_insert_with(|| TenantState { - pageserver: attach_req.node_id, - generation: 0, - }); - - if let Some(attaching_pageserver) = attach_req.node_id.as_ref() { - tenant_state.generation += 1; - tracing::info!( - tenant_id = %attach_req.tenant_id, - ps_id = %attaching_pageserver, - generation = %tenant_state.generation, - "issuing", - ); - } else if let Some(ps_id) = tenant_state.pageserver { - tracing::info!( - tenant_id = %attach_req.tenant_id, - %ps_id, - generation = %tenant_state.generation, - "dropping", - ); - } else { - tracing::info!( - tenant_id = %attach_req.tenant_id, - "no-op: tenant already has no pageserver"); - } - tenant_state.pageserver = attach_req.node_id; - let generation = tenant_state.generation; - - tracing::info!( - "handle_attach_hook: tenant {} set generation {}, pageserver {}", - attach_req.tenant_id, - tenant_state.generation, - attach_req.node_id.unwrap_or(utils::id::NodeId(0xfffffff)) - ); - - locked.save().await.map_err(ApiError::InternalServerError)?; - - json_response( - StatusCode::OK, - AttachHookResponse { - gen: attach_req.node_id.map(|_| generation), - }, - ) -} - -async fn handle_inspect(mut req: Request) -> Result, ApiError> { - let inspect_req = json_request::(&mut req).await?; - - let state = get_state(&req).inner.clone(); - let locked = state.write().await; - let tenant_state = locked.tenants.get(&inspect_req.tenant_id); - - json_response( - StatusCode::OK, - InspectResponse { - attachment: tenant_state.and_then(|s| s.pageserver.map(|ps| (s.generation, ps))), - }, - ) -} - -fn make_router(persistent_state: PersistentState) -> RouterBuilder { - endpoint::make_router() - .data(Arc::new(State::new(persistent_state))) - .post("/re-attach", |r| request_span(r, handle_re_attach)) - .post("/validate", |r| request_span(r, handle_validate)) - .post("/attach-hook", |r| request_span(r, handle_attach_hook)) - .post("/inspect", |r| request_span(r, handle_inspect)) -} - -#[tokio::main] -async fn main() -> anyhow::Result<()> { - logging::init( - LogFormat::Plain, - logging::TracingErrorLayerEnablement::Disabled, - logging::Output::Stdout, - )?; - - let args = Cli::parse(); - tracing::info!( - "Starting, state at {}, listening on {}", - args.path.to_string_lossy(), - args.listen - ); - - let persistent_state = PersistentState::load_or_new(&args.path).await; - - let http_listener = tcp_listener::bind(args.listen)?; - let router = make_router(persistent_state) - .build() - .map_err(|err| anyhow!(err))?; - let service = utils::http::RouterService::new(router).unwrap(); - let server = hyper::Server::from_tcp(http_listener)?.serve(service); - - tracing::info!("Serving on {0}", args.listen); - - tokio::task::spawn(server); - - ShutdownSignals::handle(|signal| match signal { - Signal::Interrupt | Signal::Terminate | Signal::Quit => { - tracing::info!("Got {}. Terminating", signal.name()); - // We're just a test helper: no graceful shutdown. - std::process::exit(0); - } - })?; - - Ok(()) -} diff --git a/control_plane/src/bin/neon_local.rs b/control_plane/src/bin/neon_local.rs index dd09d47655..279c47398f 100644 --- a/control_plane/src/bin/neon_local.rs +++ b/control_plane/src/bin/neon_local.rs @@ -8,19 +8,24 @@ use anyhow::{anyhow, bail, Context, Result}; use clap::{value_parser, Arg, ArgAction, ArgMatches, Command, ValueEnum}; use compute_api::spec::ComputeMode; -use control_plane::attachment_service::AttachmentService; +use control_plane::attachment_service::{ + AttachmentService, NodeAvailability, NodeConfigureRequest, NodeSchedulingPolicy, +}; use control_plane::endpoint::ComputeControlPlane; use control_plane::local_env::{InitForceMode, LocalEnv}; use control_plane::pageserver::{PageServerNode, PAGESERVER_REMOTE_STORAGE_DIR}; use control_plane::safekeeper::SafekeeperNode; -use control_plane::tenant_migration::migrate_tenant; use control_plane::{broker, local_env}; -use pageserver_api::models::TimelineInfo; +use pageserver_api::models::{ + ShardParameters, TenantCreateRequest, TimelineCreateRequest, TimelineInfo, +}; +use pageserver_api::shard::{ShardCount, ShardStripeSize, TenantShardId}; use pageserver_api::{ DEFAULT_HTTP_LISTEN_PORT as DEFAULT_PAGESERVER_HTTP_PORT, DEFAULT_PG_LISTEN_PORT as DEFAULT_PAGESERVER_PG_PORT, }; use postgres_backend::AuthType; +use postgres_connection::parse_host_port; use safekeeper_api::{ DEFAULT_HTTP_LISTEN_PORT as DEFAULT_SAFEKEEPER_HTTP_PORT, DEFAULT_PG_LISTEN_PORT as DEFAULT_SAFEKEEPER_PG_PORT, @@ -30,6 +35,7 @@ use std::path::PathBuf; use std::process::exit; use std::str::FromStr; use storage_broker::DEFAULT_LISTEN_ADDR as DEFAULT_BROKER_ADDR; +use url::Host; use utils::{ auth::{Claims, Scope}, id::{NodeId, TenantId, TenantTimelineId, TimelineId}, @@ -276,10 +282,10 @@ fn print_timeline( /// Connects to the pageserver to query this information. async fn get_timeline_infos( env: &local_env::LocalEnv, - tenant_id: &TenantId, + tenant_shard_id: &TenantShardId, ) -> Result> { Ok(get_default_pageserver(env) - .timeline_list(tenant_id) + .timeline_list(tenant_shard_id) .await? .into_iter() .map(|timeline_info| (timeline_info.timeline_id, timeline_info)) @@ -297,6 +303,20 @@ fn get_tenant_id(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> anyhow::R } } +// Helper function to parse --tenant_id option, for commands that accept a shard suffix +fn get_tenant_shard_id( + sub_match: &ArgMatches, + env: &local_env::LocalEnv, +) -> anyhow::Result { + if let Some(tenant_id_from_arguments) = parse_tenant_shard_id(sub_match).transpose() { + tenant_id_from_arguments + } else if let Some(default_id) = env.default_tenant_id { + Ok(TenantShardId::unsharded(default_id)) + } else { + anyhow::bail!("No tenant shard id. Use --tenant-id, or set a default tenant"); + } +} + fn parse_tenant_id(sub_match: &ArgMatches) -> anyhow::Result> { sub_match .get_one::("tenant-id") @@ -305,6 +325,14 @@ fn parse_tenant_id(sub_match: &ArgMatches) -> anyhow::Result> { .context("Failed to parse tenant id from the argument string") } +fn parse_tenant_shard_id(sub_match: &ArgMatches) -> anyhow::Result> { + sub_match + .get_one::("tenant-id") + .map(|id_str| TenantShardId::from_str(id_str)) + .transpose() + .context("Failed to parse tenant shard id from the argument string") +} + fn parse_timeline_id(sub_match: &ArgMatches) -> anyhow::Result> { sub_match .get_one::("timeline-id") @@ -393,47 +421,68 @@ async fn handle_tenant( Some(("create", create_match)) => { let tenant_conf: HashMap<_, _> = create_match .get_many::("config") - .map(|vals| vals.flat_map(|c| c.split_once(':')).collect()) + .map(|vals: clap::parser::ValuesRef<'_, String>| { + vals.flat_map(|c| c.split_once(':')).collect() + }) .unwrap_or_default(); + let shard_count: u8 = create_match + .get_one::("shard-count") + .cloned() + .unwrap_or(0); + + let shard_stripe_size: Option = + create_match.get_one::("shard-stripe-size").cloned(); + + let tenant_conf = PageServerNode::parse_config(tenant_conf)?; + // If tenant ID was not specified, generate one let tenant_id = parse_tenant_id(create_match)?.unwrap_or_else(TenantId::generate); - let generation = if env.control_plane_api.is_some() { - // We must register the tenant with the attachment service, so - // that when the pageserver restarts, it will be re-attached. - let attachment_service = AttachmentService::from_env(env); - attachment_service - .attach_hook(tenant_id, pageserver.conf.id) - .await? - } else { - None - }; - - pageserver - .tenant_create(tenant_id, generation, tenant_conf) + // We must register the tenant with the attachment service, so + // that when the pageserver restarts, it will be re-attached. + let attachment_service = AttachmentService::from_env(env); + attachment_service + .tenant_create(TenantCreateRequest { + // Note that ::unsharded here isn't actually because the tenant is unsharded, its because the + // attachment service expecfs a shard-naive tenant_id in this attribute, and the TenantCreateRequest + // type is used both in attachment service (for creating tenants) and in pageserver (for creating shards) + new_tenant_id: TenantShardId::unsharded(tenant_id), + generation: None, + shard_parameters: ShardParameters { + count: ShardCount(shard_count), + stripe_size: shard_stripe_size + .map(ShardStripeSize) + .unwrap_or(ShardParameters::DEFAULT_STRIPE_SIZE), + }, + config: tenant_conf, + }) .await?; println!("tenant {tenant_id} successfully created on the pageserver"); // Create an initial timeline for the new tenant - let new_timeline_id = parse_timeline_id(create_match)?; + let new_timeline_id = + parse_timeline_id(create_match)?.unwrap_or(TimelineId::generate()); let pg_version = create_match .get_one::("pg-version") .copied() .context("Failed to parse postgres version from the argument string")?; - let timeline_info = pageserver - .timeline_create( + // FIXME: passing None for ancestor_start_lsn is not kosher in a sharded world: we can't have + // different shards picking different start lsns. Maybe we have to teach attachment service + // to let shard 0 branch first and then propagate the chosen LSN to other shards. + attachment_service + .tenant_timeline_create( tenant_id, - new_timeline_id, - None, - None, - Some(pg_version), - None, + TimelineCreateRequest { + new_timeline_id, + ancestor_timeline_id: None, + ancestor_start_lsn: None, + existing_initdb_timeline_id: None, + pg_version: Some(pg_version), + }, ) .await?; - let new_timeline_id = timeline_info.timeline_id; - let last_record_lsn = timeline_info.last_record_lsn; env.register_branch_mapping( DEFAULT_BRANCH_NAME.to_string(), @@ -441,9 +490,7 @@ async fn handle_tenant( new_timeline_id, )?; - println!( - "Created an initial timeline '{new_timeline_id}' at Lsn {last_record_lsn} for tenant: {tenant_id}", - ); + println!("Created an initial timeline '{new_timeline_id}' for tenant: {tenant_id}",); if create_match.get_flag("set-default") { println!("Setting tenant {tenant_id} as a default one"); @@ -470,14 +517,64 @@ async fn handle_tenant( println!("tenant {tenant_id} successfully configured on the pageserver"); } Some(("migrate", matches)) => { - let tenant_id = get_tenant_id(matches, env)?; + let tenant_shard_id = get_tenant_shard_id(matches, env)?; let new_pageserver = get_pageserver(env, matches)?; let new_pageserver_id = new_pageserver.conf.id; - migrate_tenant(env, tenant_id, new_pageserver).await?; - println!("tenant {tenant_id} migrated to {}", new_pageserver_id); - } + let attachment_service = AttachmentService::from_env(env); + attachment_service + .tenant_migrate(tenant_shard_id, new_pageserver_id) + .await?; + println!("tenant {tenant_shard_id} migrated to {}", new_pageserver_id); + } + Some(("status", matches)) => { + let tenant_id = get_tenant_id(matches, env)?; + + let mut shard_table = comfy_table::Table::new(); + shard_table.set_header(["Shard", "Pageserver", "Physical Size"]); + + let mut tenant_synthetic_size = None; + + let attachment_service = AttachmentService::from_env(env); + for shard in attachment_service.tenant_locate(tenant_id).await?.shards { + let pageserver = + PageServerNode::from_env(env, env.get_pageserver_conf(shard.node_id)?); + + let size = pageserver + .http_client + .tenant_details(shard.shard_id) + .await? + .tenant_info + .current_physical_size + .unwrap(); + + shard_table.add_row([ + format!("{}", shard.shard_id.shard_slug()), + format!("{}", shard.node_id.0), + format!("{} MiB", size / (1024 * 1024)), + ]); + + if shard.shard_id.is_zero() { + tenant_synthetic_size = + Some(pageserver.tenant_synthetic_size(shard.shard_id).await?); + } + } + + let Some(synthetic_size) = tenant_synthetic_size else { + bail!("Shard 0 not found") + }; + + let mut tenant_table = comfy_table::Table::new(); + tenant_table.add_row(["Tenant ID".to_string(), tenant_id.to_string()]); + tenant_table.add_row([ + "Synthetic size".to_string(), + format!("{} MiB", synthetic_size.size.unwrap_or(0) / (1024 * 1024)), + ]); + + println!("{tenant_table}"); + println!("{shard_table}"); + } Some((sub_name, _)) => bail!("Unexpected tenant subcommand '{}'", sub_name), None => bail!("no tenant subcommand provided"), } @@ -489,8 +586,10 @@ async fn handle_timeline(timeline_match: &ArgMatches, env: &mut local_env::Local match timeline_match.subcommand() { Some(("list", list_match)) => { - let tenant_id = get_tenant_id(list_match, env)?; - let timelines = pageserver.timeline_list(&tenant_id).await?; + // TODO(sharding): this command shouldn't have to specify a shard ID: we should ask the attachment service + // where shard 0 is attached, and query there. + let tenant_shard_id = get_tenant_shard_id(list_match, env)?; + let timelines = pageserver.timeline_list(&tenant_shard_id).await?; print_timelines_tree(timelines, env.timeline_name_mappings())?; } Some(("create", create_match)) => { @@ -505,18 +604,19 @@ async fn handle_timeline(timeline_match: &ArgMatches, env: &mut local_env::Local .context("Failed to parse postgres version from the argument string")?; let new_timeline_id_opt = parse_timeline_id(create_match)?; + let new_timeline_id = new_timeline_id_opt.unwrap_or(TimelineId::generate()); - let timeline_info = pageserver - .timeline_create( - tenant_id, - new_timeline_id_opt, - None, - None, - Some(pg_version), - None, - ) + let attachment_service = AttachmentService::from_env(env); + let create_req = TimelineCreateRequest { + new_timeline_id, + ancestor_timeline_id: None, + existing_initdb_timeline_id: None, + ancestor_start_lsn: None, + pg_version: Some(pg_version), + }; + let timeline_info = attachment_service + .tenant_timeline_create(tenant_id, create_req) .await?; - let new_timeline_id = timeline_info.timeline_id; let last_record_lsn = timeline_info.last_record_lsn; env.register_branch_mapping(new_branch_name.to_string(), tenant_id, new_timeline_id)?; @@ -574,7 +674,6 @@ async fn handle_timeline(timeline_match: &ArgMatches, env: &mut local_env::Local None, pg_version, ComputeMode::Primary, - DEFAULT_PAGESERVER_ID, )?; println!("Done"); } @@ -598,17 +697,18 @@ async fn handle_timeline(timeline_match: &ArgMatches, env: &mut local_env::Local .map(|lsn_str| Lsn::from_str(lsn_str)) .transpose() .context("Failed to parse ancestor start Lsn from the request")?; - let timeline_info = pageserver - .timeline_create( - tenant_id, - None, - start_lsn, - Some(ancestor_timeline_id), - None, - None, - ) + let new_timeline_id = TimelineId::generate(); + let attachment_service = AttachmentService::from_env(env); + let create_req = TimelineCreateRequest { + new_timeline_id, + ancestor_timeline_id: Some(ancestor_timeline_id), + existing_initdb_timeline_id: None, + ancestor_start_lsn: start_lsn, + pg_version: None, + }; + let timeline_info = attachment_service + .tenant_timeline_create(tenant_id, create_req) .await?; - let new_timeline_id = timeline_info.timeline_id; let last_record_lsn = timeline_info.last_record_lsn; @@ -635,8 +735,10 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re match sub_name { "list" => { - let tenant_id = get_tenant_id(sub_args, env)?; - let timeline_infos = get_timeline_infos(env, &tenant_id) + // TODO(sharding): this command shouldn't have to specify a shard ID: we should ask the attachment service + // where shard 0 is attached, and query there. + let tenant_shard_id = get_tenant_shard_id(sub_args, env)?; + let timeline_infos = get_timeline_infos(env, &tenant_shard_id) .await .unwrap_or_else(|e| { eprintln!("Failed to load timeline info: {}", e); @@ -661,7 +763,7 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re for (endpoint_id, endpoint) in cplane .endpoints .iter() - .filter(|(_, endpoint)| endpoint.tenant_id == tenant_id) + .filter(|(_, endpoint)| endpoint.tenant_id == tenant_shard_id.tenant_id) { let lsn_str = match endpoint.mode { ComputeMode::Static(lsn) => { @@ -680,7 +782,10 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re }; let branch_name = timeline_name_mappings - .get(&TenantTimelineId::new(tenant_id, endpoint.timeline_id)) + .get(&TenantTimelineId::new( + tenant_shard_id.tenant_id, + endpoint.timeline_id, + )) .map(|name| name.as_str()) .unwrap_or("?"); @@ -728,13 +833,6 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re .copied() .unwrap_or(false); - let pageserver_id = - if let Some(id_str) = sub_args.get_one::("endpoint-pageserver-id") { - NodeId(id_str.parse().context("while parsing pageserver id")?) - } else { - DEFAULT_PAGESERVER_ID - }; - let mode = match (lsn, hot_standby) { (Some(lsn), false) => ComputeMode::Static(lsn), (None, true) => ComputeMode::Replica, @@ -762,7 +860,6 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re http_port, pg_version, mode, - pageserver_id, )?; } "start" => { @@ -772,9 +869,11 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re let pageserver_id = if let Some(id_str) = sub_args.get_one::("endpoint-pageserver-id") { - NodeId(id_str.parse().context("while parsing pageserver id")?) + Some(NodeId( + id_str.parse().context("while parsing pageserver id")?, + )) } else { - DEFAULT_PAGESERVER_ID + None }; let remote_ext_config = sub_args.get_one::("remote-ext-config"); @@ -805,7 +904,38 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re endpoint.timeline_id, )?; - let ps_conf = env.get_pageserver_conf(pageserver_id)?; + let (pageservers, stripe_size) = if let Some(pageserver_id) = pageserver_id { + let conf = env.get_pageserver_conf(pageserver_id).unwrap(); + let parsed = parse_host_port(&conf.listen_pg_addr).expect("Bad config"); + ( + vec![(parsed.0, parsed.1.unwrap_or(5432))], + // If caller is telling us what pageserver to use, this is not a tenant which is + // full managed by attachment service, therefore not sharded. + ShardParameters::DEFAULT_STRIPE_SIZE, + ) + } else { + // Look up the currently attached location of the tenant, and its striping metadata, + // to pass these on to postgres. + let attachment_service = AttachmentService::from_env(env); + let locate_result = attachment_service.tenant_locate(endpoint.tenant_id).await?; + let pageservers = locate_result + .shards + .into_iter() + .map(|shard| { + ( + Host::parse(&shard.listen_pg_addr) + .expect("Attachment service reported bad hostname"), + shard.listen_pg_port, + ) + }) + .collect::>(); + let stripe_size = locate_result.shard_params.stripe_size; + + (pageservers, stripe_size) + }; + assert!(!pageservers.is_empty()); + + let ps_conf = env.get_pageserver_conf(DEFAULT_PAGESERVER_ID)?; let auth_token = if matches!(ps_conf.pg_auth_type, AuthType::NeonJWT) { let claims = Claims::new(Some(endpoint.tenant_id), Scope::Tenant); @@ -816,7 +946,13 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re println!("Starting existing endpoint {endpoint_id}..."); endpoint - .start(&auth_token, safekeepers, remote_ext_config) + .start( + &auth_token, + safekeepers, + pageservers, + remote_ext_config, + stripe_size.0 as usize, + ) .await?; } "reconfigure" => { @@ -827,15 +963,31 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re .endpoints .get(endpoint_id.as_str()) .with_context(|| format!("postgres endpoint {endpoint_id} is not found"))?; - let pageserver_id = + let pageservers = if let Some(id_str) = sub_args.get_one::("endpoint-pageserver-id") { - Some(NodeId( - id_str.parse().context("while parsing pageserver id")?, - )) + let ps_id = NodeId(id_str.parse().context("while parsing pageserver id")?); + let pageserver = PageServerNode::from_env(env, env.get_pageserver_conf(ps_id)?); + vec![( + pageserver.pg_connection_config.host().clone(), + pageserver.pg_connection_config.port(), + )] } else { - None + let attachment_service = AttachmentService::from_env(env); + attachment_service + .tenant_locate(endpoint.tenant_id) + .await? + .shards + .into_iter() + .map(|shard| { + ( + Host::parse(&shard.listen_pg_addr) + .expect("Attachment service reported malformed host"), + shard.listen_pg_port, + ) + }) + .collect::>() }; - endpoint.reconfigure(pageserver_id).await?; + endpoint.reconfigure(pageservers).await?; } "stop" => { let endpoint_id = sub_args @@ -959,6 +1111,21 @@ async fn handle_pageserver(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> } } + Some(("set-state", subcommand_args)) => { + let pageserver = get_pageserver(env, subcommand_args)?; + let scheduling = subcommand_args.get_one("scheduling"); + let availability = subcommand_args.get_one("availability"); + + let attachment_service = AttachmentService::from_env(env); + attachment_service + .node_configure(NodeConfigureRequest { + node_id: pageserver.conf.id, + scheduling: scheduling.cloned(), + availability: availability.cloned(), + }) + .await?; + } + Some(("status", subcommand_args)) => { match get_pageserver(env, subcommand_args)?.check_status().await { Ok(_) => println!("Page server is up and running"), @@ -1358,6 +1525,8 @@ fn cli() -> Command { .arg(pg_version_arg.clone()) .arg(Arg::new("set-default").long("set-default").action(ArgAction::SetTrue).required(false) .help("Use this tenant in future CLI commands where tenant_id is needed, but not specified")) + .arg(Arg::new("shard-count").value_parser(value_parser!(u8)).long("shard-count").action(ArgAction::Set).help("Number of shards in the new tenant (default 1)")) + .arg(Arg::new("shard-stripe-size").value_parser(value_parser!(u32)).long("shard-stripe-size").action(ArgAction::Set).help("Sharding stripe size in pages")) ) .subcommand(Command::new("set-default").arg(tenant_id_arg.clone().required(true)) .about("Set a particular tenant as default in future CLI commands where tenant_id is needed, but not specified")) @@ -1368,6 +1537,9 @@ fn cli() -> Command { .about("Migrate a tenant from one pageserver to another") .arg(tenant_id_arg.clone()) .arg(pageserver_id_arg.clone())) + .subcommand(Command::new("status") + .about("Human readable summary of the tenant's shards and attachment locations") + .arg(tenant_id_arg.clone())) ) .subcommand( Command::new("pageserver") @@ -1387,6 +1559,12 @@ fn cli() -> Command { .about("Restart local pageserver") .arg(pageserver_config_args.clone()) ) + .subcommand(Command::new("set-state") + .arg(Arg::new("availability").value_parser(value_parser!(NodeAvailability)).long("availability").action(ArgAction::Set).help("Availability state: offline,active")) + .arg(Arg::new("scheduling").value_parser(value_parser!(NodeSchedulingPolicy)).long("scheduling").action(ArgAction::Set).help("Scheduling state: draining,pause,filling,active")) + .about("Set scheduling or availability state of pageserver node") + .arg(pageserver_config_args.clone()) + ) ) .subcommand( Command::new("attachment_service") diff --git a/control_plane/src/endpoint.rs b/control_plane/src/endpoint.rs index 3d5dfd6311..43f8ea3b43 100644 --- a/control_plane/src/endpoint.rs +++ b/control_plane/src/endpoint.rs @@ -49,10 +49,11 @@ use compute_api::spec::RemoteExtSpec; use nix::sys::signal::kill; use nix::sys::signal::Signal; use serde::{Deserialize, Serialize}; +use url::Host; use utils::id::{NodeId, TenantId, TimelineId}; +use crate::attachment_service::AttachmentService; use crate::local_env::LocalEnv; -use crate::pageserver::PageServerNode; use crate::postgresql_conf::PostgresConf; use compute_api::responses::{ComputeState, ComputeStatus}; @@ -69,7 +70,6 @@ pub struct EndpointConf { http_port: u16, pg_version: u32, skip_pg_catalog_updates: bool, - pageserver_id: NodeId, } // @@ -121,19 +121,14 @@ impl ComputeControlPlane { http_port: Option, pg_version: u32, mode: ComputeMode, - pageserver_id: NodeId, ) -> Result> { let pg_port = pg_port.unwrap_or_else(|| self.get_port()); let http_port = http_port.unwrap_or_else(|| self.get_port() + 1); - let pageserver = - PageServerNode::from_env(&self.env, self.env.get_pageserver_conf(pageserver_id)?); - let ep = Arc::new(Endpoint { endpoint_id: endpoint_id.to_owned(), pg_address: SocketAddr::new("127.0.0.1".parse().unwrap(), pg_port), http_address: SocketAddr::new("127.0.0.1".parse().unwrap(), http_port), env: self.env.clone(), - pageserver, timeline_id, mode, tenant_id, @@ -159,7 +154,6 @@ impl ComputeControlPlane { pg_port, pg_version, skip_pg_catalog_updates: true, - pageserver_id, })?, )?; std::fs::write( @@ -218,7 +212,6 @@ pub struct Endpoint { // These are not part of the endpoint as such, but the environment // the endpoint runs in. pub env: LocalEnv, - pageserver: PageServerNode, // Optimizations skip_pg_catalog_updates: bool, @@ -241,15 +234,11 @@ impl Endpoint { let conf: EndpointConf = serde_json::from_slice(&std::fs::read(entry.path().join("endpoint.json"))?)?; - let pageserver = - PageServerNode::from_env(env, env.get_pageserver_conf(conf.pageserver_id)?); - Ok(Endpoint { pg_address: SocketAddr::new("127.0.0.1".parse().unwrap(), conf.pg_port), http_address: SocketAddr::new("127.0.0.1".parse().unwrap(), conf.http_port), endpoint_id, env: env.clone(), - pageserver, timeline_id: conf.timeline_id, mode: conf.mode, tenant_id: conf.tenant_id, @@ -469,11 +458,21 @@ impl Endpoint { } } + fn build_pageserver_connstr(pageservers: &[(Host, u16)]) -> String { + pageservers + .iter() + .map(|(host, port)| format!("postgresql://no_user@{host}:{port}")) + .collect::>() + .join(",") + } + pub async fn start( &self, auth_token: &Option, safekeepers: Vec, + pageservers: Vec<(Host, u16)>, remote_ext_config: Option<&String>, + shard_stripe_size: usize, ) -> Result<()> { if self.status() == "running" { anyhow::bail!("The endpoint is already running"); @@ -487,13 +486,9 @@ impl Endpoint { std::fs::remove_dir_all(self.pgdata())?; } - let pageserver_connstring = { - let config = &self.pageserver.pg_connection_config; - let (host, port) = (config.host(), config.port()); + let pageserver_connstring = Self::build_pageserver_connstr(&pageservers); + assert!(!pageserver_connstring.is_empty()); - // NOTE: avoid spaces in connection string, because it is less error prone if we forward it somewhere. - format!("postgresql://no_user@{host}:{port}") - }; let mut safekeeper_connstrings = Vec::new(); if self.mode == ComputeMode::Primary { for sk_id in safekeepers { @@ -543,6 +538,7 @@ impl Endpoint { storage_auth_token: auth_token.clone(), remote_extensions, pgbouncer_settings: None, + shard_stripe_size: Some(shard_stripe_size), }; let spec_path = self.endpoint_path().join("spec.json"); std::fs::write(spec_path, serde_json::to_string_pretty(&spec)?)?; @@ -665,7 +661,7 @@ impl Endpoint { } } - pub async fn reconfigure(&self, pageserver_id: Option) -> Result<()> { + pub async fn reconfigure(&self, mut pageservers: Vec<(Host, u16)>) -> Result<()> { let mut spec: ComputeSpec = { let spec_path = self.endpoint_path().join("spec.json"); let file = std::fs::File::open(spec_path)?; @@ -675,25 +671,27 @@ impl Endpoint { let postgresql_conf = self.read_postgresql_conf()?; spec.cluster.postgresql_conf = Some(postgresql_conf); - if let Some(pageserver_id) = pageserver_id { - let endpoint_config_path = self.endpoint_path().join("endpoint.json"); - let mut endpoint_conf: EndpointConf = { - let file = std::fs::File::open(&endpoint_config_path)?; - serde_json::from_reader(file)? - }; - endpoint_conf.pageserver_id = pageserver_id; - std::fs::write( - endpoint_config_path, - serde_json::to_string_pretty(&endpoint_conf)?, - )?; - - let pageserver = - PageServerNode::from_env(&self.env, self.env.get_pageserver_conf(pageserver_id)?); - let ps_http_conf = &pageserver.pg_connection_config; - let (host, port) = (ps_http_conf.host(), ps_http_conf.port()); - spec.pageserver_connstring = Some(format!("postgresql://no_user@{host}:{port}")); + // If we weren't given explicit pageservers, query the attachment service + if pageservers.is_empty() { + let attachment_service = AttachmentService::from_env(&self.env); + let locate_result = attachment_service.tenant_locate(self.tenant_id).await?; + pageservers = locate_result + .shards + .into_iter() + .map(|shard| { + ( + Host::parse(&shard.listen_pg_addr) + .expect("Attachment service reported bad hostname"), + shard.listen_pg_port, + ) + }) + .collect::>(); } + let pageserver_connstr = Self::build_pageserver_connstr(&pageservers); + assert!(!pageserver_connstr.is_empty()); + spec.pageserver_connstring = Some(pageserver_connstr); + let client = reqwest::Client::new(); let response = client .post(format!( diff --git a/control_plane/src/lib.rs b/control_plane/src/lib.rs index 52a0e20429..bb79d36bfc 100644 --- a/control_plane/src/lib.rs +++ b/control_plane/src/lib.rs @@ -14,4 +14,3 @@ pub mod local_env; pub mod pageserver; pub mod postgresql_conf; pub mod safekeeper; -pub mod tenant_migration; diff --git a/control_plane/src/local_env.rs b/control_plane/src/local_env.rs index c87875457c..4460fdd3a6 100644 --- a/control_plane/src/local_env.rs +++ b/control_plane/src/local_env.rs @@ -251,7 +251,13 @@ impl LocalEnv { if let Some(conf) = self.pageservers.iter().find(|node| node.id == id) { Ok(conf) } else { - bail!("could not find pageserver {id}") + let have_ids = self + .pageservers + .iter() + .map(|node| format!("{}:{}", node.id, node.listen_http_addr)) + .collect::>(); + let joined = have_ids.join(","); + bail!("could not find pageserver {id}, have ids {joined}") } } diff --git a/control_plane/src/pageserver.rs b/control_plane/src/pageserver.rs index fb0d251722..18ccf6bd98 100644 --- a/control_plane/src/pageserver.rs +++ b/control_plane/src/pageserver.rs @@ -17,7 +17,9 @@ use std::time::Duration; use anyhow::{bail, Context}; use camino::Utf8PathBuf; use futures::SinkExt; -use pageserver_api::models::{self, LocationConfig, TenantInfo, TimelineInfo}; +use pageserver_api::models::{ + self, LocationConfig, ShardParameters, TenantHistorySize, TenantInfo, TimelineInfo, +}; use pageserver_api::shard::TenantShardId; use pageserver_client::mgmt_api; use postgres_backend::AuthType; @@ -106,6 +108,16 @@ impl PageServerNode { "control_plane_api='{}'", control_plane_api.as_str() )); + + // Attachment service uses the same auth as pageserver: if JWT is enabled + // for us, we will also need it to talk to them. + if matches!(self.conf.http_auth_type, AuthType::NeonJWT) { + let jwt_token = self + .env + .generate_auth_token(&Claims::new(None, Scope::PageServerApi)) + .unwrap(); + overrides.push(format!("control_plane_api_token='{}'", jwt_token)); + } } if !cli_overrides @@ -301,16 +313,8 @@ impl PageServerNode { pub async fn tenant_list(&self) -> mgmt_api::Result> { self.http_client.list_tenants().await } - - pub async fn tenant_create( - &self, - new_tenant_id: TenantId, - generation: Option, - settings: HashMap<&str, &str>, - ) -> anyhow::Result { - let mut settings = settings.clone(); - - let config = models::TenantConfig { + pub fn parse_config(mut settings: HashMap<&str, &str>) -> anyhow::Result { + let result = models::TenantConfig { checkpoint_distance: settings .remove("checkpoint_distance") .map(|x| x.parse::()) @@ -371,11 +375,26 @@ impl PageServerNode { .context("Failed to parse 'gc_feedback' as bool")?, heatmap_period: settings.remove("heatmap_period").map(|x| x.to_string()), }; + if !settings.is_empty() { + bail!("Unrecognized tenant settings: {settings:?}") + } else { + Ok(result) + } + } + + pub async fn tenant_create( + &self, + new_tenant_id: TenantId, + generation: Option, + settings: HashMap<&str, &str>, + ) -> anyhow::Result { + let config = Self::parse_config(settings.clone())?; let request = models::TenantCreateRequest { new_tenant_id: TenantShardId::unsharded(new_tenant_id), generation, config, + shard_parameters: ShardParameters::default(), }; if !settings.is_empty() { bail!("Unrecognized tenant settings: {settings:?}") @@ -471,18 +490,21 @@ impl PageServerNode { pub async fn location_config( &self, - tenant_id: TenantId, + tenant_shard_id: TenantShardId, config: LocationConfig, flush_ms: Option, ) -> anyhow::Result<()> { Ok(self .http_client - .location_config(tenant_id, config, flush_ms) + .location_config(tenant_shard_id, config, flush_ms) .await?) } - pub async fn timeline_list(&self, tenant_id: &TenantId) -> anyhow::Result> { - Ok(self.http_client.list_timelines(*tenant_id).await?) + pub async fn timeline_list( + &self, + tenant_shard_id: &TenantShardId, + ) -> anyhow::Result> { + Ok(self.http_client.list_timelines(*tenant_shard_id).await?) } pub async fn tenant_secondary_download(&self, tenant_id: &TenantShardId) -> anyhow::Result<()> { @@ -494,15 +516,13 @@ impl PageServerNode { pub async fn timeline_create( &self, - tenant_id: TenantId, - new_timeline_id: Option, + tenant_shard_id: TenantShardId, + new_timeline_id: TimelineId, ancestor_start_lsn: Option, ancestor_timeline_id: Option, pg_version: Option, existing_initdb_timeline_id: Option, ) -> anyhow::Result { - // If timeline ID was not specified, generate one - let new_timeline_id = new_timeline_id.unwrap_or(TimelineId::generate()); let req = models::TimelineCreateRequest { new_timeline_id, ancestor_start_lsn, @@ -510,7 +530,10 @@ impl PageServerNode { pg_version, existing_initdb_timeline_id, }; - Ok(self.http_client.timeline_create(tenant_id, &req).await?) + Ok(self + .http_client + .timeline_create(tenant_shard_id, &req) + .await?) } /// Import a basebackup prepared using either: @@ -588,4 +611,14 @@ impl PageServerNode { Ok(()) } + + pub async fn tenant_synthetic_size( + &self, + tenant_shard_id: TenantShardId, + ) -> anyhow::Result { + Ok(self + .http_client + .tenant_synthetic_size(tenant_shard_id) + .await?) + } } diff --git a/control_plane/src/tenant_migration.rs b/control_plane/src/tenant_migration.rs deleted file mode 100644 index 23ea8f4060..0000000000 --- a/control_plane/src/tenant_migration.rs +++ /dev/null @@ -1,220 +0,0 @@ -//! -//! Functionality for migrating tenants across pageservers: unlike most of neon_local, this code -//! isn't scoped to a particular physical service, as it needs to update compute endpoints to -//! point to the new pageserver. -//! -use crate::local_env::LocalEnv; -use crate::{ - attachment_service::AttachmentService, endpoint::ComputeControlPlane, - pageserver::PageServerNode, -}; -use pageserver_api::models::{ - LocationConfig, LocationConfigMode, LocationConfigSecondary, TenantConfig, -}; -use pageserver_api::shard::TenantShardId; -use std::collections::HashMap; -use std::time::Duration; -use utils::{ - id::{TenantId, TimelineId}, - lsn::Lsn, -}; - -/// Given an attached pageserver, retrieve the LSN for all timelines -async fn get_lsns( - tenant_id: TenantId, - pageserver: &PageServerNode, -) -> anyhow::Result> { - let timelines = pageserver.timeline_list(&tenant_id).await?; - Ok(timelines - .into_iter() - .map(|t| (t.timeline_id, t.last_record_lsn)) - .collect()) -} - -/// Wait for the timeline LSNs on `pageserver` to catch up with or overtake -/// `baseline`. -async fn await_lsn( - tenant_id: TenantId, - pageserver: &PageServerNode, - baseline: HashMap, -) -> anyhow::Result<()> { - loop { - let latest = match get_lsns(tenant_id, pageserver).await { - Ok(l) => l, - Err(_e) => { - println!( - "🕑 Waiting for pageserver {} to activate...", - pageserver.conf.id - ); - std::thread::sleep(Duration::from_millis(500)); - continue; - } - }; - - let mut any_behind: bool = false; - for (timeline_id, baseline_lsn) in &baseline { - match latest.get(timeline_id) { - Some(latest_lsn) => { - println!("🕑 LSN origin {baseline_lsn} vs destination {latest_lsn}"); - if latest_lsn < baseline_lsn { - any_behind = true; - } - } - None => { - // Expected timeline isn't yet visible on migration destination. - // (IRL we would have to account for timeline deletion, but this - // is just test helper) - any_behind = true; - } - } - } - - if !any_behind { - println!("✅ LSN caught up. Proceeding..."); - break; - } else { - std::thread::sleep(Duration::from_millis(500)); - } - } - - Ok(()) -} - -/// This function spans multiple services, to demonstrate live migration of a tenant -/// between pageservers: -/// - Coordinate attach/secondary/detach on pageservers -/// - call into attachment_service for generations -/// - reconfigure compute endpoints to point to new attached pageserver -pub async fn migrate_tenant( - env: &LocalEnv, - tenant_id: TenantId, - dest_ps: PageServerNode, -) -> anyhow::Result<()> { - println!("🤔 Checking existing status..."); - let attachment_service = AttachmentService::from_env(env); - - fn build_location_config( - mode: LocationConfigMode, - generation: Option, - secondary_conf: Option, - ) -> LocationConfig { - LocationConfig { - mode, - generation, - secondary_conf, - tenant_conf: TenantConfig::default(), - shard_number: 0, - shard_count: 0, - shard_stripe_size: 0, - } - } - - let previous = attachment_service.inspect(tenant_id).await?; - let mut baseline_lsns = None; - if let Some((generation, origin_ps_id)) = &previous { - let origin_ps = PageServerNode::from_env(env, env.get_pageserver_conf(*origin_ps_id)?); - - if origin_ps_id == &dest_ps.conf.id { - println!("🔁 Already attached to {origin_ps_id}, freshening..."); - let gen = attachment_service - .attach_hook(tenant_id, dest_ps.conf.id) - .await?; - let dest_conf = build_location_config(LocationConfigMode::AttachedSingle, gen, None); - dest_ps.location_config(tenant_id, dest_conf, None).await?; - println!("✅ Migration complete"); - return Ok(()); - } - - println!("🔁 Switching origin pageserver {origin_ps_id} to stale mode"); - - let stale_conf = - build_location_config(LocationConfigMode::AttachedStale, Some(*generation), None); - origin_ps - .location_config(tenant_id, stale_conf, Some(Duration::from_secs(10))) - .await?; - - baseline_lsns = Some(get_lsns(tenant_id, &origin_ps).await?); - } - - println!( - "🔁 Downloading latest layers to destination pageserver {}", - dest_ps.conf.id - ); - match dest_ps - .tenant_secondary_download(&TenantShardId::unsharded(tenant_id)) - .await - { - Ok(()) => {} - Err(_) => { - println!(" (skipping, destination wasn't in secondary mode)") - } - } - - let gen = attachment_service - .attach_hook(tenant_id, dest_ps.conf.id) - .await?; - let dest_conf = build_location_config(LocationConfigMode::AttachedMulti, gen, None); - - println!("🔁 Attaching to pageserver {}", dest_ps.conf.id); - dest_ps.location_config(tenant_id, dest_conf, None).await?; - - if let Some(baseline) = baseline_lsns { - println!("🕑 Waiting for LSN to catch up..."); - await_lsn(tenant_id, &dest_ps, baseline).await?; - } - - let cplane = ComputeControlPlane::load(env.clone())?; - for (endpoint_name, endpoint) in &cplane.endpoints { - if endpoint.tenant_id == tenant_id { - println!( - "🔁 Reconfiguring endpoint {} to use pageserver {}", - endpoint_name, dest_ps.conf.id - ); - endpoint.reconfigure(Some(dest_ps.conf.id)).await?; - } - } - - for other_ps_conf in &env.pageservers { - if other_ps_conf.id == dest_ps.conf.id { - continue; - } - - let other_ps = PageServerNode::from_env(env, other_ps_conf); - let other_ps_tenants = other_ps.tenant_list().await?; - - // Check if this tenant is attached - let found = other_ps_tenants - .into_iter() - .map(|t| t.id) - .any(|i| i.tenant_id == tenant_id); - if !found { - continue; - } - - // Downgrade to a secondary location - let secondary_conf = build_location_config( - LocationConfigMode::Secondary, - None, - Some(LocationConfigSecondary { warm: true }), - ); - - println!( - "💤 Switching to secondary mode on pageserver {}", - other_ps.conf.id - ); - other_ps - .location_config(tenant_id, secondary_conf, None) - .await?; - } - - println!( - "🔁 Switching to AttachedSingle mode on pageserver {}", - dest_ps.conf.id - ); - let dest_conf = build_location_config(LocationConfigMode::AttachedSingle, gen, None); - dest_ps.location_config(tenant_id, dest_conf, None).await?; - - println!("✅ Migration complete"); - - Ok(()) -} diff --git a/libs/compute_api/src/spec.rs b/libs/compute_api/src/spec.rs index 810b5c90d4..13ac18e0c5 100644 --- a/libs/compute_api/src/spec.rs +++ b/libs/compute_api/src/spec.rs @@ -75,6 +75,10 @@ pub struct ComputeSpec { pub remote_extensions: Option, pub pgbouncer_settings: Option>, + + // Stripe size for pageserver sharding, in pages + #[serde(default)] + pub shard_stripe_size: Option, } /// Feature flag to signal `compute_ctl` to enable certain experimental functionality. diff --git a/libs/pageserver_api/Cargo.toml b/libs/pageserver_api/Cargo.toml index 4146597d8d..96c6c10d3e 100644 --- a/libs/pageserver_api/Cargo.toml +++ b/libs/pageserver_api/Cargo.toml @@ -19,6 +19,7 @@ strum.workspace = true strum_macros.workspace = true hex.workspace = true thiserror.workspace = true +humantime-serde.workspace = true workspace_hack.workspace = true diff --git a/libs/pageserver_api/src/models.rs b/libs/pageserver_api/src/models.rs index 25530c9fc9..86d2c2a7ca 100644 --- a/libs/pageserver_api/src/models.rs +++ b/libs/pageserver_api/src/models.rs @@ -4,7 +4,7 @@ use std::{ collections::HashMap, io::{BufRead, Read}, num::{NonZeroU64, NonZeroUsize}, - time::SystemTime, + time::{Duration, SystemTime}, }; use byteorder::{BigEndian, ReadBytesExt}; @@ -18,7 +18,10 @@ use utils::{ lsn::Lsn, }; -use crate::{reltag::RelTag, shard::TenantShardId}; +use crate::{ + reltag::RelTag, + shard::{ShardCount, ShardStripeSize, TenantShardId}, +}; use anyhow::bail; use bytes::{Buf, BufMut, Bytes, BytesMut}; @@ -188,6 +191,31 @@ pub struct TimelineCreateRequest { pub pg_version: Option, } +/// Parameters that apply to all shards in a tenant. Used during tenant creation. +#[derive(Serialize, Deserialize, Debug)] +#[serde(deny_unknown_fields)] +pub struct ShardParameters { + pub count: ShardCount, + pub stripe_size: ShardStripeSize, +} + +impl ShardParameters { + pub const DEFAULT_STRIPE_SIZE: ShardStripeSize = ShardStripeSize(256 * 1024 / 8); + + pub fn is_unsharded(&self) -> bool { + self.count == ShardCount(0) + } +} + +impl Default for ShardParameters { + fn default() -> Self { + Self { + count: ShardCount(0), + stripe_size: Self::DEFAULT_STRIPE_SIZE, + } + } +} + #[derive(Serialize, Deserialize, Debug)] #[serde(deny_unknown_fields)] pub struct TenantCreateRequest { @@ -195,6 +223,12 @@ pub struct TenantCreateRequest { #[serde(default)] #[serde(skip_serializing_if = "Option::is_none")] pub generation: Option, + + // If omitted, create a single shard with TenantShardId::unsharded() + #[serde(default)] + #[serde(skip_serializing_if = "ShardParameters::is_unsharded")] + pub shard_parameters: ShardParameters, + #[serde(flatten)] pub config: TenantConfig, // as we have a flattened field, we should reject all unknown fields in it } @@ -217,7 +251,7 @@ impl std::ops::Deref for TenantCreateRequest { /// An alternative representation of `pageserver::tenant::TenantConf` with /// simpler types. -#[derive(Serialize, Deserialize, Debug, Default)] +#[derive(Serialize, Deserialize, Debug, Default, Clone, Eq, PartialEq)] pub struct TenantConfig { pub checkpoint_distance: Option, pub checkpoint_timeout: Option, @@ -232,21 +266,41 @@ pub struct TenantConfig { pub lagging_wal_timeout: Option, pub max_lsn_wal_lag: Option, pub trace_read_requests: Option, - // We defer the parsing of the eviction_policy field to the request handler. - // Otherwise we'd have to move the types for eviction policy into this package. - // We might do that once the eviction feature has stabilizied. - // For now, this field is not even documented in the openapi_spec.yml. - pub eviction_policy: Option, + pub eviction_policy: Option, pub min_resident_size_override: Option, pub evictions_low_residence_duration_metric_threshold: Option, pub gc_feedback: Option, pub heatmap_period: Option, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind")] +pub enum EvictionPolicy { + NoEviction, + LayerAccessThreshold(EvictionPolicyLayerAccessThreshold), +} + +impl EvictionPolicy { + pub fn discriminant_str(&self) -> &'static str { + match self { + EvictionPolicy::NoEviction => "NoEviction", + EvictionPolicy::LayerAccessThreshold(_) => "LayerAccessThreshold", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct EvictionPolicyLayerAccessThreshold { + #[serde(with = "humantime_serde")] + pub period: Duration, + #[serde(with = "humantime_serde")] + pub threshold: Duration, +} + /// A flattened analog of a `pagesever::tenant::LocationMode`, which /// lists out all possible states (and the virtual "Detached" state) /// in a flat form rather than using rust-style enums. -#[derive(Serialize, Deserialize, Debug)] +#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq)] pub enum LocationConfigMode { AttachedSingle, AttachedMulti, @@ -255,19 +309,21 @@ pub enum LocationConfigMode { Detached, } -#[derive(Serialize, Deserialize, Debug)] +#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq)] pub struct LocationConfigSecondary { pub warm: bool, } /// An alternative representation of `pageserver::tenant::LocationConf`, /// for use in external-facing APIs. -#[derive(Serialize, Deserialize, Debug)] +#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq)] pub struct LocationConfig { pub mode: LocationConfigMode, /// If attaching, in what generation? #[serde(default)] pub generation: Option, + + // If requesting mode `Secondary`, configuration for that. #[serde(default)] pub secondary_conf: Option, @@ -280,11 +336,17 @@ pub struct LocationConfig { #[serde(default)] pub shard_stripe_size: u32, - // If requesting mode `Secondary`, configuration for that. - // Custom storage configuration for the tenant, if any + // This configuration only affects attached mode, but should be provided irrespective + // of the mode, as a secondary location might transition on startup if the response + // to the `/re-attach` control plane API requests it. pub tenant_conf: TenantConfig, } +#[derive(Serialize, Deserialize)] +pub struct LocationConfigListResponse { + pub tenant_shards: Vec<(TenantShardId, Option)>, +} + #[derive(Serialize, Deserialize)] #[serde(transparent)] pub struct TenantCreateResponse(pub TenantId); @@ -297,7 +359,7 @@ pub struct StatusResponse { #[derive(Serialize, Deserialize, Debug)] #[serde(deny_unknown_fields)] pub struct TenantLocationConfigRequest { - pub tenant_id: TenantId, + pub tenant_id: TenantShardId, #[serde(flatten)] pub config: LocationConfig, // as we have a flattened field, we should reject all unknown fields in it } @@ -660,6 +722,17 @@ pub struct PagestreamDbSizeResponse { pub db_size: i64, } +// This is a cut-down version of TenantHistorySize from the pageserver crate, omitting fields +// that require pageserver-internal types. It is sufficient to get the total size. +#[derive(Serialize, Deserialize, Debug)] +pub struct TenantHistorySize { + pub id: TenantId, + /// Size is a mixture of WAL and logical size, so the unit is bytes. + /// + /// Will be none if `?inputs_only=true` was given. + pub size: Option, +} + impl PagestreamFeMessage { pub fn serialize(&self) -> Bytes { let mut bytes = BytesMut::new(); diff --git a/libs/pageserver_api/src/shard.rs b/libs/pageserver_api/src/shard.rs index 18ef2be523..e27aad8156 100644 --- a/libs/pageserver_api/src/shard.rs +++ b/libs/pageserver_api/src/shard.rs @@ -1,6 +1,9 @@ use std::{ops::RangeInclusive, str::FromStr}; -use crate::key::{is_rel_block_key, Key}; +use crate::{ + key::{is_rel_block_key, Key}, + models::ShardParameters, +}; use hex::FromHex; use serde::{Deserialize, Serialize}; use thiserror; @@ -85,6 +88,12 @@ impl TenantShardId { pub fn is_unsharded(&self) -> bool { self.shard_number == ShardNumber(0) && self.shard_count == ShardCount(0) } + pub fn to_index(&self) -> ShardIndex { + ShardIndex { + shard_number: self.shard_number, + shard_count: self.shard_count, + } + } } /// Formatting helper @@ -333,7 +342,7 @@ const DEFAULT_STRIPE_SIZE: ShardStripeSize = ShardStripeSize(256 * 1024 / 8); pub struct ShardIdentity { pub number: ShardNumber, pub count: ShardCount, - stripe_size: ShardStripeSize, + pub stripe_size: ShardStripeSize, layout: ShardLayout, } @@ -403,6 +412,17 @@ impl ShardIdentity { } } + /// For use when creating ShardIdentity instances for new shards, where a creation request + /// specifies the ShardParameters that apply to all shards. + pub fn from_params(number: ShardNumber, params: &ShardParameters) -> Self { + Self { + number, + count: params.count, + layout: LAYOUT_V1, + stripe_size: params.stripe_size, + } + } + fn is_broken(&self) -> bool { self.layout == LAYOUT_BROKEN } diff --git a/libs/remote_storage/src/azure_blob.rs b/libs/remote_storage/src/azure_blob.rs index 18cf5d97ba..7895a21f66 100644 --- a/libs/remote_storage/src/azure_blob.rs +++ b/libs/remote_storage/src/azure_blob.rs @@ -5,7 +5,9 @@ use std::collections::HashMap; use std::env; use std::num::NonZeroU32; use std::pin::Pin; +use std::str::FromStr; use std::sync::Arc; +use std::time::Duration; use super::REMOTE_STORAGE_PREFIX_SEPARATOR; use anyhow::Result; @@ -13,12 +15,14 @@ use azure_core::request_options::{MaxResults, Metadata, Range}; use azure_core::RetryOptions; use azure_identity::DefaultAzureCredential; use azure_storage::StorageCredentials; +use azure_storage_blobs::blob::CopyStatus; use azure_storage_blobs::prelude::ClientBuilder; use azure_storage_blobs::{blob::operations::GetBlobBuilder, prelude::ContainerClient}; use bytes::Bytes; use futures::stream::Stream; use futures_util::StreamExt; -use http_types::StatusCode; +use http_types::{StatusCode, Url}; +use tokio::time::Instant; use tracing::debug; use crate::s3_bucket::RequestKind; @@ -323,10 +327,49 @@ impl RemoteStorage for AzureBlobStorage { Ok(()) } - async fn copy(&self, _from: &RemotePath, _to: &RemotePath) -> anyhow::Result<()> { - Err(anyhow::anyhow!( - "copy for azure blob storage is not implemented" - )) + async fn copy(&self, from: &RemotePath, to: &RemotePath) -> anyhow::Result<()> { + let _permit = self.permit(RequestKind::Copy).await; + let blob_client = self.client.blob_client(self.relative_path_to_name(to)); + + let source_url = format!( + "{}/{}", + self.client.url()?, + self.relative_path_to_name(from) + ); + let builder = blob_client.copy(Url::from_str(&source_url)?); + + let result = builder.into_future().await?; + + let mut copy_status = result.copy_status; + let start_time = Instant::now(); + const MAX_WAIT_TIME: Duration = Duration::from_secs(60); + loop { + match copy_status { + CopyStatus::Aborted => { + anyhow::bail!("Received abort for copy from {from} to {to}."); + } + CopyStatus::Failed => { + anyhow::bail!("Received failure response for copy from {from} to {to}."); + } + CopyStatus::Success => return Ok(()), + CopyStatus::Pending => (), + } + // The copy is taking longer. Waiting a second and then re-trying. + // TODO estimate time based on copy_progress and adjust time based on that + tokio::time::sleep(Duration::from_millis(1000)).await; + let properties = blob_client.get_properties().into_future().await?; + let Some(status) = properties.blob.properties.copy_status else { + tracing::warn!("copy_status for copy is None!, from={from}, to={to}"); + return Ok(()); + }; + if start_time.elapsed() > MAX_WAIT_TIME { + anyhow::bail!("Copy from from {from} to {to} took longer than limit MAX_WAIT_TIME={}s. copy_pogress={:?}.", + MAX_WAIT_TIME.as_secs_f32(), + properties.blob.properties.copy_progress, + ); + } + copy_status = status; + } } } diff --git a/libs/remote_storage/tests/common/tests.rs b/libs/remote_storage/tests/common/tests.rs new file mode 100644 index 0000000000..abccc24c97 --- /dev/null +++ b/libs/remote_storage/tests/common/tests.rs @@ -0,0 +1,288 @@ +use anyhow::Context; +use camino::Utf8Path; +use remote_storage::RemotePath; +use std::collections::HashSet; +use std::sync::Arc; +use test_context::test_context; +use tracing::debug; + +use crate::common::{download_to_vec, upload_stream, wrap_stream}; + +use super::{ + MaybeEnabledStorage, MaybeEnabledStorageWithSimpleTestBlobs, MaybeEnabledStorageWithTestBlobs, +}; + +/// Tests that S3 client can list all prefixes, even if the response come paginated and requires multiple S3 queries. +/// Uses real S3 and requires [`ENABLE_REAL_S3_REMOTE_STORAGE_ENV_VAR_NAME`] and related S3 cred env vars specified. +/// See the client creation in [`create_s3_client`] for details on the required env vars. +/// If real S3 tests are disabled, the test passes, skipping any real test run: currently, there's no way to mark the test ignored in runtime with the +/// deafult test framework, see https://github.com/rust-lang/rust/issues/68007 for details. +/// +/// First, the test creates a set of S3 objects with keys `/${random_prefix_part}/${base_prefix_str}/sub_prefix_${i}/blob_${i}` in [`upload_remote_data`] +/// where +/// * `random_prefix_part` is set for the entire S3 client during the S3 client creation in [`create_s3_client`], to avoid multiple test runs interference +/// * `base_prefix_str` is a common prefix to use in the client requests: we would want to ensure that the client is able to list nested prefixes inside the bucket +/// +/// Then, verifies that the client does return correct prefixes when queried: +/// * with no prefix, it lists everything after its `${random_prefix_part}/` — that should be `${base_prefix_str}` value only +/// * with `${base_prefix_str}/` prefix, it lists every `sub_prefix_${i}` +/// +/// With the real S3 enabled and `#[cfg(test)]` Rust configuration used, the S3 client test adds a `max-keys` param to limit the response keys. +/// This way, we are able to test the pagination implicitly, by ensuring all results are returned from the remote storage and avoid uploading too many blobs to S3, +/// since current default AWS S3 pagination limit is 1000. +/// (see https://docs.aws.amazon.com/AmazonS3/latest/API/API_ListObjectsV2.html#API_ListObjectsV2_RequestSyntax) +/// +/// Lastly, the test attempts to clean up and remove all uploaded S3 files. +/// If any errors appear during the clean up, they get logged, but the test is not failed or stopped until clean up is finished. +#[test_context(MaybeEnabledStorageWithTestBlobs)] +#[tokio::test] +async fn pagination_should_work(ctx: &mut MaybeEnabledStorageWithTestBlobs) -> anyhow::Result<()> { + let ctx = match ctx { + MaybeEnabledStorageWithTestBlobs::Enabled(ctx) => ctx, + MaybeEnabledStorageWithTestBlobs::Disabled => return Ok(()), + MaybeEnabledStorageWithTestBlobs::UploadsFailed(e, _) => { + anyhow::bail!("S3 init failed: {e:?}") + } + }; + + let test_client = Arc::clone(&ctx.enabled.client); + let expected_remote_prefixes = ctx.remote_prefixes.clone(); + + let base_prefix = RemotePath::new(Utf8Path::new(ctx.enabled.base_prefix)) + .context("common_prefix construction")?; + let root_remote_prefixes = test_client + .list_prefixes(None) + .await + .context("client list root prefixes failure")? + .into_iter() + .collect::>(); + assert_eq!( + root_remote_prefixes, HashSet::from([base_prefix.clone()]), + "remote storage root prefixes list mismatches with the uploads. Returned prefixes: {root_remote_prefixes:?}" + ); + + let nested_remote_prefixes = test_client + .list_prefixes(Some(&base_prefix)) + .await + .context("client list nested prefixes failure")? + .into_iter() + .collect::>(); + let remote_only_prefixes = nested_remote_prefixes + .difference(&expected_remote_prefixes) + .collect::>(); + let missing_uploaded_prefixes = expected_remote_prefixes + .difference(&nested_remote_prefixes) + .collect::>(); + assert_eq!( + remote_only_prefixes.len() + missing_uploaded_prefixes.len(), 0, + "remote storage nested prefixes list mismatches with the uploads. Remote only prefixes: {remote_only_prefixes:?}, missing uploaded prefixes: {missing_uploaded_prefixes:?}", + ); + + Ok(()) +} + +/// Tests that S3 client can list all files in a folder, even if the response comes paginated and requirees multiple S3 queries. +/// Uses real S3 and requires [`ENABLE_REAL_S3_REMOTE_STORAGE_ENV_VAR_NAME`] and related S3 cred env vars specified. Test will skip real code and pass if env vars not set. +/// See `s3_pagination_should_work` for more information. +/// +/// First, create a set of S3 objects with keys `random_prefix/folder{j}/blob_{i}.txt` in [`upload_remote_data`] +/// Then performs the following queries: +/// 1. `list_files(None)`. This should return all files `random_prefix/folder{j}/blob_{i}.txt` +/// 2. `list_files("folder1")`. This should return all files `random_prefix/folder1/blob_{i}.txt` +#[test_context(MaybeEnabledStorageWithSimpleTestBlobs)] +#[tokio::test] +async fn list_files_works(ctx: &mut MaybeEnabledStorageWithSimpleTestBlobs) -> anyhow::Result<()> { + let ctx = match ctx { + MaybeEnabledStorageWithSimpleTestBlobs::Enabled(ctx) => ctx, + MaybeEnabledStorageWithSimpleTestBlobs::Disabled => return Ok(()), + MaybeEnabledStorageWithSimpleTestBlobs::UploadsFailed(e, _) => { + anyhow::bail!("S3 init failed: {e:?}") + } + }; + let test_client = Arc::clone(&ctx.enabled.client); + let base_prefix = + RemotePath::new(Utf8Path::new("folder1")).context("common_prefix construction")?; + let root_files = test_client + .list_files(None) + .await + .context("client list root files failure")? + .into_iter() + .collect::>(); + assert_eq!( + root_files, + ctx.remote_blobs.clone(), + "remote storage list_files on root mismatches with the uploads." + ); + let nested_remote_files = test_client + .list_files(Some(&base_prefix)) + .await + .context("client list nested files failure")? + .into_iter() + .collect::>(); + let trim_remote_blobs: HashSet<_> = ctx + .remote_blobs + .iter() + .map(|x| x.get_path()) + .filter(|x| x.starts_with("folder1")) + .map(|x| RemotePath::new(x).expect("must be valid path")) + .collect(); + assert_eq!( + nested_remote_files, trim_remote_blobs, + "remote storage list_files on subdirrectory mismatches with the uploads." + ); + Ok(()) +} + +#[test_context(MaybeEnabledStorage)] +#[tokio::test] +async fn delete_non_exising_works(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<()> { + let ctx = match ctx { + MaybeEnabledStorage::Enabled(ctx) => ctx, + MaybeEnabledStorage::Disabled => return Ok(()), + }; + + let path = RemotePath::new(Utf8Path::new( + format!("{}/for_sure_there_is_nothing_there_really", ctx.base_prefix).as_str(), + )) + .with_context(|| "RemotePath conversion")?; + + ctx.client.delete(&path).await.expect("should succeed"); + + Ok(()) +} + +#[test_context(MaybeEnabledStorage)] +#[tokio::test] +async fn delete_objects_works(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<()> { + let ctx = match ctx { + MaybeEnabledStorage::Enabled(ctx) => ctx, + MaybeEnabledStorage::Disabled => return Ok(()), + }; + + let path1 = RemotePath::new(Utf8Path::new(format!("{}/path1", ctx.base_prefix).as_str())) + .with_context(|| "RemotePath conversion")?; + + let path2 = RemotePath::new(Utf8Path::new(format!("{}/path2", ctx.base_prefix).as_str())) + .with_context(|| "RemotePath conversion")?; + + let path3 = RemotePath::new(Utf8Path::new(format!("{}/path3", ctx.base_prefix).as_str())) + .with_context(|| "RemotePath conversion")?; + + let (data, len) = upload_stream("remote blob data1".as_bytes().into()); + ctx.client.upload(data, len, &path1, None).await?; + + let (data, len) = upload_stream("remote blob data2".as_bytes().into()); + ctx.client.upload(data, len, &path2, None).await?; + + let (data, len) = upload_stream("remote blob data3".as_bytes().into()); + ctx.client.upload(data, len, &path3, None).await?; + + ctx.client.delete_objects(&[path1, path2]).await?; + + let prefixes = ctx.client.list_prefixes(None).await?; + + assert_eq!(prefixes.len(), 1); + + ctx.client.delete_objects(&[path3]).await?; + + Ok(()) +} + +#[test_context(MaybeEnabledStorage)] +#[tokio::test] +async fn upload_download_works(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<()> { + let MaybeEnabledStorage::Enabled(ctx) = ctx else { + return Ok(()); + }; + + let path = RemotePath::new(Utf8Path::new(format!("{}/file", ctx.base_prefix).as_str())) + .with_context(|| "RemotePath conversion")?; + + let orig = bytes::Bytes::from_static("remote blob data here".as_bytes()); + + let (data, len) = wrap_stream(orig.clone()); + + ctx.client.upload(data, len, &path, None).await?; + + // Normal download request + let dl = ctx.client.download(&path).await?; + let buf = download_to_vec(dl).await?; + assert_eq!(&buf, &orig); + + // Full range (end specified) + let dl = ctx + .client + .download_byte_range(&path, 0, Some(len as u64)) + .await?; + let buf = download_to_vec(dl).await?; + assert_eq!(&buf, &orig); + + // partial range (end specified) + let dl = ctx.client.download_byte_range(&path, 4, Some(10)).await?; + let buf = download_to_vec(dl).await?; + assert_eq!(&buf, &orig[4..10]); + + // partial range (end beyond real end) + let dl = ctx + .client + .download_byte_range(&path, 8, Some(len as u64 * 100)) + .await?; + let buf = download_to_vec(dl).await?; + assert_eq!(&buf, &orig[8..]); + + // Partial range (end unspecified) + let dl = ctx.client.download_byte_range(&path, 4, None).await?; + let buf = download_to_vec(dl).await?; + assert_eq!(&buf, &orig[4..]); + + // Full range (end unspecified) + let dl = ctx.client.download_byte_range(&path, 0, None).await?; + let buf = download_to_vec(dl).await?; + assert_eq!(&buf, &orig); + + debug!("Cleanup: deleting file at path {path:?}"); + ctx.client + .delete(&path) + .await + .with_context(|| format!("{path:?} removal"))?; + + Ok(()) +} + +#[test_context(MaybeEnabledStorage)] +#[tokio::test] +async fn copy_works(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<()> { + let MaybeEnabledStorage::Enabled(ctx) = ctx else { + return Ok(()); + }; + + let path = RemotePath::new(Utf8Path::new( + format!("{}/file_to_copy", ctx.base_prefix).as_str(), + )) + .with_context(|| "RemotePath conversion")?; + let path_dest = RemotePath::new(Utf8Path::new( + format!("{}/file_dest", ctx.base_prefix).as_str(), + )) + .with_context(|| "RemotePath conversion")?; + + let orig = bytes::Bytes::from_static("remote blob data content".as_bytes()); + + let (data, len) = wrap_stream(orig.clone()); + + ctx.client.upload(data, len, &path, None).await?; + + // Normal download request + ctx.client.copy_object(&path, &path_dest).await?; + + let dl = ctx.client.download(&path_dest).await?; + let buf = download_to_vec(dl).await?; + assert_eq!(&buf, &orig); + + debug!("Cleanup: deleting file at path {path:?}"); + ctx.client + .delete_objects(&[path.clone(), path_dest.clone()]) + .await + .with_context(|| format!("{path:?} removal"))?; + + Ok(()) +} diff --git a/libs/remote_storage/tests/test_real_azure.rs b/libs/remote_storage/tests/test_real_azure.rs index 0387dc30e7..6f9a1ec6f7 100644 --- a/libs/remote_storage/tests/test_real_azure.rs +++ b/libs/remote_storage/tests/test_real_azure.rs @@ -6,263 +6,23 @@ use std::sync::Arc; use std::time::UNIX_EPOCH; use anyhow::Context; -use camino::Utf8Path; use remote_storage::{ AzureConfig, GenericRemoteStorage, RemotePath, RemoteStorageConfig, RemoteStorageKind, }; -use test_context::{test_context, AsyncTestContext}; -use tracing::{debug, info}; +use test_context::AsyncTestContext; +use tracing::info; mod common; -use common::{ - cleanup, download_to_vec, ensure_logging_ready, upload_remote_data, upload_simple_remote_data, - upload_stream, wrap_stream, -}; +#[path = "common/tests.rs"] +mod tests_azure; + +use common::{cleanup, ensure_logging_ready, upload_remote_data, upload_simple_remote_data}; const ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME: &str = "ENABLE_REAL_AZURE_REMOTE_STORAGE"; const BASE_PREFIX: &str = "test"; -/// Tests that the Azure client can list all prefixes, even if the response comes paginated and requires multiple HTTP queries. -/// Uses real Azure and requires [`ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME`] and related Azure cred env vars specified. -/// See the client creation in [`create_azure_client`] for details on the required env vars. -/// If real Azure tests are disabled, the test passes, skipping any real test run: currently, there's no way to mark the test ignored in runtime with the -/// deafult test framework, see https://github.com/rust-lang/rust/issues/68007 for details. -/// -/// First, the test creates a set of Azure blobs with keys `/${random_prefix_part}/${base_prefix_str}/sub_prefix_${i}/blob_${i}` in [`upload_remote_data`] -/// where -/// * `random_prefix_part` is set for the entire Azure client during the Azure client creation in [`create_azure_client`], to avoid multiple test runs interference -/// * `base_prefix_str` is a common prefix to use in the client requests: we would want to ensure that the client is able to list nested prefixes inside the bucket -/// -/// Then, verifies that the client does return correct prefixes when queried: -/// * with no prefix, it lists everything after its `${random_prefix_part}/` — that should be `${base_prefix_str}` value only -/// * with `${base_prefix_str}/` prefix, it lists every `sub_prefix_${i}` -/// -/// With the real Azure enabled and `#[cfg(test)]` Rust configuration used, the Azure client test adds a `max-keys` param to limit the response keys. -/// This way, we are able to test the pagination implicitly, by ensuring all results are returned from the remote storage and avoid uploading too many blobs to Azure. -/// -/// Lastly, the test attempts to clean up and remove all uploaded Azure files. -/// If any errors appear during the clean up, they get logged, but the test is not failed or stopped until clean up is finished. -#[test_context(MaybeEnabledAzureWithTestBlobs)] -#[tokio::test] -async fn azure_pagination_should_work( - ctx: &mut MaybeEnabledAzureWithTestBlobs, -) -> anyhow::Result<()> { - let ctx = match ctx { - MaybeEnabledAzureWithTestBlobs::Enabled(ctx) => ctx, - MaybeEnabledAzureWithTestBlobs::Disabled => return Ok(()), - MaybeEnabledAzureWithTestBlobs::UploadsFailed(e, _) => { - anyhow::bail!("Azure init failed: {e:?}") - } - }; - - let test_client = Arc::clone(&ctx.enabled.client); - let expected_remote_prefixes = ctx.remote_prefixes.clone(); - - let base_prefix = RemotePath::new(Utf8Path::new(ctx.enabled.base_prefix)) - .context("common_prefix construction")?; - let root_remote_prefixes = test_client - .list_prefixes(None) - .await - .context("client list root prefixes failure")? - .into_iter() - .collect::>(); - assert_eq!( - root_remote_prefixes, HashSet::from([base_prefix.clone()]), - "remote storage root prefixes list mismatches with the uploads. Returned prefixes: {root_remote_prefixes:?}" - ); - - let nested_remote_prefixes = test_client - .list_prefixes(Some(&base_prefix)) - .await - .context("client list nested prefixes failure")? - .into_iter() - .collect::>(); - let remote_only_prefixes = nested_remote_prefixes - .difference(&expected_remote_prefixes) - .collect::>(); - let missing_uploaded_prefixes = expected_remote_prefixes - .difference(&nested_remote_prefixes) - .collect::>(); - assert_eq!( - remote_only_prefixes.len() + missing_uploaded_prefixes.len(), 0, - "remote storage nested prefixes list mismatches with the uploads. Remote only prefixes: {remote_only_prefixes:?}, missing uploaded prefixes: {missing_uploaded_prefixes:?}", - ); - - Ok(()) -} - -/// Tests that Azure client can list all files in a folder, even if the response comes paginated and requirees multiple Azure queries. -/// Uses real Azure and requires [`ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME`] and related Azure cred env vars specified. Test will skip real code and pass if env vars not set. -/// See `Azure_pagination_should_work` for more information. -/// -/// First, create a set of Azure objects with keys `random_prefix/folder{j}/blob_{i}.txt` in [`upload_remote_data`] -/// Then performs the following queries: -/// 1. `list_files(None)`. This should return all files `random_prefix/folder{j}/blob_{i}.txt` -/// 2. `list_files("folder1")`. This should return all files `random_prefix/folder1/blob_{i}.txt` -#[test_context(MaybeEnabledAzureWithSimpleTestBlobs)] -#[tokio::test] -async fn azure_list_files_works( - ctx: &mut MaybeEnabledAzureWithSimpleTestBlobs, -) -> anyhow::Result<()> { - let ctx = match ctx { - MaybeEnabledAzureWithSimpleTestBlobs::Enabled(ctx) => ctx, - MaybeEnabledAzureWithSimpleTestBlobs::Disabled => return Ok(()), - MaybeEnabledAzureWithSimpleTestBlobs::UploadsFailed(e, _) => { - anyhow::bail!("Azure init failed: {e:?}") - } - }; - let test_client = Arc::clone(&ctx.enabled.client); - let base_prefix = - RemotePath::new(Utf8Path::new("folder1")).context("common_prefix construction")?; - let root_files = test_client - .list_files(None) - .await - .context("client list root files failure")? - .into_iter() - .collect::>(); - assert_eq!( - root_files, - ctx.remote_blobs.clone(), - "remote storage list_files on root mismatches with the uploads." - ); - let nested_remote_files = test_client - .list_files(Some(&base_prefix)) - .await - .context("client list nested files failure")? - .into_iter() - .collect::>(); - let trim_remote_blobs: HashSet<_> = ctx - .remote_blobs - .iter() - .map(|x| x.get_path()) - .filter(|x| x.starts_with("folder1")) - .map(|x| RemotePath::new(x).expect("must be valid path")) - .collect(); - assert_eq!( - nested_remote_files, trim_remote_blobs, - "remote storage list_files on subdirrectory mismatches with the uploads." - ); - Ok(()) -} - -#[test_context(MaybeEnabledAzure)] -#[tokio::test] -async fn azure_delete_non_exising_works(ctx: &mut MaybeEnabledAzure) -> anyhow::Result<()> { - let ctx = match ctx { - MaybeEnabledAzure::Enabled(ctx) => ctx, - MaybeEnabledAzure::Disabled => return Ok(()), - }; - - let path = RemotePath::new(Utf8Path::new( - format!("{}/for_sure_there_is_nothing_there_really", ctx.base_prefix).as_str(), - )) - .with_context(|| "RemotePath conversion")?; - - ctx.client.delete(&path).await.expect("should succeed"); - - Ok(()) -} - -#[test_context(MaybeEnabledAzure)] -#[tokio::test] -async fn azure_delete_objects_works(ctx: &mut MaybeEnabledAzure) -> anyhow::Result<()> { - let ctx = match ctx { - MaybeEnabledAzure::Enabled(ctx) => ctx, - MaybeEnabledAzure::Disabled => return Ok(()), - }; - - let path1 = RemotePath::new(Utf8Path::new(format!("{}/path1", ctx.base_prefix).as_str())) - .with_context(|| "RemotePath conversion")?; - - let path2 = RemotePath::new(Utf8Path::new(format!("{}/path2", ctx.base_prefix).as_str())) - .with_context(|| "RemotePath conversion")?; - - let path3 = RemotePath::new(Utf8Path::new(format!("{}/path3", ctx.base_prefix).as_str())) - .with_context(|| "RemotePath conversion")?; - - let (data, len) = upload_stream("remote blob data1".as_bytes().into()); - ctx.client.upload(data, len, &path1, None).await?; - - let (data, len) = upload_stream("remote blob data2".as_bytes().into()); - ctx.client.upload(data, len, &path2, None).await?; - - let (data, len) = upload_stream("remote blob data3".as_bytes().into()); - ctx.client.upload(data, len, &path3, None).await?; - - ctx.client.delete_objects(&[path1, path2]).await?; - - let prefixes = ctx.client.list_prefixes(None).await?; - - assert_eq!(prefixes.len(), 1); - - ctx.client.delete_objects(&[path3]).await?; - - Ok(()) -} - -#[test_context(MaybeEnabledAzure)] -#[tokio::test] -async fn azure_upload_download_works(ctx: &mut MaybeEnabledAzure) -> anyhow::Result<()> { - let MaybeEnabledAzure::Enabled(ctx) = ctx else { - return Ok(()); - }; - - let path = RemotePath::new(Utf8Path::new(format!("{}/file", ctx.base_prefix).as_str())) - .with_context(|| "RemotePath conversion")?; - - let orig = bytes::Bytes::from_static("remote blob data here".as_bytes()); - - let (data, len) = wrap_stream(orig.clone()); - - ctx.client.upload(data, len, &path, None).await?; - - // Normal download request - let dl = ctx.client.download(&path).await?; - let buf = download_to_vec(dl).await?; - assert_eq!(&buf, &orig); - - // Full range (end specified) - let dl = ctx - .client - .download_byte_range(&path, 0, Some(len as u64)) - .await?; - let buf = download_to_vec(dl).await?; - assert_eq!(&buf, &orig); - - // partial range (end specified) - let dl = ctx.client.download_byte_range(&path, 4, Some(10)).await?; - let buf = download_to_vec(dl).await?; - assert_eq!(&buf, &orig[4..10]); - - // partial range (end beyond real end) - let dl = ctx - .client - .download_byte_range(&path, 8, Some(len as u64 * 100)) - .await?; - let buf = download_to_vec(dl).await?; - assert_eq!(&buf, &orig[8..]); - - // Partial range (end unspecified) - let dl = ctx.client.download_byte_range(&path, 4, None).await?; - let buf = download_to_vec(dl).await?; - assert_eq!(&buf, &orig[4..]); - - // Full range (end unspecified) - let dl = ctx.client.download_byte_range(&path, 0, None).await?; - let buf = download_to_vec(dl).await?; - assert_eq!(&buf, &orig); - - debug!("Cleanup: deleting file at path {path:?}"); - ctx.client - .delete(&path) - .await - .with_context(|| format!("{path:?} removal"))?; - - Ok(()) -} - struct EnabledAzure { client: Arc, base_prefix: &'static str, @@ -281,13 +41,13 @@ impl EnabledAzure { } } -enum MaybeEnabledAzure { +enum MaybeEnabledStorage { Enabled(EnabledAzure), Disabled, } #[async_trait::async_trait] -impl AsyncTestContext for MaybeEnabledAzure { +impl AsyncTestContext for MaybeEnabledStorage { async fn setup() -> Self { ensure_logging_ready(); @@ -303,7 +63,7 @@ impl AsyncTestContext for MaybeEnabledAzure { } } -enum MaybeEnabledAzureWithTestBlobs { +enum MaybeEnabledStorageWithTestBlobs { Enabled(AzureWithTestBlobs), Disabled, UploadsFailed(anyhow::Error, AzureWithTestBlobs), @@ -316,7 +76,7 @@ struct AzureWithTestBlobs { } #[async_trait::async_trait] -impl AsyncTestContext for MaybeEnabledAzureWithTestBlobs { +impl AsyncTestContext for MaybeEnabledStorageWithTestBlobs { async fn setup() -> Self { ensure_logging_ready(); if env::var(ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME).is_err() { @@ -367,7 +127,7 @@ impl AsyncTestContext for MaybeEnabledAzureWithTestBlobs { // However, they are not idential. The list_prefixes function is concerned with listing prefixes, // whereas the list_files function is concerned with listing files. // See `RemoteStorage::list_files` documentation for more details -enum MaybeEnabledAzureWithSimpleTestBlobs { +enum MaybeEnabledStorageWithSimpleTestBlobs { Enabled(AzureWithSimpleTestBlobs), Disabled, UploadsFailed(anyhow::Error, AzureWithSimpleTestBlobs), @@ -378,7 +138,7 @@ struct AzureWithSimpleTestBlobs { } #[async_trait::async_trait] -impl AsyncTestContext for MaybeEnabledAzureWithSimpleTestBlobs { +impl AsyncTestContext for MaybeEnabledStorageWithSimpleTestBlobs { async fn setup() -> Self { ensure_logging_ready(); if env::var(ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME).is_err() { diff --git a/libs/remote_storage/tests/test_real_s3.rs b/libs/remote_storage/tests/test_real_s3.rs index 8f46b2abd6..4a999d115e 100644 --- a/libs/remote_storage/tests/test_real_s3.rs +++ b/libs/remote_storage/tests/test_real_s3.rs @@ -6,259 +6,23 @@ use std::sync::Arc; use std::time::UNIX_EPOCH; use anyhow::Context; -use camino::Utf8Path; use remote_storage::{ GenericRemoteStorage, RemotePath, RemoteStorageConfig, RemoteStorageKind, S3Config, }; -use test_context::{test_context, AsyncTestContext}; -use tracing::{debug, info}; +use test_context::AsyncTestContext; +use tracing::info; mod common; -use common::{ - cleanup, download_to_vec, ensure_logging_ready, upload_remote_data, upload_simple_remote_data, - upload_stream, wrap_stream, -}; +#[path = "common/tests.rs"] +mod tests_s3; + +use common::{cleanup, ensure_logging_ready, upload_remote_data, upload_simple_remote_data}; const ENABLE_REAL_S3_REMOTE_STORAGE_ENV_VAR_NAME: &str = "ENABLE_REAL_S3_REMOTE_STORAGE"; const BASE_PREFIX: &str = "test"; -/// Tests that S3 client can list all prefixes, even if the response come paginated and requires multiple S3 queries. -/// Uses real S3 and requires [`ENABLE_REAL_S3_REMOTE_STORAGE_ENV_VAR_NAME`] and related S3 cred env vars specified. -/// See the client creation in [`create_s3_client`] for details on the required env vars. -/// If real S3 tests are disabled, the test passes, skipping any real test run: currently, there's no way to mark the test ignored in runtime with the -/// deafult test framework, see https://github.com/rust-lang/rust/issues/68007 for details. -/// -/// First, the test creates a set of S3 objects with keys `/${random_prefix_part}/${base_prefix_str}/sub_prefix_${i}/blob_${i}` in [`upload_remote_data`] -/// where -/// * `random_prefix_part` is set for the entire S3 client during the S3 client creation in [`create_s3_client`], to avoid multiple test runs interference -/// * `base_prefix_str` is a common prefix to use in the client requests: we would want to ensure that the client is able to list nested prefixes inside the bucket -/// -/// Then, verifies that the client does return correct prefixes when queried: -/// * with no prefix, it lists everything after its `${random_prefix_part}/` — that should be `${base_prefix_str}` value only -/// * with `${base_prefix_str}/` prefix, it lists every `sub_prefix_${i}` -/// -/// With the real S3 enabled and `#[cfg(test)]` Rust configuration used, the S3 client test adds a `max-keys` param to limit the response keys. -/// This way, we are able to test the pagination implicitly, by ensuring all results are returned from the remote storage and avoid uploading too many blobs to S3, -/// since current default AWS S3 pagination limit is 1000. -/// (see https://docs.aws.amazon.com/AmazonS3/latest/API/API_ListObjectsV2.html#API_ListObjectsV2_RequestSyntax) -/// -/// Lastly, the test attempts to clean up and remove all uploaded S3 files. -/// If any errors appear during the clean up, they get logged, but the test is not failed or stopped until clean up is finished. -#[test_context(MaybeEnabledS3WithTestBlobs)] -#[tokio::test] -async fn s3_pagination_should_work(ctx: &mut MaybeEnabledS3WithTestBlobs) -> anyhow::Result<()> { - let ctx = match ctx { - MaybeEnabledS3WithTestBlobs::Enabled(ctx) => ctx, - MaybeEnabledS3WithTestBlobs::Disabled => return Ok(()), - MaybeEnabledS3WithTestBlobs::UploadsFailed(e, _) => anyhow::bail!("S3 init failed: {e:?}"), - }; - - let test_client = Arc::clone(&ctx.enabled.client); - let expected_remote_prefixes = ctx.remote_prefixes.clone(); - - let base_prefix = RemotePath::new(Utf8Path::new(ctx.enabled.base_prefix)) - .context("common_prefix construction")?; - let root_remote_prefixes = test_client - .list_prefixes(None) - .await - .context("client list root prefixes failure")? - .into_iter() - .collect::>(); - assert_eq!( - root_remote_prefixes, HashSet::from([base_prefix.clone()]), - "remote storage root prefixes list mismatches with the uploads. Returned prefixes: {root_remote_prefixes:?}" - ); - - let nested_remote_prefixes = test_client - .list_prefixes(Some(&base_prefix)) - .await - .context("client list nested prefixes failure")? - .into_iter() - .collect::>(); - let remote_only_prefixes = nested_remote_prefixes - .difference(&expected_remote_prefixes) - .collect::>(); - let missing_uploaded_prefixes = expected_remote_prefixes - .difference(&nested_remote_prefixes) - .collect::>(); - assert_eq!( - remote_only_prefixes.len() + missing_uploaded_prefixes.len(), 0, - "remote storage nested prefixes list mismatches with the uploads. Remote only prefixes: {remote_only_prefixes:?}, missing uploaded prefixes: {missing_uploaded_prefixes:?}", - ); - - Ok(()) -} - -/// Tests that S3 client can list all files in a folder, even if the response comes paginated and requirees multiple S3 queries. -/// Uses real S3 and requires [`ENABLE_REAL_S3_REMOTE_STORAGE_ENV_VAR_NAME`] and related S3 cred env vars specified. Test will skip real code and pass if env vars not set. -/// See `s3_pagination_should_work` for more information. -/// -/// First, create a set of S3 objects with keys `random_prefix/folder{j}/blob_{i}.txt` in [`upload_remote_data`] -/// Then performs the following queries: -/// 1. `list_files(None)`. This should return all files `random_prefix/folder{j}/blob_{i}.txt` -/// 2. `list_files("folder1")`. This should return all files `random_prefix/folder1/blob_{i}.txt` -#[test_context(MaybeEnabledS3WithSimpleTestBlobs)] -#[tokio::test] -async fn s3_list_files_works(ctx: &mut MaybeEnabledS3WithSimpleTestBlobs) -> anyhow::Result<()> { - let ctx = match ctx { - MaybeEnabledS3WithSimpleTestBlobs::Enabled(ctx) => ctx, - MaybeEnabledS3WithSimpleTestBlobs::Disabled => return Ok(()), - MaybeEnabledS3WithSimpleTestBlobs::UploadsFailed(e, _) => { - anyhow::bail!("S3 init failed: {e:?}") - } - }; - let test_client = Arc::clone(&ctx.enabled.client); - let base_prefix = - RemotePath::new(Utf8Path::new("folder1")).context("common_prefix construction")?; - let root_files = test_client - .list_files(None) - .await - .context("client list root files failure")? - .into_iter() - .collect::>(); - assert_eq!( - root_files, - ctx.remote_blobs.clone(), - "remote storage list_files on root mismatches with the uploads." - ); - let nested_remote_files = test_client - .list_files(Some(&base_prefix)) - .await - .context("client list nested files failure")? - .into_iter() - .collect::>(); - let trim_remote_blobs: HashSet<_> = ctx - .remote_blobs - .iter() - .map(|x| x.get_path()) - .filter(|x| x.starts_with("folder1")) - .map(|x| RemotePath::new(x).expect("must be valid path")) - .collect(); - assert_eq!( - nested_remote_files, trim_remote_blobs, - "remote storage list_files on subdirrectory mismatches with the uploads." - ); - Ok(()) -} - -#[test_context(MaybeEnabledS3)] -#[tokio::test] -async fn s3_delete_non_exising_works(ctx: &mut MaybeEnabledS3) -> anyhow::Result<()> { - let ctx = match ctx { - MaybeEnabledS3::Enabled(ctx) => ctx, - MaybeEnabledS3::Disabled => return Ok(()), - }; - - let path = RemotePath::new(Utf8Path::new( - format!("{}/for_sure_there_is_nothing_there_really", ctx.base_prefix).as_str(), - )) - .with_context(|| "RemotePath conversion")?; - - ctx.client.delete(&path).await.expect("should succeed"); - - Ok(()) -} - -#[test_context(MaybeEnabledS3)] -#[tokio::test] -async fn s3_delete_objects_works(ctx: &mut MaybeEnabledS3) -> anyhow::Result<()> { - let ctx = match ctx { - MaybeEnabledS3::Enabled(ctx) => ctx, - MaybeEnabledS3::Disabled => return Ok(()), - }; - - let path1 = RemotePath::new(Utf8Path::new(format!("{}/path1", ctx.base_prefix).as_str())) - .with_context(|| "RemotePath conversion")?; - - let path2 = RemotePath::new(Utf8Path::new(format!("{}/path2", ctx.base_prefix).as_str())) - .with_context(|| "RemotePath conversion")?; - - let path3 = RemotePath::new(Utf8Path::new(format!("{}/path3", ctx.base_prefix).as_str())) - .with_context(|| "RemotePath conversion")?; - - let (data, len) = upload_stream("remote blob data1".as_bytes().into()); - ctx.client.upload(data, len, &path1, None).await?; - - let (data, len) = upload_stream("remote blob data2".as_bytes().into()); - ctx.client.upload(data, len, &path2, None).await?; - - let (data, len) = upload_stream("remote blob data3".as_bytes().into()); - ctx.client.upload(data, len, &path3, None).await?; - - ctx.client.delete_objects(&[path1, path2]).await?; - - let prefixes = ctx.client.list_prefixes(None).await?; - - assert_eq!(prefixes.len(), 1); - - ctx.client.delete_objects(&[path3]).await?; - - Ok(()) -} - -#[test_context(MaybeEnabledS3)] -#[tokio::test] -async fn s3_upload_download_works(ctx: &mut MaybeEnabledS3) -> anyhow::Result<()> { - let MaybeEnabledS3::Enabled(ctx) = ctx else { - return Ok(()); - }; - - let path = RemotePath::new(Utf8Path::new(format!("{}/file", ctx.base_prefix).as_str())) - .with_context(|| "RemotePath conversion")?; - - let orig = bytes::Bytes::from_static("remote blob data here".as_bytes()); - - let (data, len) = wrap_stream(orig.clone()); - - ctx.client.upload(data, len, &path, None).await?; - - // Normal download request - let dl = ctx.client.download(&path).await?; - let buf = download_to_vec(dl).await?; - assert_eq!(&buf, &orig); - - // Full range (end specified) - let dl = ctx - .client - .download_byte_range(&path, 0, Some(len as u64)) - .await?; - let buf = download_to_vec(dl).await?; - assert_eq!(&buf, &orig); - - // partial range (end specified) - let dl = ctx.client.download_byte_range(&path, 4, Some(10)).await?; - let buf = download_to_vec(dl).await?; - assert_eq!(&buf, &orig[4..10]); - - // partial range (end beyond real end) - let dl = ctx - .client - .download_byte_range(&path, 8, Some(len as u64 * 100)) - .await?; - let buf = download_to_vec(dl).await?; - assert_eq!(&buf, &orig[8..]); - - // Partial range (end unspecified) - let dl = ctx.client.download_byte_range(&path, 4, None).await?; - let buf = download_to_vec(dl).await?; - assert_eq!(&buf, &orig[4..]); - - // Full range (end unspecified) - let dl = ctx.client.download_byte_range(&path, 0, None).await?; - let buf = download_to_vec(dl).await?; - assert_eq!(&buf, &orig); - - debug!("Cleanup: deleting file at path {path:?}"); - ctx.client - .delete(&path) - .await - .with_context(|| format!("{path:?} removal"))?; - - Ok(()) -} - struct EnabledS3 { client: Arc, base_prefix: &'static str, @@ -277,13 +41,13 @@ impl EnabledS3 { } } -enum MaybeEnabledS3 { +enum MaybeEnabledStorage { Enabled(EnabledS3), Disabled, } #[async_trait::async_trait] -impl AsyncTestContext for MaybeEnabledS3 { +impl AsyncTestContext for MaybeEnabledStorage { async fn setup() -> Self { ensure_logging_ready(); @@ -299,7 +63,7 @@ impl AsyncTestContext for MaybeEnabledS3 { } } -enum MaybeEnabledS3WithTestBlobs { +enum MaybeEnabledStorageWithTestBlobs { Enabled(S3WithTestBlobs), Disabled, UploadsFailed(anyhow::Error, S3WithTestBlobs), @@ -312,7 +76,7 @@ struct S3WithTestBlobs { } #[async_trait::async_trait] -impl AsyncTestContext for MaybeEnabledS3WithTestBlobs { +impl AsyncTestContext for MaybeEnabledStorageWithTestBlobs { async fn setup() -> Self { ensure_logging_ready(); if env::var(ENABLE_REAL_S3_REMOTE_STORAGE_ENV_VAR_NAME).is_err() { @@ -363,7 +127,7 @@ impl AsyncTestContext for MaybeEnabledS3WithTestBlobs { // However, they are not idential. The list_prefixes function is concerned with listing prefixes, // whereas the list_files function is concerned with listing files. // See `RemoteStorage::list_files` documentation for more details -enum MaybeEnabledS3WithSimpleTestBlobs { +enum MaybeEnabledStorageWithSimpleTestBlobs { Enabled(S3WithSimpleTestBlobs), Disabled, UploadsFailed(anyhow::Error, S3WithSimpleTestBlobs), @@ -374,7 +138,7 @@ struct S3WithSimpleTestBlobs { } #[async_trait::async_trait] -impl AsyncTestContext for MaybeEnabledS3WithSimpleTestBlobs { +impl AsyncTestContext for MaybeEnabledStorageWithSimpleTestBlobs { async fn setup() -> Self { ensure_logging_ready(); if env::var(ENABLE_REAL_S3_REMOTE_STORAGE_ENV_VAR_NAME).is_err() { diff --git a/libs/utils/src/id.rs b/libs/utils/src/id.rs index 57dcc27719..0409001f4f 100644 --- a/libs/utils/src/id.rs +++ b/libs/utils/src/id.rs @@ -1,3 +1,4 @@ +use std::num::ParseIntError; use std::{fmt, str::FromStr}; use anyhow::Context; @@ -374,6 +375,13 @@ impl fmt::Display for NodeId { } } +impl FromStr for NodeId { + type Err = ParseIntError; + fn from_str(s: &str) -> Result { + Ok(NodeId(u64::from_str(s)?)) + } +} + #[cfg(test)] mod tests { use serde_assert::{Deserializer, Serializer, Token, Tokens}; diff --git a/pageserver/client/src/mgmt_api.rs b/pageserver/client/src/mgmt_api.rs index adfd99e8f9..077c3909e1 100644 --- a/pageserver/client/src/mgmt_api.rs +++ b/pageserver/client/src/mgmt_api.rs @@ -1,5 +1,5 @@ use pageserver_api::{models::*, shard::TenantShardId}; -use reqwest::{IntoUrl, Method}; +use reqwest::{IntoUrl, Method, StatusCode}; use utils::{ http::error::HttpErrorBody, id::{TenantId, TimelineId}, @@ -22,14 +22,14 @@ pub enum Error { #[error("receive error body: {0}")] ReceiveErrorBody(String), - #[error("pageserver API: {0}")] - ApiError(String), + #[error("pageserver API: {1}")] + ApiError(StatusCode, String), } pub type Result = std::result::Result; -pub(crate) trait ResponseErrorMessageExt: Sized { - async fn error_from_body(self) -> Result; +pub trait ResponseErrorMessageExt: Sized { + fn error_from_body(self) -> impl std::future::Future> + Send; } impl ResponseErrorMessageExt for reqwest::Response { @@ -41,7 +41,7 @@ impl ResponseErrorMessageExt for reqwest::Response { let url = self.url().to_owned(); Err(match self.json::().await { - Ok(HttpErrorBody { msg }) => Error::ApiError(msg), + Ok(HttpErrorBody { msg }) => Error::ApiError(status, msg), Err(_) => { Error::ReceiveErrorBody(format!("Http error ({}) at {}.", status.as_u16(), url)) } @@ -71,9 +71,9 @@ impl Client { pub async fn tenant_details( &self, - tenant_id: TenantId, + tenant_shard_id: TenantShardId, ) -> Result { - let uri = format!("{}/v1/tenant/{tenant_id}", self.mgmt_api_endpoint); + let uri = format!("{}/v1/tenant/{tenant_shard_id}", self.mgmt_api_endpoint); self.get(uri) .await? .json() @@ -83,9 +83,12 @@ impl Client { pub async fn list_timelines( &self, - tenant_id: TenantId, + tenant_shard_id: TenantShardId, ) -> Result> { - let uri = format!("{}/v1/tenant/{tenant_id}/timeline", self.mgmt_api_endpoint); + let uri = format!( + "{}/v1/tenant/{tenant_shard_id}/timeline", + self.mgmt_api_endpoint + ); self.get(&uri) .await? .json() @@ -179,23 +182,23 @@ impl Client { "{}/v1/tenant/{}/secondary/download", self.mgmt_api_endpoint, tenant_id ); - self.request(Method::POST, &uri, ()) - .await? - .error_for_status() - .map(|_| ()) - .map_err(|e| Error::ApiError(format!("{}", e))) + self.request(Method::POST, &uri, ()).await?; + Ok(()) } pub async fn location_config( &self, - tenant_id: TenantId, + tenant_shard_id: TenantShardId, config: LocationConfig, flush_ms: Option, ) -> Result<()> { - let req_body = TenantLocationConfigRequest { tenant_id, config }; + let req_body = TenantLocationConfigRequest { + tenant_id: tenant_shard_id, + config, + }; let path = format!( "{}/v1/tenant/{}/location_config", - self.mgmt_api_endpoint, tenant_id + self.mgmt_api_endpoint, tenant_shard_id ); let path = if let Some(flush_ms) = flush_ms { format!("{}?flush_ms={}", path, flush_ms.as_millis()) @@ -206,14 +209,23 @@ impl Client { Ok(()) } + pub async fn list_location_config(&self) -> Result { + let path = format!("{}/v1/location_config", self.mgmt_api_endpoint); + self.request(Method::GET, &path, ()) + .await? + .json() + .await + .map_err(Error::ReceiveBody) + } + pub async fn timeline_create( &self, - tenant_id: TenantId, + tenant_shard_id: TenantShardId, req: &TimelineCreateRequest, ) -> Result { let uri = format!( "{}/v1/tenant/{}/timeline", - self.mgmt_api_endpoint, tenant_id + self.mgmt_api_endpoint, tenant_shard_id ); self.request(Method::POST, &uri, req) .await? @@ -233,4 +245,34 @@ impl Client { .await .map_err(Error::ReceiveBody) } + + pub async fn timeline_list( + &self, + tenant_shard_id: &TenantShardId, + ) -> Result> { + let uri = format!( + "{}/v1/tenant/{}/timeline", + self.mgmt_api_endpoint, tenant_shard_id + ); + self.get(&uri) + .await? + .json() + .await + .map_err(Error::ReceiveBody) + } + + pub async fn tenant_synthetic_size( + &self, + tenant_shard_id: TenantShardId, + ) -> Result { + let uri = format!( + "{}/v1/tenant/{}/synthetic_size", + self.mgmt_api_endpoint, tenant_shard_id + ); + self.get(&uri) + .await? + .json() + .await + .map_err(Error::ReceiveBody) + } } diff --git a/pageserver/client/src/mgmt_api/util.rs b/pageserver/client/src/mgmt_api/util.rs index 048a3bb7cd..bd85506d10 100644 --- a/pageserver/client/src/mgmt_api/util.rs +++ b/pageserver/client/src/mgmt_api/util.rs @@ -2,6 +2,7 @@ use std::sync::Arc; +use pageserver_api::shard::TenantShardId; use tokio::task::JoinSet; use utils::id::{TenantId, TenantTimelineId}; @@ -31,7 +32,10 @@ pub async fn get_pageserver_tenant_timelines_unsharded( async move { ( tenant_id, - mgmt_api_client.tenant_details(tenant_id).await.unwrap(), + mgmt_api_client + .tenant_details(TenantShardId::unsharded(tenant_id)) + .await + .unwrap(), ) } }); diff --git a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs index 46b8bd10d7..d8957ddd6b 100644 --- a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs +++ b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs @@ -349,10 +349,10 @@ async fn main_impl( let work_sender_task = tokio::spawn(work_sender); + info!("waiting for everything to become ready"); + start_work_barrier.wait().await; + info!("work started"); if let Some(runtime) = args.runtime { - info!("waiting for everything to become ready"); - start_work_barrier.wait().await; - info!("work started"); tokio::time::sleep(runtime.into()).await; info!("runtime over, signalling cancellation"); cancel.cancel(); diff --git a/pageserver/src/bin/pageserver.rs b/pageserver/src/bin/pageserver.rs index 621ad050f4..15e3359c06 100644 --- a/pageserver/src/bin/pageserver.rs +++ b/pageserver/src/bin/pageserver.rs @@ -527,6 +527,7 @@ fn start_pageserver( conf, remote_storage.clone(), disk_usage_eviction_state.clone(), + tenant_manager.clone(), background_jobs_barrier.clone(), )?; } diff --git a/pageserver/src/config.rs b/pageserver/src/config.rs index 7c03dc1bdd..52277d7f24 100644 --- a/pageserver/src/config.rs +++ b/pageserver/src/config.rs @@ -1126,11 +1126,12 @@ mod tests { }; use camino_tempfile::{tempdir, Utf8TempDir}; + use pageserver_api::models::EvictionPolicy; use remote_storage::{RemoteStorageKind, S3Config}; use utils::serde_percent::Percent; use super::*; - use crate::{tenant::config::EvictionPolicy, DEFAULT_PG_VERSION}; + use crate::DEFAULT_PG_VERSION; const ALL_BASE_VALUES_TOML: &str = r#" # Initial configuration file created by 'pageserver --init' diff --git a/pageserver/src/disk_usage_eviction_task.rs b/pageserver/src/disk_usage_eviction_task.rs index 7d74da1196..4cb20a1bb1 100644 --- a/pageserver/src/disk_usage_eviction_task.rs +++ b/pageserver/src/disk_usage_eviction_task.rs @@ -47,21 +47,24 @@ use std::{ }; use anyhow::Context; -use camino::Utf8Path; +use pageserver_api::shard::TenantShardId; use remote_storage::GenericRemoteStorage; use serde::{Deserialize, Serialize}; use tokio::time::Instant; use tokio_util::sync::CancellationToken; use tracing::{debug, error, info, instrument, warn, Instrument}; -use utils::completion; use utils::serde_percent::Percent; +use utils::{completion, id::TimelineId}; use crate::{ config::PageServerConf, task_mgr::{self, TaskKind, BACKGROUND_RUNTIME}, tenant::{ self, - storage_layer::{AsLayerDesc, EvictionError, Layer}, + mgr::TenantManager, + remote_timeline_client::LayerFileMetadata, + secondary::SecondaryTenant, + storage_layer::{AsLayerDesc, EvictionError, Layer, LayerFileName}, Timeline, }, }; @@ -125,6 +128,7 @@ pub fn launch_disk_usage_global_eviction_task( conf: &'static PageServerConf, storage: GenericRemoteStorage, state: Arc, + tenant_manager: Arc, background_jobs_barrier: completion::Barrier, ) -> anyhow::Result<()> { let Some(task_config) = &conf.disk_usage_based_eviction else { @@ -150,8 +154,7 @@ pub fn launch_disk_usage_global_eviction_task( _ = background_jobs_barrier.wait() => { } }; - disk_usage_eviction_task(&state, task_config, &storage, &conf.tenants_path(), cancel) - .await; + disk_usage_eviction_task(&state, task_config, &storage, tenant_manager, cancel).await; Ok(()) }, ); @@ -164,7 +167,7 @@ async fn disk_usage_eviction_task( state: &State, task_config: &DiskUsageEvictionTaskConfig, storage: &GenericRemoteStorage, - tenants_dir: &Utf8Path, + tenant_manager: Arc, cancel: CancellationToken, ) { scopeguard::defer! { @@ -191,7 +194,7 @@ async fn disk_usage_eviction_task( state, task_config, storage, - tenants_dir, + &tenant_manager, &cancel, ) .await; @@ -226,15 +229,17 @@ async fn disk_usage_eviction_task_iteration( state: &State, task_config: &DiskUsageEvictionTaskConfig, storage: &GenericRemoteStorage, - tenants_dir: &Utf8Path, + tenant_manager: &Arc, cancel: &CancellationToken, ) -> anyhow::Result<()> { - let usage_pre = filesystem_level_usage::get(tenants_dir, task_config) + let tenants_dir = tenant_manager.get_conf().tenants_path(); + let usage_pre = filesystem_level_usage::get(&tenants_dir, task_config) .context("get filesystem-level disk usage before evictions")?; let res = disk_usage_eviction_task_iteration_impl( state, storage, usage_pre, + tenant_manager, task_config.eviction_order, cancel, ) @@ -248,7 +253,7 @@ async fn disk_usage_eviction_task_iteration( } IterationOutcome::Finished(outcome) => { // Verify with statvfs whether we made any real progress - let after = filesystem_level_usage::get(tenants_dir, task_config) + let after = filesystem_level_usage::get(&tenants_dir, task_config) // It's quite unlikely to hit the error here. Keep the code simple and bail out. .context("get filesystem-level disk usage after evictions")?; @@ -324,6 +329,7 @@ pub(crate) async fn disk_usage_eviction_task_iteration_impl( state: &State, _storage: &GenericRemoteStorage, usage_pre: U, + tenant_manager: &Arc, eviction_order: EvictionOrder, cancel: &CancellationToken, ) -> anyhow::Result> { @@ -344,29 +350,29 @@ pub(crate) async fn disk_usage_eviction_task_iteration_impl( "running disk usage based eviction due to pressure" ); - let candidates = match collect_eviction_candidates(eviction_order, cancel).await? { - EvictionCandidates::Cancelled => { - return Ok(IterationOutcome::Cancelled); - } - EvictionCandidates::Finished(partitioned) => partitioned, - }; + let candidates = + match collect_eviction_candidates(tenant_manager, eviction_order, cancel).await? { + EvictionCandidates::Cancelled => { + return Ok(IterationOutcome::Cancelled); + } + EvictionCandidates::Finished(partitioned) => partitioned, + }; // Debug-log the list of candidates let now = SystemTime::now(); for (i, (partition, candidate)) in candidates.iter().enumerate() { let nth = i + 1; - let desc = candidate.layer.layer_desc(); let total_candidates = candidates.len(); - let size = desc.file_size; + let size = candidate.layer.get_file_size(); let rel = candidate.relative_last_activity; debug!( "cand {nth}/{total_candidates}: size={size}, rel_last_activity={rel}, no_access_for={}us, partition={partition:?}, {}/{}/{}", now.duration_since(candidate.last_activity_ts) .unwrap() .as_micros(), - desc.tenant_shard_id, - desc.timeline_id, - candidate.layer, + candidate.layer.get_tenant_shard_id(), + candidate.layer.get_timeline_id(), + candidate.layer.get_name(), ); } @@ -398,7 +404,7 @@ pub(crate) async fn disk_usage_eviction_task_iteration_impl( warned = Some(usage_planned); } - usage_planned.add_available_bytes(candidate.layer.layer_desc().file_size); + usage_planned.add_available_bytes(candidate.layer.get_file_size()); evicted_amount += 1; } @@ -463,19 +469,30 @@ pub(crate) async fn disk_usage_eviction_task_iteration_impl( continue; }; - js.spawn(async move { - let rtc = candidate.timeline.remote_client.as_ref().expect( - "holding the witness, all timelines must have a remote timeline client", - ); - let file_size = candidate.layer.layer_desc().file_size; - candidate - .layer - .evict_and_wait(rtc) - .await - .map(|()| file_size) - .map_err(|e| (file_size, e)) - }); + match candidate.layer { + EvictionLayer::Attached(layer) => { + let file_size = layer.layer_desc().file_size; + js.spawn(async move { + layer + .evict_and_wait() + .await + .map(|()| file_size) + .map_err(|e| (file_size, e)) + }); + } + EvictionLayer::Secondary(layer) => { + let file_size = layer.metadata.file_size(); + let tenant_manager = tenant_manager.clone(); + js.spawn(async move { + layer + .secondary_tenant + .evict_layer(tenant_manager.get_conf(), layer.timeline_id, layer.name) + .await; + Ok(file_size) + }); + } + } tokio::task::yield_now().await; } @@ -502,11 +519,100 @@ pub(crate) async fn disk_usage_eviction_task_iteration_impl( } #[derive(Clone)] -struct EvictionCandidate { - timeline: Arc, - layer: Layer, - last_activity_ts: SystemTime, - relative_last_activity: finite_f32::FiniteF32, +pub(crate) struct EvictionSecondaryLayer { + pub(crate) secondary_tenant: Arc, + pub(crate) timeline_id: TimelineId, + pub(crate) name: LayerFileName, + pub(crate) metadata: LayerFileMetadata, +} + +/// Full [`Layer`] objects are specific to tenants in attached mode. This type is a layer +/// of indirection to store either a `Layer`, or a reference to a secondary tenant and a layer name. +#[derive(Clone)] +pub(crate) enum EvictionLayer { + Attached(Layer), + #[allow(dead_code)] + Secondary(EvictionSecondaryLayer), +} + +impl From for EvictionLayer { + fn from(value: Layer) -> Self { + Self::Attached(value) + } +} + +impl EvictionLayer { + pub(crate) fn get_tenant_shard_id(&self) -> &TenantShardId { + match self { + Self::Attached(l) => &l.layer_desc().tenant_shard_id, + Self::Secondary(sl) => sl.secondary_tenant.get_tenant_shard_id(), + } + } + + pub(crate) fn get_timeline_id(&self) -> &TimelineId { + match self { + Self::Attached(l) => &l.layer_desc().timeline_id, + Self::Secondary(sl) => &sl.timeline_id, + } + } + + pub(crate) fn get_name(&self) -> LayerFileName { + match self { + Self::Attached(l) => l.layer_desc().filename(), + Self::Secondary(sl) => sl.name.clone(), + } + } + + pub(crate) fn get_file_size(&self) -> u64 { + match self { + Self::Attached(l) => l.layer_desc().file_size, + Self::Secondary(sl) => sl.metadata.file_size(), + } + } +} + +#[derive(Clone)] +pub(crate) struct EvictionCandidate { + pub(crate) layer: EvictionLayer, + pub(crate) last_activity_ts: SystemTime, + pub(crate) relative_last_activity: finite_f32::FiniteF32, +} + +impl std::fmt::Display for EvictionLayer { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match self { + Self::Attached(l) => l.fmt(f), + Self::Secondary(sl) => { + write!(f, "{}/{}", sl.timeline_id, sl.name) + } + } + } +} + +pub(crate) struct DiskUsageEvictionInfo { + /// Timeline's largest layer (remote or resident) + pub max_layer_size: Option, + /// Timeline's resident layers + pub resident_layers: Vec, +} + +impl std::fmt::Debug for EvictionCandidate { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // format the tv_sec, tv_nsec into rfc3339 in case someone is looking at it + // having to allocate a string to this is bad, but it will rarely be formatted + let ts = chrono::DateTime::::from(self.last_activity_ts); + let ts = ts.to_rfc3339_opts(chrono::SecondsFormat::Nanos, true); + struct DisplayIsDebug<'a, T>(&'a T); + impl<'a, T: std::fmt::Display> std::fmt::Debug for DisplayIsDebug<'a, T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } + } + f.debug_struct("LocalLayerInfoForDiskUsageEviction") + .field("layer", &DisplayIsDebug(&self.layer)) + .field("last_activity", &ts) + .finish() + } } #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] @@ -623,6 +729,7 @@ enum EvictionCandidates { /// - tenant B 1 layer /// - tenant C 8 layers async fn collect_eviction_candidates( + tenant_manager: &Arc, eviction_order: EvictionOrder, cancel: &CancellationToken, ) -> anyhow::Result { @@ -631,13 +738,16 @@ async fn collect_eviction_candidates( .await .context("get list of tenants")?; + // TODO: avoid listing every layer in every tenant: this loop can block the executor, + // and the resulting data structure can be huge. + // (https://github.com/neondatabase/neon/issues/6224) let mut candidates = Vec::new(); - for (tenant_id, _state, _gen) in &tenants { + for (tenant_id, _state, _gen) in tenants { if cancel.is_cancelled() { return Ok(EvictionCandidates::Cancelled); } - let tenant = match tenant::mgr::get_tenant(*tenant_id, true) { + let tenant = match tenant::mgr::get_tenant(tenant_id, true) { Ok(tenant) => tenant, Err(e) => { // this can happen if tenant has lifecycle transition after we fetched it @@ -665,11 +775,7 @@ async fn collect_eviction_candidates( } let info = tl.get_local_layers_for_disk_usage_eviction().await; debug!(tenant_id=%tl.tenant_shard_id.tenant_id, shard_id=%tl.tenant_shard_id.shard_slug(), timeline_id=%tl.timeline_id, "timeline resident layers count: {}", info.resident_layers.len()); - tenant_candidates.extend( - info.resident_layers - .into_iter() - .map(|layer_infos| (tl.clone(), layer_infos)), - ); + tenant_candidates.extend(info.resident_layers.into_iter()); max_layer_size = max_layer_size.max(info.max_layer_size.unwrap_or(0)); if cancel.is_cancelled() { @@ -707,7 +813,7 @@ async fn collect_eviction_candidates( // Sort layers most-recently-used first, then partition by // cumsum above/below min_resident_size. tenant_candidates - .sort_unstable_by_key(|(_, layer_info)| std::cmp::Reverse(layer_info.last_activity_ts)); + .sort_unstable_by_key(|layer_info| std::cmp::Reverse(layer_info.last_activity_ts)); let mut cumsum: i128 = 0; // keeping the -1 or not decides if every tenant should lose their least recently accessed @@ -741,12 +847,10 @@ async fn collect_eviction_candidates( .unwrap_or(1); let divider = total as f32; - for (i, (timeline, layer_info)) in tenant_candidates.into_iter().enumerate() { - let file_size = layer_info.file_size(); - + for (i, mut candidate) in tenant_candidates.into_iter().enumerate() { // as we iterate this reverse sorted list, the most recently accessed layer will always // be 1.0; this is for us to evict it last. - let relative_last_activity = if matches!( + candidate.relative_last_activity = if matches!( eviction_order, EvictionOrder::RelativeAccessed { .. } ) { @@ -761,22 +865,46 @@ async fn collect_eviction_candidates( finite_f32::FiniteF32::ZERO }; - let candidate = EvictionCandidate { - timeline, - last_activity_ts: layer_info.last_activity_ts, - layer: layer_info.layer, - relative_last_activity, - }; let partition = if cumsum > min_resident_size as i128 { MinResidentSizePartition::Above } else { MinResidentSizePartition::Below }; + cumsum += i128::from(candidate.layer.get_file_size()); candidates.push((partition, candidate)); - cumsum += i128::from(file_size); } } + // Note: the same tenant ID might be hit twice, if it transitions from attached to + // secondary while we run. That is okay: when we eventually try and run the eviction, + // the `Gate` on the object will ensure that whichever one has already been shut down + // will not delete anything. + + let mut secondary_tenants = Vec::new(); + tenant_manager.foreach_secondary_tenants( + |_tenant_shard_id: &TenantShardId, state: &Arc| { + secondary_tenants.push(state.clone()); + }, + ); + + for secondary_tenant in secondary_tenants { + let mut layer_info = secondary_tenant.get_layers_for_eviction(); + + layer_info + .resident_layers + .sort_unstable_by_key(|layer_info| std::cmp::Reverse(layer_info.last_activity_ts)); + + candidates.extend(layer_info.resident_layers.into_iter().map(|candidate| { + ( + // Secondary locations' layers are always considered above the min resident size, + // i.e. secondary locations are permitted to be trimmed to zero layers if all + // the layers have sufficiently old access times. + MinResidentSizePartition::Above, + candidate, + ) + })); + } + debug_assert!(MinResidentSizePartition::Above < MinResidentSizePartition::Below, "as explained in the function's doc comment, layers that aren't in the tenant's min_resident_size are evicted first"); @@ -821,7 +949,7 @@ impl std::ops::Deref for TimelineKey { } /// A totally ordered f32 subset we can use with sorting functions. -mod finite_f32 { +pub(crate) mod finite_f32 { /// A totally ordered f32 subset we can use with sorting functions. #[derive(Clone, Copy, PartialEq)] diff --git a/pageserver/src/http/routes.rs b/pageserver/src/http/routes.rs index 5db6119a09..0a18ac4af1 100644 --- a/pageserver/src/http/routes.rs +++ b/pageserver/src/http/routes.rs @@ -14,6 +14,8 @@ use hyper::header; use hyper::StatusCode; use hyper::{Body, Request, Response, Uri}; use metrics::launch_timestamp::LaunchTimestamp; +use pageserver_api::models::LocationConfigListResponse; +use pageserver_api::models::ShardParameters; use pageserver_api::models::TenantDetails; use pageserver_api::models::TenantState; use pageserver_api::models::{ @@ -38,11 +40,11 @@ use crate::pgdatadir_mapping::LsnForTimestamp; use crate::task_mgr::TaskKind; use crate::tenant::config::{LocationConf, TenantConfOpt}; use crate::tenant::mgr::GetActiveTenantError; -use crate::tenant::mgr::UpsertLocationError; use crate::tenant::mgr::{ GetTenantError, SetNewTenantConfigError, TenantManager, TenantMapError, TenantMapInsertError, TenantSlotError, TenantSlotUpsertError, TenantStateError, }; +use crate::tenant::mgr::{TenantSlot, UpsertLocationError}; use crate::tenant::secondary::SecondaryController; use crate::tenant::size::ModelInputs; use crate::tenant::storage_layer::LayerAccessStatsReset; @@ -265,7 +267,7 @@ impl From for ApiError { SetNewTenantConfigError::GetTenant(tid) => { ApiError::NotFound(anyhow!("tenant {}", tid).into()) } - e @ SetNewTenantConfigError::Persist(_) => { + e @ (SetNewTenantConfigError::Persist(_) | SetNewTenantConfigError::Other(_)) => { ApiError::InternalServerError(anyhow::Error::new(e)) } } @@ -704,7 +706,9 @@ async fn tenant_attach_handler( } let tenant_shard_id = TenantShardId::unsharded(tenant_id); - let location_conf = LocationConf::attached_single(tenant_conf, generation); + let shard_params = ShardParameters::default(); + let location_conf = LocationConf::attached_single(tenant_conf, generation, &shard_params); + let tenant = state .tenant_manager .upsert_location( @@ -1194,7 +1198,8 @@ async fn tenant_create_handler( let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Warn); - let location_conf = LocationConf::attached_single(tenant_conf, generation); + let location_conf = + LocationConf::attached_single(tenant_conf, generation, &request_data.shard_parameters); let new_tenant = state .tenant_manager @@ -1213,7 +1218,6 @@ async fn tenant_create_handler( "Upsert succeeded but didn't return tenant!" ))); }; - // We created the tenant. Existing API semantics are that the tenant // is Active when this function returns. if let res @ Err(_) = new_tenant @@ -1351,6 +1355,28 @@ async fn put_tenant_location_config_handler( json_response(StatusCode::OK, ()) } +async fn list_location_config_handler( + request: Request, + _cancel: CancellationToken, +) -> Result, ApiError> { + let state = get_state(&request); + let slots = state.tenant_manager.list(); + let result = LocationConfigListResponse { + tenant_shards: slots + .into_iter() + .map(|(tenant_shard_id, slot)| { + let v = match slot { + TenantSlot::Attached(t) => Some(t.get_location_conf()), + TenantSlot::Secondary(s) => Some(s.get_location_conf()), + TenantSlot::InProgress(_) => None, + }; + (tenant_shard_id, v) + }) + .collect(), + }; + json_response(StatusCode::OK, result) +} + /// Testing helper to transition a tenant to [`crate::tenant::TenantState::Broken`]. async fn handle_tenant_break( r: Request, @@ -1652,12 +1678,13 @@ async fn disk_usage_eviction_run( ))); }; - let state = state.disk_usage_eviction_state.clone(); + let eviction_state = state.disk_usage_eviction_state.clone(); let res = crate::disk_usage_eviction_task::disk_usage_eviction_task_iteration_impl( - &state, + &eviction_state, storage, usage, + &state.tenant_manager, config.eviction_order, &cancel, ) @@ -1892,6 +1919,9 @@ pub fn make_router( .put("/v1/tenant/:tenant_shard_id/location_config", |r| { api_handler(r, put_tenant_location_config_handler) }) + .get("/v1/location_config", |r| { + api_handler(r, list_location_config_handler) + }) .get("/v1/tenant/:tenant_shard_id/timeline", |r| { api_handler(r, timeline_list_handler) }) diff --git a/pageserver/src/metrics.rs b/pageserver/src/metrics.rs index 6f4431c3cf..07f9049ca5 100644 --- a/pageserver/src/metrics.rs +++ b/pageserver/src/metrics.rs @@ -337,15 +337,6 @@ pub(crate) mod page_cache_eviction_metrics { } } -pub(crate) static PAGE_CACHE_ACQUIRE_PINNED_SLOT_TIME: Lazy = Lazy::new(|| { - register_histogram!( - "pageserver_page_cache_acquire_pinned_slot_seconds", - "Time spent acquiring a pinned slot in the page cache", - CRITICAL_OP_BUCKETS.into(), - ) - .expect("failed to define a metric") -}); - static PAGE_CACHE_ERRORS: Lazy = Lazy::new(|| { register_int_counter_vec!( "page_cache_errors_total", diff --git a/pageserver/src/page_cache.rs b/pageserver/src/page_cache.rs index c3c98af406..28d2584bf4 100644 --- a/pageserver/src/page_cache.rs +++ b/pageserver/src/page_cache.rs @@ -550,7 +550,6 @@ impl PageCache { // not require changes. async fn try_get_pinned_slot_permit(&self) -> anyhow::Result { - let timer = crate::metrics::PAGE_CACHE_ACQUIRE_PINNED_SLOT_TIME.start_timer(); match tokio::time::timeout( // Choose small timeout, neon_smgr does its own retries. // https://neondb.slack.com/archives/C04DGM6SMTM/p1694786876476869 @@ -563,7 +562,6 @@ impl PageCache { res.expect("this semaphore is never closed"), )), Err(_timeout) => { - timer.stop_and_discard(); crate::metrics::page_cache_errors_inc( crate::metrics::PageCacheErrorKind::AcquirePinnedSlotTimeout, ); diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs index 9b4b333a92..a296dde07d 100644 --- a/pageserver/src/page_service.rs +++ b/pageserver/src/page_service.rs @@ -13,7 +13,10 @@ use anyhow::Context; use async_compression::tokio::write::GzipEncoder; use bytes::Buf; use bytes::Bytes; +use futures::stream::FuturesUnordered; use futures::Stream; +use futures::StreamExt; +use pageserver_api::key::Key; use pageserver_api::models::TenantState; use pageserver_api::models::{ PagestreamBeMessage, PagestreamDbSizeRequest, PagestreamDbSizeResponse, @@ -21,11 +24,14 @@ use pageserver_api::models::{ PagestreamFeMessage, PagestreamGetPageRequest, PagestreamGetPageResponse, PagestreamNblocksRequest, PagestreamNblocksResponse, }; +use pageserver_api::shard::ShardIndex; +use pageserver_api::shard::{ShardCount, ShardNumber}; use postgres_backend::{self, is_expected_io_error, AuthType, PostgresBackend, QueryError}; use pq_proto::framed::ConnectionError; use pq_proto::FeStartupPacket; use pq_proto::{BeMessage, FeMessage, RowDescriptor}; use std::borrow::Cow; +use std::collections::HashMap; use std::io; use std::net::TcpListener; use std::pin::pin; @@ -40,6 +46,7 @@ use tokio_util::sync::CancellationToken; use tracing::field; use tracing::*; use utils::id::ConnectionId; +use utils::sync::gate::GateGuard; use utils::{ auth::{Claims, Scope, SwappableJwtAuth}, id::{TenantId, TimelineId}, @@ -274,6 +281,13 @@ async fn page_service_conn_main( } } +/// While a handler holds a reference to a Timeline, it also holds a the +/// timeline's Gate open. +struct HandlerTimeline { + timeline: Arc, + _guard: GateGuard, +} + struct PageServerHandler { _conf: &'static PageServerConf, broker_client: storage_broker::BrokerClientChannel, @@ -285,6 +299,14 @@ struct PageServerHandler { /// For each query received over the connection, /// `process_query` creates a child context from this one. connection_ctx: RequestContext, + + /// See [`Self::cache_timeline`] for usage. + /// + /// Note on size: the typical size of this map is 1. The largest size we expect + /// to see is the number of shards divided by the number of pageservers (typically < 2), + /// or the ratio used when splitting shards (i.e. how many children created from one) + /// parent shard, where a "large" number might be ~8. + shard_timelines: HashMap, } #[derive(thiserror::Error, Debug)] @@ -358,13 +380,46 @@ impl PageServerHandler { auth, claims: None, connection_ctx, + shard_timelines: HashMap::new(), } } - /// Wrap PostgresBackend::flush to respect our CancellationToken: it is important to use - /// this rather than naked flush() in order to shut down promptly. Without this, we would - /// block shutdown of a tenant if a postgres client was failing to consume bytes we send - /// in the flush. + /// Analogous to calling cancelled() on a Timeline's cancellation token: waits for cancellation. + /// + /// We use many Timeline objects, and hold GateGuards on all of them. We must therefore respect + /// all of their cancellation tokens. + async fn timeline_cancelled(&self) { + // A short wait before we expend the cycles to walk our timeline map. This avoids incurring + // that cost every time we check for cancellation. + tokio::time::sleep(Duration::from_millis(10)).await; + + // This function is never called concurrently with code that adds timelines to shard_timelines, + // which is enforced by the borrow checker (the future returned by this function carries the + // immutable &self). So it's fine to evaluate shard_timelines after the sleep, we don't risk + // missing any inserts to the map. + + let mut futs = self + .shard_timelines + .values() + .map(|ht| ht.timeline.cancel.cancelled()) + .collect::>(); + + futs.next().await; + } + + /// Analogous to calling is_cancelled() on a Timeline's cancellation token + fn timeline_is_cancelled(&self) -> bool { + self.shard_timelines + .values() + .any(|ht| ht.timeline.cancel.is_cancelled() || ht.timeline.is_stopping()) + } + + /// This function always respects cancellation of any timeline in `[Self::shard_timelines]`. Pass in + /// a cancellation token at the next scope up (such as a tenant cancellation token) to ensure we respect + /// cancellation if there aren't any timelines in the cache. + /// + /// If calling from a function that doesn't use the `[Self::shard_timelines]` cache, then pass in the + /// timeline cancellation token. async fn flush_cancellable( &self, pgb: &mut PostgresBackend, @@ -377,6 +432,9 @@ impl PageServerHandler { flush_r = pgb.flush() => { Ok(flush_r?) }, + _ = self.timeline_cancelled() => { + Err(QueryError::Shutdown) + } _ = cancel.cancelled() => { Err(QueryError::Shutdown) } @@ -452,7 +510,7 @@ impl PageServerHandler { #[instrument(skip_all)] async fn handle_pagerequests( - &self, + &mut self, pgb: &mut PostgresBackend, tenant_id: TenantId, timeline_id: TimelineId, @@ -463,10 +521,6 @@ impl PageServerHandler { { debug_assert_current_span_has_tenant_and_timeline_id(); - // Note that since one connection may contain getpage requests that target different - // shards (e.g. during splitting when the compute is not yet aware of the split), the tenant - // that we look up here may not be the one that serves all the actual requests: we will double - // check the mapping of key->shard later before calling into Timeline for getpage requests. let tenant = mgr::get_active_tenant_with_timeout( tenant_id, ShardSelector::First, @@ -487,19 +541,9 @@ impl PageServerHandler { None }; - // Check that the timeline exists - let timeline = tenant - .get_timeline(timeline_id, true) - .map_err(|e| QueryError::NotFound(format!("{e}").into()))?; - - // Avoid starting new requests if the timeline has already started shutting down, - // and block timeline shutdown until this request is complete, or drops out due - // to cancellation. - let _timeline_guard = timeline.gate.enter().map_err(|_| QueryError::Shutdown)?; - // switch client to COPYBOTH pgb.write_message_noflush(&BeMessage::CopyBothResponse)?; - self.flush_cancellable(pgb, &timeline.cancel).await?; + self.flush_cancellable(pgb, &tenant.cancel).await?; let metrics = metrics::SmgrQueryTimePerTimeline::new(&tenant_id, &timeline_id); @@ -507,7 +551,7 @@ impl PageServerHandler { let msg = tokio::select! { biased; - _ = timeline.cancel.cancelled() => { + _ = self.timeline_cancelled() => { // We were requested to shut down. info!("shutdown request received in page handler"); return Err(QueryError::Shutdown) @@ -544,7 +588,7 @@ impl PageServerHandler { let _timer = metrics.start_timer(metrics::SmgrQueryType::GetRelExists); let span = tracing::info_span!("handle_get_rel_exists_request", rel = %req.rel, req_lsn = %req.lsn); ( - self.handle_get_rel_exists_request(&timeline, &req, &ctx) + self.handle_get_rel_exists_request(tenant_id, timeline_id, &req, &ctx) .instrument(span.clone()) .await, span, @@ -554,7 +598,7 @@ impl PageServerHandler { let _timer = metrics.start_timer(metrics::SmgrQueryType::GetRelSize); let span = tracing::info_span!("handle_get_nblocks_request", rel = %req.rel, req_lsn = %req.lsn); ( - self.handle_get_nblocks_request(&timeline, &req, &ctx) + self.handle_get_nblocks_request(tenant_id, timeline_id, &req, &ctx) .instrument(span.clone()) .await, span, @@ -564,7 +608,7 @@ impl PageServerHandler { let _timer = metrics.start_timer(metrics::SmgrQueryType::GetPageAtLsn); let span = tracing::info_span!("handle_get_page_at_lsn_request", rel = %req.rel, blkno = %req.blkno, req_lsn = %req.lsn); ( - self.handle_get_page_at_lsn_request(&timeline, &req, &ctx) + self.handle_get_page_at_lsn_request(tenant_id, timeline_id, &req, &ctx) .instrument(span.clone()) .await, span, @@ -574,7 +618,7 @@ impl PageServerHandler { let _timer = metrics.start_timer(metrics::SmgrQueryType::GetDbSize); let span = tracing::info_span!("handle_db_size_request", dbnode = %req.dbnode, req_lsn = %req.lsn); ( - self.handle_db_size_request(&timeline, &req, &ctx) + self.handle_db_size_request(tenant_id, timeline_id, &req, &ctx) .instrument(span.clone()) .await, span, @@ -594,7 +638,7 @@ impl PageServerHandler { span.in_scope(|| info!("handler requested reconnect: {reason}")); return Err(QueryError::Reconnect); } - Err(e) if timeline.cancel.is_cancelled() || timeline.is_stopping() => { + Err(e) if self.timeline_is_cancelled() => { // This branch accomodates code within request handlers that returns an anyhow::Error instead of a clean // shutdown error, this may be buried inside a PageReconstructError::Other for example. // @@ -617,7 +661,7 @@ impl PageServerHandler { }); pgb.write_message_noflush(&BeMessage::CopyData(&response_msg.serialize()))?; - self.flush_cancellable(pgb, &timeline.cancel).await?; + self.flush_cancellable(pgb, &tenant.cancel).await?; } } } @@ -814,11 +858,14 @@ impl PageServerHandler { } async fn handle_get_rel_exists_request( - &self, - timeline: &Timeline, + &mut self, + tenant_id: TenantId, + timeline_id: TimelineId, req: &PagestreamExistsRequest, ctx: &RequestContext, ) -> Result { + let timeline = self.get_timeline_shard_zero(tenant_id, timeline_id).await?; + let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn(); let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn, ctx) @@ -834,11 +881,13 @@ impl PageServerHandler { } async fn handle_get_nblocks_request( - &self, - timeline: &Timeline, + &mut self, + tenant_id: TenantId, + timeline_id: TimelineId, req: &PagestreamNblocksRequest, ctx: &RequestContext, ) -> Result { + let timeline = self.get_timeline_shard_zero(tenant_id, timeline_id).await?; let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn(); let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn, ctx) @@ -854,11 +903,13 @@ impl PageServerHandler { } async fn handle_db_size_request( - &self, - timeline: &Timeline, + &mut self, + tenant_id: TenantId, + timeline_id: TimelineId, req: &PagestreamDbSizeRequest, ctx: &RequestContext, ) -> Result { + let timeline = self.get_timeline_shard_zero(tenant_id, timeline_id).await?; let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn(); let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn, ctx) @@ -880,16 +931,160 @@ impl PageServerHandler { })) } - async fn do_handle_get_page_at_lsn_request( - &self, - timeline: &Timeline, + /// For most getpage requests, we will already have a Timeline to serve the request: this function + /// looks up such a Timeline synchronously and without touching any global state. + fn get_cached_timeline_for_page( + &mut self, + req: &PagestreamGetPageRequest, + ) -> Result<&Arc, Key> { + let key = if let Some((first_idx, first_timeline)) = self.shard_timelines.iter().next() { + // Fastest path: single sharded case + if first_idx.shard_count < ShardCount(2) { + return Ok(&first_timeline.timeline); + } + + let key = rel_block_to_key(req.rel, req.blkno); + let shard_num = first_timeline + .timeline + .get_shard_identity() + .get_shard_number(&key); + + // Fast path: matched the first timeline in our local handler map. This case is common if + // only one shard per tenant is attached to this pageserver. + if first_timeline.timeline.get_shard_identity().number == shard_num { + return Ok(&first_timeline.timeline); + } + + let shard_index = ShardIndex { + shard_number: shard_num, + shard_count: first_timeline.timeline.get_shard_identity().count, + }; + + // Fast-ish path: timeline is in the connection handler's local cache + if let Some(found) = self.shard_timelines.get(&shard_index) { + return Ok(&found.timeline); + } + + key + } else { + rel_block_to_key(req.rel, req.blkno) + }; + + Err(key) + } + + /// Having looked up the [`Timeline`] instance for a particular shard, cache it to enable + /// use in future requests without having to traverse [`crate::tenant::mgr::TenantManager`] + /// again. + /// + /// Note that all the Timelines in this cache are for the same timeline_id: they're differ + /// in which shard they belong to. When we serve a getpage@lsn request, we choose a shard + /// based on key. + /// + /// The typical size of this cache is 1, as we generally create shards to distribute work + /// across pageservers, so don't tend to have multiple shards for the same tenant on the + /// same pageserver. + fn cache_timeline( + &mut self, + timeline: Arc, + ) -> Result<&Arc, GetActiveTimelineError> { + let gate_guard = timeline + .gate + .enter() + .map_err(|_| GetActiveTimelineError::Tenant(GetActiveTenantError::Cancelled))?; + + let shard_index = timeline.tenant_shard_id.to_index(); + let entry = self + .shard_timelines + .entry(shard_index) + .or_insert(HandlerTimeline { + timeline, + _guard: gate_guard, + }); + + Ok(&entry.timeline) + } + + /// If [`Self::get_cached_timeline_for_page`] missed, then this function is used to populate the cache with + /// a Timeline to serve requests for this key, if such a Timeline is present on this pageserver. If no such + /// Timeline is found, then we will return an error (this indicates that the client is talking to the wrong node). + async fn load_timeline_for_page( + &mut self, + tenant_id: TenantId, + timeline_id: TimelineId, + key: Key, + ) -> anyhow::Result<&Arc, GetActiveTimelineError> { + // Slow path: we must call out to the TenantManager to find the timeline for this Key + let timeline = self + .get_active_tenant_timeline(tenant_id, timeline_id, ShardSelector::Page(key)) + .await?; + + self.cache_timeline(timeline) + } + + async fn get_timeline_shard_zero( + &mut self, + tenant_id: TenantId, + timeline_id: TimelineId, + ) -> anyhow::Result<&Arc, GetActiveTimelineError> { + // This is a borrow-checker workaround: we can't return from inside of the `if let Some` because + // that would be an immutable-borrow-self return, whereas later in the function we will use a mutable + // ref to salf. So instead, we first build a bool, and then return while not borrowing self. + let have_cached = if let Some((idx, _tl)) = self.shard_timelines.iter().next() { + idx.shard_number == ShardNumber(0) + } else { + false + }; + + if have_cached { + let entry = self.shard_timelines.iter().next().unwrap(); + Ok(&entry.1.timeline) + } else { + let timeline = self + .get_active_tenant_timeline(tenant_id, timeline_id, ShardSelector::Zero) + .await?; + Ok(self.cache_timeline(timeline)?) + } + } + + async fn handle_get_page_at_lsn_request( + &mut self, + tenant_id: TenantId, + timeline_id: TimelineId, req: &PagestreamGetPageRequest, ctx: &RequestContext, ) -> Result { + let timeline = match self.get_cached_timeline_for_page(req) { + Ok(tl) => tl, + Err(key) => { + match self + .load_timeline_for_page(tenant_id, timeline_id, key) + .await + { + Ok(t) => t, + Err(GetActiveTimelineError::Tenant(GetActiveTenantError::NotFound(_))) => { + // We already know this tenant exists in general, because we resolved it at + // start of connection. Getting a NotFound here indicates that the shard containing + // the requested page is not present on this node: the client's knowledge of shard->pageserver + // mapping is out of date. + // + // Closing the connection by returning ``::Reconnect` has the side effect of rate-limiting above message, via + // client's reconnect backoff, as well as hopefully prompting the client to load its updated configuration + // and talk to a different pageserver. + return Err(PageStreamError::Reconnect( + "getpage@lsn request routed to wrong shard".into(), + )); + } + Err(e) => return Err(e.into()), + } + } + }; + let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn(); let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn, ctx) .await?; + let page = timeline .get_rel_page_at_lsn(req.rel, req.blkno, Version::Lsn(lsn), req.latest, ctx) .await?; @@ -899,60 +1094,6 @@ impl PageServerHandler { })) } - async fn handle_get_page_at_lsn_request( - &self, - timeline: &Timeline, - req: &PagestreamGetPageRequest, - ctx: &RequestContext, - ) -> Result { - let key = rel_block_to_key(req.rel, req.blkno); - if timeline.get_shard_identity().is_key_local(&key) { - self.do_handle_get_page_at_lsn_request(timeline, req, ctx) - .await - } else { - // The Tenant shard we looked up at connection start does not hold this particular - // key: look for other shards in this tenant. This scenario occurs if a pageserver - // has multiple shards for the same tenant. - // - // TODO: optimize this (https://github.com/neondatabase/neon/pull/6037) - let timeline = match self - .get_active_tenant_timeline( - timeline.tenant_shard_id.tenant_id, - timeline.timeline_id, - ShardSelector::Page(key), - ) - .await - { - Ok(t) => t, - Err(GetActiveTimelineError::Tenant(GetActiveTenantError::NotFound(_))) => { - // We already know this tenant exists in general, because we resolved it at - // start of connection. Getting a NotFound here indicates that the shard containing - // the requested page is not present on this node: the client's knowledge of shard->pageserver - // mapping is out of date. - tracing::info!("Page request routed to wrong shard: my identity {:?}, should go to shard {}, key {}", - timeline.get_shard_identity(), timeline.get_shard_identity().get_shard_number(&key).0, key); - // Closing the connection by returning ``::Reconnect` has the side effect of rate-limiting above message, via - // client's reconnect backoff, as well as hopefully prompting the client to load its updated configuration - // and talk to a different pageserver. - return Err(PageStreamError::Reconnect( - "getpage@lsn request routed to wrong shard".into(), - )); - } - Err(e) => return Err(e.into()), - }; - - // Take a GateGuard for the duration of this request. If we were using our main Timeline object, - // the GateGuard was already held over the whole connection. - let _timeline_guard = timeline - .gate - .enter() - .map_err(|_| PageStreamError::Shutdown)?; - - self.do_handle_get_page_at_lsn_request(&timeline, req, ctx) - .await - } - } - #[allow(clippy::too_many_arguments)] #[instrument(skip_all, fields(?lsn, ?prev_lsn, %full_backup))] async fn handle_basebackup_request( diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index c1cd823957..9d17efa48c 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -18,6 +18,7 @@ use enumset::EnumSet; use futures::stream::FuturesUnordered; use futures::FutureExt; use futures::StreamExt; +use pageserver_api::models; use pageserver_api::models::TimelineState; use pageserver_api::shard::ShardIdentity; use pageserver_api::shard::TenantShardId; @@ -73,6 +74,7 @@ use crate::tenant::config::LocationMode; use crate::tenant::config::TenantConfOpt; use crate::tenant::metadata::load_metadata; pub use crate::tenant::remote_timeline_client::index::IndexPart; +use crate::tenant::remote_timeline_client::remote_initdb_archive_path; use crate::tenant::remote_timeline_client::MaybeDeletedIndexPart; use crate::tenant::remote_timeline_client::INITDB_PATH; use crate::tenant::storage_layer::DeltaLayer; @@ -2324,6 +2326,32 @@ impl Tenant { .clone() } + /// For API access: generate a LocationConfig equivalent to the one that would be used to + /// create a Tenant in the same state. Do not use this in hot paths: it's for relatively + /// rare external API calls, like a reconciliation at startup. + pub(crate) fn get_location_conf(&self) -> models::LocationConfig { + let conf = self.tenant_conf.read().unwrap(); + + let location_config_mode = match conf.location.attach_mode { + AttachmentMode::Single => models::LocationConfigMode::AttachedSingle, + AttachmentMode::Multi => models::LocationConfigMode::AttachedMulti, + AttachmentMode::Stale => models::LocationConfigMode::AttachedStale, + }; + + // We have a pageserver TenantConf, we need the API-facing TenantConfig. + let tenant_config: models::TenantConfig = conf.tenant_conf.into(); + + models::LocationConfig { + mode: location_config_mode, + generation: self.generation.into(), + secondary_conf: None, + shard_number: self.shard_identity.number.0, + shard_count: self.shard_identity.count.0, + shard_stripe_size: self.shard_identity.stripe_size.0, + tenant_conf: tenant_config, + } + } + pub(crate) fn get_tenant_shard_id(&self) -> &TenantShardId { &self.tenant_shard_id } @@ -2674,10 +2702,11 @@ impl Tenant { } } - // Legacy configs are implicitly in attached state + // Legacy configs are implicitly in attached state, and do not support sharding Ok(LocationConf::attached_single( tenant_conf, Generation::none(), + &models::ShardParameters::default(), )) } else { // FIXME If the config file is not found, assume that we're attaching @@ -3270,6 +3299,18 @@ impl Tenant { let Some(storage) = &self.remote_storage else { bail!("no storage configured but load_existing_initdb set to {existing_initdb_timeline_id}"); }; + if existing_initdb_timeline_id != timeline_id { + let source_path = &remote_initdb_archive_path( + &self.tenant_shard_id.tenant_id, + &existing_initdb_timeline_id, + ); + let dest_path = + &remote_initdb_archive_path(&self.tenant_shard_id.tenant_id, &timeline_id); + storage + .copy_object(source_path, dest_path) + .await + .context("copy initdb tar")?; + } let (initdb_tar_zst_path, initdb_tar_zst) = self::remote_timeline_client::download_initdb_tar_zst( self.conf, @@ -3293,7 +3334,7 @@ impl Tenant { .await .context("extract initdb tar")?; } else { - // Init temporarily repo to get bootstrap data, this creates a directory in the `initdb_path` path + // Init temporarily repo to get bootstrap data, this creates a directory in the `pgdata_path` path run_initdb(self.conf, &pgdata_path, pg_version, &self.cancel).await?; // Upload the created data dir to S3 @@ -3778,6 +3819,7 @@ pub(crate) mod harness { use bytes::{Bytes, BytesMut}; use camino::Utf8PathBuf; use once_cell::sync::OnceCell; + use pageserver_api::models::ShardParameters; use pageserver_api::shard::ShardIndex; use std::fs; use std::sync::Arc; @@ -3962,6 +4004,7 @@ pub(crate) mod harness { AttachedTenantConf::try_from(LocationConf::attached_single( TenantConfOpt::from(self.tenant_conf), self.generation, + &ShardParameters::default(), )) .unwrap(), // This is a legacy/test code path: sharding isn't supported here. diff --git a/pageserver/src/tenant/config.rs b/pageserver/src/tenant/config.rs index 2d4cd350d7..c44164c12d 100644 --- a/pageserver/src/tenant/config.rs +++ b/pageserver/src/tenant/config.rs @@ -10,6 +10,7 @@ //! use anyhow::bail; use pageserver_api::models; +use pageserver_api::models::EvictionPolicy; use pageserver_api::shard::{ShardCount, ShardIdentity, ShardNumber, ShardStripeSize}; use serde::de::IntoDeserializer; use serde::{Deserialize, Serialize}; @@ -167,14 +168,17 @@ impl LocationConf { /// For use when loading from a legacy configuration: presence of a tenant /// implies it is in AttachmentMode::Single, which used to be the only /// possible state. This function should eventually be removed. - pub(crate) fn attached_single(tenant_conf: TenantConfOpt, generation: Generation) -> Self { + pub(crate) fn attached_single( + tenant_conf: TenantConfOpt, + generation: Generation, + shard_params: &models::ShardParameters, + ) -> Self { Self { mode: LocationMode::Attached(AttachedLocationConfig { generation, attach_mode: AttachmentMode::Single, }), - // Legacy configuration loads are always from tenants created before sharding existed. - shard: ShardIdentity::unsharded(), + shard: ShardIdentity::from_params(ShardNumber(0), shard_params), tenant_conf, } } @@ -428,30 +432,6 @@ pub struct TenantConfOpt { pub heatmap_period: Option, } -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(tag = "kind")] -pub enum EvictionPolicy { - NoEviction, - LayerAccessThreshold(EvictionPolicyLayerAccessThreshold), -} - -impl EvictionPolicy { - pub fn discriminant_str(&self) -> &'static str { - match self { - EvictionPolicy::NoEviction => "NoEviction", - EvictionPolicy::LayerAccessThreshold(_) => "LayerAccessThreshold", - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub struct EvictionPolicyLayerAccessThreshold { - #[serde(with = "humantime_serde")] - pub period: Duration, - #[serde(with = "humantime_serde")] - pub threshold: Duration, -} - impl TenantConfOpt { pub fn merge(&self, global_conf: TenantConf) -> TenantConf { TenantConf { @@ -576,6 +556,38 @@ impl TryFrom for TenantConfOpt { } } +/// This is a conversion from our internal tenant config object to the one used +/// in external APIs. +impl From for models::TenantConfig { + fn from(value: TenantConfOpt) -> Self { + fn humantime(d: Duration) -> String { + format!("{}s", d.as_secs()) + } + Self { + checkpoint_distance: value.checkpoint_distance, + checkpoint_timeout: value.checkpoint_timeout.map(humantime), + compaction_target_size: value.compaction_target_size, + compaction_period: value.compaction_period.map(humantime), + compaction_threshold: value.compaction_threshold, + gc_horizon: value.gc_horizon, + gc_period: value.gc_period.map(humantime), + image_creation_threshold: value.image_creation_threshold, + pitr_interval: value.pitr_interval.map(humantime), + walreceiver_connect_timeout: value.walreceiver_connect_timeout.map(humantime), + lagging_wal_timeout: value.lagging_wal_timeout.map(humantime), + max_lsn_wal_lag: value.max_lsn_wal_lag, + trace_read_requests: value.trace_read_requests, + eviction_policy: value.eviction_policy, + min_resident_size_override: value.min_resident_size_override, + evictions_low_residence_duration_metric_threshold: value + .evictions_low_residence_duration_metric_threshold + .map(humantime), + gc_feedback: value.gc_feedback, + heatmap_period: value.heatmap_period.map(humantime), + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/pageserver/src/tenant/mgr.rs b/pageserver/src/tenant/mgr.rs index 39c886cab6..50b895aca1 100644 --- a/pageserver/src/tenant/mgr.rs +++ b/pageserver/src/tenant/mgr.rs @@ -3,7 +3,8 @@ use camino::{Utf8DirEntry, Utf8Path, Utf8PathBuf}; use pageserver_api::key::Key; -use pageserver_api::shard::{ShardIdentity, ShardNumber, TenantShardId}; +use pageserver_api::models::ShardParameters; +use pageserver_api::shard::{ShardCount, ShardIdentity, ShardNumber, TenantShardId}; use rand::{distributions::Alphanumeric, Rng}; use std::borrow::Cow; use std::collections::{BTreeMap, HashMap}; @@ -56,6 +57,7 @@ use super::TenantSharedResources; /// that way we avoid having to carefully switch a tenant's ingestion etc on and off during /// its lifetime, and we can preserve some important safety invariants like `Tenant` always /// having a properly acquired generation (Secondary doesn't need a generation) +#[derive(Clone)] pub(crate) enum TenantSlot { Attached(Arc), Secondary(Arc), @@ -476,6 +478,8 @@ pub async fn init_tenant_mgr( tenant_shard_id, TenantSlot::Secondary(SecondaryTenant::new( tenant_shard_id, + location_conf.shard, + location_conf.tenant_conf, secondary_config, )), ); @@ -760,6 +764,8 @@ pub(crate) enum SetNewTenantConfigError { GetTenant(#[from] GetTenantError), #[error(transparent)] Persist(anyhow::Error), + #[error(transparent)] + Other(anyhow::Error), } pub(crate) async fn set_new_tenant_config( @@ -773,10 +779,21 @@ pub(crate) async fn set_new_tenant_config( info!("configuring tenant {tenant_id}"); let tenant = get_tenant(tenant_shard_id, true)?; + if tenant.tenant_shard_id().shard_count > ShardCount(0) { + // Note that we use ShardParameters::default below. + return Err(SetNewTenantConfigError::Other(anyhow::anyhow!( + "This API may only be used on single-sharded tenants, use the /location_config API for sharded tenants" + ))); + } + // This is a legacy API that only operates on attached tenants: the preferred // API to use is the location_config/ endpoint, which lets the caller provide // the full LocationConf. - let location_conf = LocationConf::attached_single(new_tenant_conf, tenant.generation); + let location_conf = LocationConf::attached_single( + new_tenant_conf, + tenant.generation, + &ShardParameters::default(), + ); Tenant::persist_tenant_config(conf, &tenant_shard_id, &location_conf) .await @@ -907,6 +924,7 @@ impl TenantManager { Some(TenantSlot::Secondary(secondary_tenant)), ) => { secondary_tenant.set_config(secondary_conf); + secondary_tenant.set_tenant_conf(&new_location_config.tenant_conf); Some(FastPathModified::Secondary(secondary_tenant.clone())) } _ => { @@ -1039,16 +1057,36 @@ impl TenantManager { let new_slot = match &new_location_config.mode { LocationMode::Secondary(secondary_config) => { - TenantSlot::Secondary(SecondaryTenant::new(tenant_shard_id, secondary_config)) + let shard_identity = new_location_config.shard; + TenantSlot::Secondary(SecondaryTenant::new( + tenant_shard_id, + shard_identity, + new_location_config.tenant_conf, + secondary_config, + )) } LocationMode::Attached(_attach_config) => { let shard_identity = new_location_config.shard; + + // Testing hack: if we are configured with no control plane, then drop the generation + // from upserts. This enables creating generation-less tenants even though neon_local + // always uses generations when calling the location conf API. + let attached_conf = if cfg!(feature = "testing") { + let mut conf = AttachedTenantConf::try_from(new_location_config)?; + if self.conf.control_plane_api.is_none() { + conf.location.generation = Generation::none(); + } + conf + } else { + AttachedTenantConf::try_from(new_location_config)? + }; + let tenant = tenant_spawn( self.conf, tenant_shard_id, &tenant_path, self.resources.clone(), - AttachedTenantConf::try_from(new_location_config)?, + attached_conf, shard_identity, None, self.tenants, @@ -1189,6 +1227,17 @@ impl TenantManager { } } + /// Total list of all tenant slots: this includes attached, secondary, and InProgress. + pub(crate) fn list(&self) -> Vec<(TenantShardId, TenantSlot)> { + let locked = self.tenants.read().unwrap(); + match &*locked { + TenantsMap::Initializing => Vec::new(), + TenantsMap::Open(map) | TenantsMap::ShuttingDown(map) => { + map.iter().map(|(k, v)| (*k, v.clone())).collect() + } + } + } + pub(crate) async fn delete_tenant( &self, tenant_shard_id: TenantShardId, diff --git a/pageserver/src/tenant/remote_timeline_client.rs b/pageserver/src/tenant/remote_timeline_client.rs index ec2a6efef6..7935209252 100644 --- a/pageserver/src/tenant/remote_timeline_client.rs +++ b/pageserver/src/tenant/remote_timeline_client.rs @@ -182,7 +182,7 @@ pub(crate) mod download; pub mod index; -mod upload; +pub(crate) mod upload; use anyhow::Context; use camino::Utf8Path; @@ -691,7 +691,10 @@ impl RemoteTimelineClient { .insert(layer.layer_desc().filename(), metadata.clone()); upload_queue.latest_files_changes_since_metadata_upload_scheduled += 1; - info!("scheduled layer file upload {layer}"); + info!( + "scheduled layer file upload {layer} gen={:?} shard={:?}", + metadata.generation, metadata.shard + ); let op = UploadOp::UploadLayer(layer, metadata); self.calls_unfinished_metric_begin(&op); upload_queue.queued_operations.push_back(op); diff --git a/pageserver/src/tenant/secondary.rs b/pageserver/src/tenant/secondary.rs index 2331447266..d00d901be6 100644 --- a/pageserver/src/tenant/secondary.rs +++ b/pageserver/src/tenant/secondary.rs @@ -3,22 +3,36 @@ pub mod heatmap; mod heatmap_uploader; mod scheduler; -use std::sync::Arc; +use std::{sync::Arc, time::SystemTime}; -use crate::task_mgr::{self, TaskKind, BACKGROUND_RUNTIME}; +use crate::{ + config::PageServerConf, + disk_usage_eviction_task::DiskUsageEvictionInfo, + task_mgr::{self, TaskKind, BACKGROUND_RUNTIME}, + virtual_file::MaybeFatalIo, +}; use self::{ downloader::{downloader_task, SecondaryDetail}, heatmap_uploader::heatmap_uploader_task, }; -use super::{config::SecondaryLocationConfig, mgr::TenantManager}; +use super::{ + config::{SecondaryLocationConfig, TenantConfOpt}, + mgr::TenantManager, + span::debug_assert_current_span_has_tenant_id, + storage_layer::LayerFileName, +}; -use pageserver_api::shard::TenantShardId; +use pageserver_api::{ + models, + shard::{ShardIdentity, TenantShardId}, +}; use remote_storage::GenericRemoteStorage; use tokio_util::sync::CancellationToken; -use utils::{completion::Barrier, sync::gate::Gate}; +use tracing::instrument; +use utils::{completion::Barrier, fs_ext, id::TimelineId, sync::gate::Gate}; enum DownloadCommand { Download(TenantShardId), @@ -75,12 +89,20 @@ pub(crate) struct SecondaryTenant { pub(crate) gate: Gate, + // Secondary mode does not need the full shard identity or the TenantConfOpt. However, + // storing these enables us to report our full LocationConf, enabling convenient reconciliation + // by the control plane (see [`Self::get_location_conf`]) + shard_identity: ShardIdentity, + tenant_conf: std::sync::Mutex, + detail: std::sync::Mutex, } impl SecondaryTenant { pub(crate) fn new( tenant_shard_id: TenantShardId, + shard_identity: ShardIdentity, + tenant_conf: TenantConfOpt, config: &SecondaryLocationConfig, ) -> Arc { Arc::new(Self { @@ -92,6 +114,9 @@ impl SecondaryTenant { cancel: CancellationToken::new(), gate: Gate::new(format!("SecondaryTenant {tenant_shard_id}")), + shard_identity, + tenant_conf: std::sync::Mutex::new(tenant_conf), + detail: std::sync::Mutex::new(SecondaryDetail::new(config.clone())), }) } @@ -107,9 +132,91 @@ impl SecondaryTenant { self.detail.lock().unwrap().config = config.clone(); } - fn get_tenant_shard_id(&self) -> &TenantShardId { + pub(crate) fn set_tenant_conf(&self, config: &TenantConfOpt) { + *(self.tenant_conf.lock().unwrap()) = *config; + } + + /// For API access: generate a LocationConfig equivalent to the one that would be used to + /// create a Tenant in the same state. Do not use this in hot paths: it's for relatively + /// rare external API calls, like a reconciliation at startup. + pub(crate) fn get_location_conf(&self) -> models::LocationConfig { + let conf = self.detail.lock().unwrap().config.clone(); + + let conf = models::LocationConfigSecondary { warm: conf.warm }; + + let tenant_conf = *self.tenant_conf.lock().unwrap(); + models::LocationConfig { + mode: models::LocationConfigMode::Secondary, + generation: None, + secondary_conf: Some(conf), + shard_number: self.tenant_shard_id.shard_number.0, + shard_count: self.tenant_shard_id.shard_count.0, + shard_stripe_size: self.shard_identity.stripe_size.0, + tenant_conf: tenant_conf.into(), + } + } + + pub(crate) fn get_tenant_shard_id(&self) -> &TenantShardId { &self.tenant_shard_id } + + pub(crate) fn get_layers_for_eviction(self: &Arc) -> DiskUsageEvictionInfo { + self.detail.lock().unwrap().get_layers_for_eviction(self) + } + + #[instrument(skip_all, fields(tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug(), timeline_id=%timeline_id, name=%name))] + pub(crate) async fn evict_layer( + &self, + conf: &PageServerConf, + timeline_id: TimelineId, + name: LayerFileName, + ) { + debug_assert_current_span_has_tenant_id(); + + let _guard = match self.gate.enter() { + Ok(g) => g, + Err(_) => { + tracing::debug!("Dropping layer evictions, secondary tenant shutting down",); + return; + } + }; + + let now = SystemTime::now(); + + let path = conf + .timeline_path(&self.tenant_shard_id, &timeline_id) + .join(name.file_name()); + + // We tolerate ENOENT, because between planning eviction and executing + // it, the secondary downloader could have seen an updated heatmap that + // resulted in a layer being deleted. + // Other local I/O errors are process-fatal: these should never happen. + tokio::fs::remove_file(path) + .await + .or_else(fs_ext::ignore_not_found) + .fatal_err("Deleting layer during eviction"); + + // Update the timeline's state. This does not have to be synchronized with + // the download process, because: + // - If downloader is racing with us to remove a file (e.g. because it is + // removed from heatmap), then our mutual .remove() operations will both + // succeed. + // - If downloader is racing with us to download the object (this would require + // multiple eviction iterations to race with multiple download iterations), then + // if we remove it from the state, the worst that happens is the downloader + // downloads it again before re-inserting, or we delete the file but it remains + // in the state map (in which case it will be downloaded if this secondary + // tenant transitions to attached and tries to access it) + // + // The important assumption here is that the secondary timeline state does not + // have to 100% match what is on disk, because it's a best-effort warming + // of the cache. + let mut detail = self.detail.lock().unwrap(); + if let Some(timeline_detail) = detail.timelines.get_mut(&timeline_id) { + timeline_detail.on_disk_layers.remove(&name); + timeline_detail.evicted_at.insert(name, now); + } + } } /// The SecondaryController is a pseudo-rpc client for administrative control of secondary mode downloads, diff --git a/pageserver/src/tenant/secondary/downloader.rs b/pageserver/src/tenant/secondary/downloader.rs index 2a79c406cf..702c0b1ec1 100644 --- a/pageserver/src/tenant/secondary/downloader.rs +++ b/pageserver/src/tenant/secondary/downloader.rs @@ -8,6 +8,9 @@ use std::{ use crate::{ config::PageServerConf, + disk_usage_eviction_task::{ + finite_f32, DiskUsageEvictionInfo, EvictionCandidate, EvictionLayer, EvictionSecondaryLayer, + }, metrics::SECONDARY_MODE, tenant::{ config::SecondaryLocationConfig, @@ -142,6 +145,46 @@ impl SecondaryDetail { timelines: HashMap::new(), } } + + pub(super) fn get_layers_for_eviction( + &self, + parent: &Arc, + ) -> DiskUsageEvictionInfo { + let mut result = DiskUsageEvictionInfo { + max_layer_size: None, + resident_layers: Vec::new(), + }; + for (timeline_id, timeline_detail) in &self.timelines { + result + .resident_layers + .extend(timeline_detail.on_disk_layers.iter().map(|(name, ods)| { + EvictionCandidate { + layer: EvictionLayer::Secondary(EvictionSecondaryLayer { + secondary_tenant: parent.clone(), + timeline_id: *timeline_id, + name: name.clone(), + metadata: ods.metadata.clone(), + }), + last_activity_ts: ods.access_time, + relative_last_activity: finite_f32::FiniteF32::ZERO, + } + })); + } + result.max_layer_size = result + .resident_layers + .iter() + .map(|l| l.layer.get_file_size()) + .max(); + + tracing::debug!( + "eviction: secondary tenant {} found {} timelines, {} layers", + parent.get_tenant_shard_id(), + self.timelines.len(), + result.resident_layers.len() + ); + + result + } } struct PendingDownload { diff --git a/pageserver/src/tenant/storage_layer/layer.rs b/pageserver/src/tenant/storage_layer/layer.rs index 3f29e9f6a5..12af866810 100644 --- a/pageserver/src/tenant/storage_layer/layer.rs +++ b/pageserver/src/tenant/storage_layer/layer.rs @@ -15,7 +15,7 @@ use utils::sync::heavier_once_cell; use crate::config::PageServerConf; use crate::context::RequestContext; use crate::repository::Key; -use crate::tenant::{remote_timeline_client::LayerFileMetadata, RemoteTimelineClient, Timeline}; +use crate::tenant::{remote_timeline_client::LayerFileMetadata, Timeline}; use super::delta_layer::{self, DeltaEntry}; use super::image_layer; @@ -204,17 +204,14 @@ impl Layer { /// /// Technically cancellation safe, but cancelling might shift the viewpoint of what generation /// of download-evict cycle on retry. - pub(crate) async fn evict_and_wait( - &self, - rtc: &RemoteTimelineClient, - ) -> Result<(), EvictionError> { - self.0.evict_and_wait(rtc).await + pub(crate) async fn evict_and_wait(&self) -> Result<(), EvictionError> { + self.0.evict_and_wait().await } /// Delete the layer file when the `self` gets dropped, also try to schedule a remote index upload /// then. /// - /// On drop, this will cause a call to [`RemoteTimelineClient::schedule_deletion_of_unlinked`]. + /// On drop, this will cause a call to [`crate::tenant::remote_timeline_client::RemoteTimelineClient::schedule_deletion_of_unlinked`]. /// This means that the unlinking by [gc] or [compaction] must have happened strictly before /// the value this is called on gets dropped. /// @@ -606,10 +603,7 @@ impl LayerInner { /// Cancellation safe, however dropping the future and calling this method again might result /// in a new attempt to evict OR join the previously started attempt. - pub(crate) async fn evict_and_wait( - &self, - _: &RemoteTimelineClient, - ) -> Result<(), EvictionError> { + pub(crate) async fn evict_and_wait(&self) -> Result<(), EvictionError> { use tokio::sync::broadcast::error::RecvError; assert!(self.have_remote_client); diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index 8b96c30522..6b1487259f 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -15,9 +15,10 @@ use fail::fail_point; use itertools::Itertools; use pageserver_api::{ models::{ - DownloadRemoteLayersTaskInfo, DownloadRemoteLayersTaskSpawnRequest, LayerMapInfo, - TimelineState, + DownloadRemoteLayersTaskInfo, DownloadRemoteLayersTaskSpawnRequest, EvictionPolicy, + LayerMapInfo, TimelineState, }, + reltag::BlockNumber, shard::{ShardIdentity, TenantShardId}, }; use rand::Rng; @@ -42,33 +43,39 @@ use std::{ ops::ControlFlow, }; -use crate::context::{ - AccessStatsBehavior, DownloadBehavior, RequestContext, RequestContextBuilder, -}; -use crate::tenant::storage_layer::delta_layer::DeltaEntry; -use crate::tenant::storage_layer::{ - AsLayerDesc, DeltaLayerWriter, EvictionError, ImageLayerWriter, InMemoryLayer, Layer, - LayerAccessStatsReset, LayerFileName, ResidentLayer, ValueReconstructResult, - ValueReconstructState, -}; -use crate::tenant::tasks::BackgroundLoopKind; use crate::tenant::timeline::logical_size::CurrentLogicalSize; use crate::tenant::{ layer_map::{LayerMap, SearchResult}, metadata::{save_metadata, TimelineMetadata}, par_fsync, }; +use crate::{ + context::{AccessStatsBehavior, DownloadBehavior, RequestContext, RequestContextBuilder}, + disk_usage_eviction_task::DiskUsageEvictionInfo, +}; use crate::{deletion_queue::DeletionQueueClient, tenant::remote_timeline_client::StopError}; +use crate::{ + disk_usage_eviction_task::finite_f32, + tenant::storage_layer::{ + AsLayerDesc, DeltaLayerWriter, EvictionError, ImageLayerWriter, InMemoryLayer, Layer, + LayerAccessStatsReset, LayerFileName, ResidentLayer, ValueReconstructResult, + ValueReconstructState, + }, +}; +use crate::{ + disk_usage_eviction_task::EvictionCandidate, tenant::storage_layer::delta_layer::DeltaEntry, +}; +use crate::{pgdatadir_mapping::LsnForTimestamp, tenant::tasks::BackgroundLoopKind}; use crate::config::PageServerConf; use crate::keyspace::{KeyPartitioning, KeySpace, KeySpaceRandomAccum}; use crate::metrics::{ TimelineMetrics, MATERIALIZED_PAGE_CACHE_HIT, MATERIALIZED_PAGE_CACHE_HIT_DIRECT, }; +use crate::pgdatadir_mapping::CalculateLogicalSizeError; use crate::pgdatadir_mapping::{is_inherited_key, is_rel_fsm_block_key, is_rel_vm_block_key}; -use crate::pgdatadir_mapping::{CalculateLogicalSizeError, LsnForTimestamp}; -use crate::tenant::config::{EvictionPolicy, TenantConfOpt}; -use pageserver_api::reltag::{BlockNumber, RelTag}; +use crate::tenant::config::TenantConfOpt; +use pageserver_api::reltag::RelTag; use pageserver_api::shard::ShardIndex; use postgres_connection::PgConnectionConfig; @@ -1133,12 +1140,7 @@ impl Timeline { return Ok(None); }; - let rtc = self - .remote_client - .as_ref() - .ok_or_else(|| anyhow::anyhow!("remote storage not configured; cannot evict"))?; - - match local_layer.evict_and_wait(rtc).await { + match local_layer.evict_and_wait().await { Ok(()) => Ok(Some(true)), Err(EvictionError::NotFound) => Ok(Some(false)), Err(EvictionError::Downloaded) => Ok(Some(false)), @@ -2102,7 +2104,7 @@ impl Timeline { let layer_file_names = eviction_info .resident_layers .iter() - .map(|l| l.layer.layer_desc().filename()) + .map(|l| l.layer.get_name()) .collect::>(); let decorated = match remote_client.get_layers_metadata(layer_file_names) { @@ -2120,7 +2122,7 @@ impl Timeline { .filter_map(|(layer, remote_info)| { remote_info.map(|remote_info| { HeatMapLayer::new( - layer.layer.layer_desc().filename(), + layer.layer.get_name(), IndexLayerMetadata::from(remote_info), layer.last_activity_ts, ) @@ -4423,43 +4425,6 @@ impl Timeline { } } -pub(crate) struct DiskUsageEvictionInfo { - /// Timeline's largest layer (remote or resident) - pub max_layer_size: Option, - /// Timeline's resident layers - pub resident_layers: Vec, -} - -pub(crate) struct LocalLayerInfoForDiskUsageEviction { - pub layer: Layer, - pub last_activity_ts: SystemTime, -} - -impl std::fmt::Debug for LocalLayerInfoForDiskUsageEviction { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - // format the tv_sec, tv_nsec into rfc3339 in case someone is looking at it - // having to allocate a string to this is bad, but it will rarely be formatted - let ts = chrono::DateTime::::from(self.last_activity_ts); - let ts = ts.to_rfc3339_opts(chrono::SecondsFormat::Nanos, true); - struct DisplayIsDebug<'a, T>(&'a T); - impl<'a, T: std::fmt::Display> std::fmt::Debug for DisplayIsDebug<'a, T> { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.0) - } - } - f.debug_struct("LocalLayerInfoForDiskUsageEviction") - .field("layer", &DisplayIsDebug(&self.layer)) - .field("last_activity", &ts) - .finish() - } -} - -impl LocalLayerInfoForDiskUsageEviction { - pub fn file_size(&self) -> u64 { - self.layer.layer_desc().file_size - } -} - impl Timeline { /// Returns non-remote layers for eviction. pub(crate) async fn get_local_layers_for_disk_usage_eviction(&self) -> DiskUsageEvictionInfo { @@ -4493,9 +4458,10 @@ impl Timeline { SystemTime::now() }); - resident_layers.push(LocalLayerInfoForDiskUsageEviction { - layer: l.drop_eviction_guard(), + resident_layers.push(EvictionCandidate { + layer: l.drop_eviction_guard().into(), last_activity_ts, + relative_last_activity: finite_f32::FiniteF32::ZERO, }); } @@ -4652,11 +4618,6 @@ mod tests { .await .unwrap(); - let rtc = timeline - .remote_client - .clone() - .expect("just configured this"); - let layer = find_some_layer(&timeline).await; let layer = layer .keep_resident() @@ -4665,8 +4626,8 @@ mod tests { .expect("should had been resident") .drop_eviction_guard(); - let first = async { layer.evict_and_wait(&rtc).await }; - let second = async { layer.evict_and_wait(&rtc).await }; + let first = async { layer.evict_and_wait().await }; + let second = async { layer.evict_and_wait().await }; let (first, second) = tokio::join!(first, second); diff --git a/pageserver/src/tenant/timeline/eviction_task.rs b/pageserver/src/tenant/timeline/eviction_task.rs index ea5f5f5fa7..01a5bfc32b 100644 --- a/pageserver/src/tenant/timeline/eviction_task.rs +++ b/pageserver/src/tenant/timeline/eviction_task.rs @@ -20,6 +20,7 @@ use std::{ time::{Duration, SystemTime}, }; +use pageserver_api::models::{EvictionPolicy, EvictionPolicyLayerAccessThreshold}; use tokio::time::Instant; use tokio_util::sync::CancellationToken; use tracing::{debug, error, info, info_span, instrument, warn, Instrument}; @@ -29,10 +30,7 @@ use crate::{ pgdatadir_mapping::CollectKeySpaceError, task_mgr::{self, TaskKind, BACKGROUND_RUNTIME}, tenant::{ - config::{EvictionPolicy, EvictionPolicyLayerAccessThreshold}, - tasks::BackgroundLoopKind, - timeline::EvictionError, - LogicalSizeCalculationCause, Tenant, + tasks::BackgroundLoopKind, timeline::EvictionError, LogicalSizeCalculationCause, Tenant, }, }; @@ -215,13 +213,10 @@ impl Timeline { // So, we just need to deal with this. - let remote_client = match self.remote_client.as_ref() { - Some(c) => c, - None => { - error!("no remote storage configured, cannot evict layers"); - return ControlFlow::Continue(()); - } - }; + if self.remote_client.is_none() { + error!("no remote storage configured, cannot evict layers"); + return ControlFlow::Continue(()); + } let mut js = tokio::task::JoinSet::new(); { @@ -274,9 +269,8 @@ impl Timeline { }; let layer = guard.drop_eviction_guard(); if no_activity_for > p.threshold { - let remote_client = remote_client.clone(); // this could cause a lot of allocations in some cases - js.spawn(async move { layer.evict_and_wait(&remote_client).await }); + js.spawn(async move { layer.evict_and_wait().await }); stats.candidates += 1; } } diff --git a/pgxn/neon/relsize_cache.c b/pgxn/neon/relsize_cache.c index b13134b5c3..cc7ac2c394 100644 --- a/pgxn/neon/relsize_cache.c +++ b/pgxn/neon/relsize_cache.c @@ -40,11 +40,23 @@ typedef struct { RelTag tag; BlockNumber size; + dlist_node lru_node; /* LRU list node */ } RelSizeEntry; +typedef struct +{ + size_t size; + uint64 hits; + uint64 misses; + uint64 writes; + dlist_head lru; /* double linked list for LRU replacement + * algorithm */ +} RelSizeHashControl; + static HTAB *relsize_hash; static LWLockId relsize_lock; static int relsize_hash_size; +static RelSizeHashControl* relsize_ctl; static shmem_startup_hook_type prev_shmem_startup_hook = NULL; #if PG_VERSION_NUM >= 150000 static shmem_request_hook_type prev_shmem_request_hook = NULL; @@ -52,7 +64,7 @@ static void relsize_shmem_request(void); #endif /* - * Size of a cache entry is 20 bytes. So this default will take about 1.2 MB, + * Size of a cache entry is 36 bytes. So this default will take about 2.3 MB, * which seems reasonable. */ #define DEFAULT_RELSIZE_HASH_SIZE (64 * 1024) @@ -61,19 +73,29 @@ static void neon_smgr_shmem_startup(void) { static HASHCTL info; + bool found; if (prev_shmem_startup_hook) prev_shmem_startup_hook(); LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE); - relsize_lock = (LWLockId) GetNamedLWLockTranche("neon_relsize"); - info.keysize = sizeof(RelTag); - info.entrysize = sizeof(RelSizeEntry); - relsize_hash = ShmemInitHash("neon_relsize", - relsize_hash_size, relsize_hash_size, - &info, - HASH_ELEM | HASH_BLOBS); - LWLockRelease(AddinShmemInitLock); + relsize_ctl = (RelSizeHashControl *) ShmemInitStruct("relsize_hash", sizeof(RelSizeHashControl), &found); + if (!found) + { + relsize_lock = (LWLockId) GetNamedLWLockTranche("neon_relsize"); + info.keysize = sizeof(RelTag); + info.entrysize = sizeof(RelSizeEntry); + relsize_hash = ShmemInitHash("neon_relsize", + relsize_hash_size, relsize_hash_size, + &info, + HASH_ELEM | HASH_BLOBS); + LWLockRelease(AddinShmemInitLock); + relsize_ctl->size = 0; + relsize_ctl->hits = 0; + relsize_ctl->misses = 0; + relsize_ctl->writes = 0; + dlist_init(&relsize_ctl->lru); + } } bool @@ -93,7 +115,15 @@ get_cached_relsize(NRelFileInfo rinfo, ForkNumber forknum, BlockNumber *size) if (entry != NULL) { *size = entry->size; + relsize_ctl->hits += 1; found = true; + /* Move entry to the LRU list tail */ + dlist_delete(&entry->lru_node); + dlist_push_tail(&relsize_ctl->lru, &entry->lru_node); + } + else + { + relsize_ctl->misses += 1; } LWLockRelease(relsize_lock); } @@ -107,12 +137,43 @@ set_cached_relsize(NRelFileInfo rinfo, ForkNumber forknum, BlockNumber size) { RelTag tag; RelSizeEntry *entry; + bool found = false; tag.rinfo = rinfo; tag.forknum = forknum; LWLockAcquire(relsize_lock, LW_EXCLUSIVE); - entry = hash_search(relsize_hash, &tag, HASH_ENTER, NULL); + /* + * This should actually never happen! Below we check if hash is full and delete least recently user item in this case. + * But for further safety we also perform check here. + */ + while ((entry = hash_search(relsize_hash, &tag, HASH_ENTER_NULL, &found)) == NULL) + { + RelSizeEntry *victim = dlist_container(RelSizeEntry, lru_node, dlist_pop_head_node(&relsize_ctl->lru)); + hash_search(relsize_hash, &victim->tag, HASH_REMOVE, NULL); + Assert(relsize_ctl->size > 0); + relsize_ctl->size -= 1; + } entry->size = size; + if (!found) + { + if (++relsize_ctl->size == relsize_hash_size) + { + /* + * Remove least recently used elment from the hash. + * Hash size after is becomes `relsize_hash_size-1`. + * But it is not considered to be a problem, because size of this hash is expecrted large enough and +-1 doesn't matter. + */ + RelSizeEntry *victim = dlist_container(RelSizeEntry, lru_node, dlist_pop_head_node(&relsize_ctl->lru)); + hash_search(relsize_hash, &victim->tag, HASH_REMOVE, NULL); + relsize_ctl->size -= 1; + } + } + else + { + dlist_delete(&entry->lru_node); + } + dlist_push_tail(&relsize_ctl->lru, &entry->lru_node); + relsize_ctl->writes += 1; LWLockRelease(relsize_lock); } } @@ -132,6 +193,21 @@ update_cached_relsize(NRelFileInfo rinfo, ForkNumber forknum, BlockNumber size) entry = hash_search(relsize_hash, &tag, HASH_ENTER, &found); if (!found || entry->size < size) entry->size = size; + if (!found) + { + if (++relsize_ctl->size == relsize_hash_size) + { + RelSizeEntry *victim = dlist_container(RelSizeEntry, lru_node, dlist_pop_head_node(&relsize_ctl->lru)); + hash_search(relsize_hash, &victim->tag, HASH_REMOVE, NULL); + relsize_ctl->size -= 1; + } + } + else + { + dlist_delete(&entry->lru_node); + } + relsize_ctl->writes += 1; + dlist_push_tail(&relsize_ctl->lru, &entry->lru_node); LWLockRelease(relsize_lock); } } @@ -142,11 +218,16 @@ forget_cached_relsize(NRelFileInfo rinfo, ForkNumber forknum) if (relsize_hash_size > 0) { RelTag tag; - + RelSizeEntry *entry; tag.rinfo = rinfo; tag.forknum = forknum; LWLockAcquire(relsize_lock, LW_EXCLUSIVE); - hash_search(relsize_hash, &tag, HASH_REMOVE, NULL); + entry = hash_search(relsize_hash, &tag, HASH_REMOVE, NULL); + if (entry) + { + dlist_delete(&entry->lru_node); + relsize_ctl->size -= 1; + } LWLockRelease(relsize_lock); } } @@ -191,7 +272,7 @@ relsize_shmem_request(void) if (prev_shmem_request_hook) prev_shmem_request_hook(); - RequestAddinShmemSpace(hash_estimate_size(relsize_hash_size, sizeof(RelSizeEntry))); + RequestAddinShmemSpace(sizeof(RelSizeHashControl) + hash_estimate_size(relsize_hash_size, sizeof(RelSizeEntry))); RequestNamedLWLockTranche("neon_relsize", 1); } #endif diff --git a/proxy/src/serverless/sql_over_http.rs b/proxy/src/serverless/sql_over_http.rs index 719559ed48..391fb95e9e 100644 --- a/proxy/src/serverless/sql_over_http.rs +++ b/proxy/src/serverless/sql_over_http.rs @@ -15,6 +15,7 @@ use serde_json::Map; use serde_json::Value; use smol_str::SmolStr; use tokio_postgres::error::DbError; +use tokio_postgres::error::ErrorPosition; use tokio_postgres::types::Kind; use tokio_postgres::types::Type; use tokio_postgres::GenericClient; @@ -59,6 +60,7 @@ enum Payload { const MAX_RESPONSE_SIZE: usize = 10 * 1024 * 1024; // 10 MiB const MAX_REQUEST_SIZE: u64 = 10 * 1024 * 1024; // 10 MiB +const SERVERLESS_DRIVER_SNI_HOSTNAME_FIRST_PART: &str = "api"; static RAW_TEXT_OUTPUT: HeaderName = HeaderName::from_static("neon-raw-text-output"); static ARRAY_MODE: HeaderName = HeaderName::from_static("neon-array-mode"); @@ -176,10 +178,11 @@ fn get_conn_info( .and_then(|h| h.to_str().ok()) .and_then(|h| h.split(':').next()); - if hostname != sni_hostname { + // sni_hostname has to be either the same as hostname or the one used in serverless driver. + if !check_matches(&sni_hostname, hostname)? { return Err(anyhow::anyhow!("mismatched SNI hostname and hostname")); } else if let Some(h) = host_header { - if h != hostname { + if h != sni_hostname { return Err(anyhow::anyhow!("mismatched host header and hostname")); } } @@ -213,6 +216,20 @@ fn get_conn_info( }) } +fn check_matches(sni_hostname: &str, hostname: &str) -> Result { + if sni_hostname == hostname { + return Ok(true); + } + let (sni_hostname_first, sni_hostname_rest) = sni_hostname + .split_once('.') + .ok_or_else(|| anyhow::anyhow!("Unexpected sni format."))?; + let (_, hostname_rest) = hostname + .split_once('.') + .ok_or_else(|| anyhow::anyhow!("Unexpected hostname format."))?; + Ok(sni_hostname_rest == hostname_rest + && sni_hostname_first == SERVERLESS_DRIVER_SNI_HOSTNAME_FIRST_PART) +} + // TODO: return different http error codes pub async fn handle( tls: &'static TlsConfig, @@ -231,7 +248,7 @@ pub async fn handle( Ok(r) => match r { Ok(r) => r, Err(e) => { - let message = format!("{:?}", e); + let mut message = format!("{:?}", e); let db_error = e .downcast_ref::() .and_then(|e| e.as_db_error()); @@ -244,7 +261,25 @@ pub async fn handle( .unwrap_or_default() } - // TODO(conrad): db_error.position() + if let Some(db_error) = db_error { + db_error.message().clone_into(&mut message); + } + + let position = db_error.and_then(|db| db.position()); + let (position, internal_position, internal_query) = match position { + Some(ErrorPosition::Original(position)) => ( + Value::String(position.to_string()), + Value::Null, + Value::Null, + ), + Some(ErrorPosition::Internal { position, query }) => ( + Value::Null, + Value::String(position.to_string()), + Value::String(query.clone()), + ), + None => (Value::Null, Value::Null, Value::Null), + }; + let code = get(db_error, |db| db.code().code()); let severity = get(db_error, |db| db.severity()); let detail = get(db_error, |db| db.detail()); @@ -256,7 +291,7 @@ pub async fn handle( let datatype = get(db_error, |db| db.datatype()); let constraint = get(db_error, |db| db.constraint()); let file = get(db_error, |db| db.file()); - let line = get(db_error, |db| db.line()); + let line = get(db_error, |db| db.line().map(|l| l.to_string())); let routine = get(db_error, |db| db.routine()); error!( @@ -271,12 +306,15 @@ pub async fn handle( "code": code, "detail": detail, "hint": hint, + "position": position, + "internalPosition": internal_position, + "internalQuery": internal_query, "severity": severity, "where": where_, "table": table, "column": column, "schema": schema, - "datatype": datatype, + "dataType": datatype, "constraint": constraint, "file": file, "line": line, diff --git a/test_runner/fixtures/metrics.py b/test_runner/fixtures/metrics.py index a6a25da332..7c489bda67 100644 --- a/test_runner/fixtures/metrics.py +++ b/test_runner/fixtures/metrics.py @@ -16,6 +16,7 @@ class Metrics: def query_all(self, name: str, filter: Optional[Dict[str, str]] = None) -> List[Sample]: filter = filter or {} res = [] + for sample in self.metrics[name]: try: if all(sample.labels[k] == v for k, v in filter.items()): diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index c139fa215c..ffd93004d2 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -19,7 +19,7 @@ from functools import cached_property from itertools import chain, product from pathlib import Path from types import TracebackType -from typing import Any, Dict, Iterator, List, Optional, Tuple, Type, cast +from typing import Any, Dict, Iterator, List, Optional, Tuple, Type, Union, cast from urllib.parse import urlparse import asyncpg @@ -61,7 +61,7 @@ from fixtures.remote_storage import ( default_remote_storage, remote_storage_to_toml_inline_table, ) -from fixtures.types import Lsn, TenantId, TimelineId +from fixtures.types import Lsn, TenantId, TenantShardId, TimelineId from fixtures.utils import ( ATTACHMENT_NAME_REGEX, allure_add_grafana_links, @@ -495,6 +495,8 @@ class NeonEnvBuilder: self, initial_tenant_conf: Optional[Dict[str, str]] = None, default_remote_storage_if_missing: bool = True, + initial_tenant_shard_count: Optional[int] = None, + initial_tenant_shard_stripe_size: Optional[int] = None, ) -> NeonEnv: """ Default way to create and start NeonEnv. Also creates the initial_tenant with root initial_timeline. @@ -512,7 +514,11 @@ class NeonEnvBuilder: f"Services started, creating initial tenant {env.initial_tenant} and its initial timeline" ) initial_tenant, initial_timeline = env.neon_cli.create_tenant( - tenant_id=env.initial_tenant, conf=initial_tenant_conf, timeline_id=env.initial_timeline + tenant_id=env.initial_tenant, + conf=initial_tenant_conf, + timeline_id=env.initial_timeline, + shard_count=initial_tenant_shard_count, + shard_stripe_size=initial_tenant_shard_stripe_size, ) assert env.initial_tenant == initial_tenant assert env.initial_timeline == initial_timeline @@ -861,7 +867,9 @@ class NeonEnv: attachment_service_port = self.port_distributor.get_port() self.control_plane_api: str = f"http://127.0.0.1:{attachment_service_port}" - self.attachment_service: NeonAttachmentService = NeonAttachmentService(self) + self.attachment_service: NeonAttachmentService = NeonAttachmentService( + self, config.auth_enabled + ) # Create a config file corresponding to the options cfg: Dict[str, Any] = { @@ -983,6 +991,16 @@ class NeonEnv: raise RuntimeError(f"Pageserver with ID {id} not found") + def get_tenant_pageserver(self, tenant_id: Union[TenantId, TenantShardId]): + """ + Get the NeonPageserver where this tenant shard is currently attached, according + to the attachment service. + """ + meta = self.attachment_service.inspect(tenant_id) + assert meta is not None, f"{tenant_id} attachment location not found" + pageserver_id = meta[1] + return self.get_pageserver(pageserver_id) + def get_safekeeper_connstrs(self) -> str: """Get list of safekeeper endpoints suitable for safekeepers GUC""" return ",".join(f"localhost:{wa.port.pg}" for wa in self.safekeepers) @@ -1226,15 +1244,29 @@ class AbstractNeonCli(abc.ABC): env_vars[var] = val # Intercept CalledProcessError and print more info - res = subprocess.run( - args, - env=env_vars, - check=False, - universal_newlines=True, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - timeout=timeout, - ) + try: + res = subprocess.run( + args, + env=env_vars, + check=False, + universal_newlines=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + timeout=timeout, + ) + except subprocess.TimeoutExpired as e: + if e.stderr: + stderr = e.stderr.decode(errors="replace") + else: + stderr = "" + + if e.stdout: + stdout = e.stdout.decode(errors="replace") + else: + stdout = "" + + log.warn(f"CLI timeout: stderr={stderr}, stdout={stdout}") + raise indent = " " if not res.returncode: @@ -1285,6 +1317,8 @@ class NeonCli(AbstractNeonCli): tenant_id: Optional[TenantId] = None, timeline_id: Optional[TimelineId] = None, conf: Optional[Dict[str, str]] = None, + shard_count: Optional[int] = None, + shard_stripe_size: Optional[int] = None, set_default: bool = False, ) -> Tuple[TenantId, TimelineId]: """ @@ -1312,6 +1346,12 @@ class NeonCli(AbstractNeonCli): if set_default: args.append("--set-default") + if shard_count is not None: + args.extend(["--shard-count", str(shard_count)]) + + if shard_stripe_size is not None: + args.extend(["--shard-stripe-size", str(shard_stripe_size)]) + res = self.raw_cli(args) res.check_returncode() return tenant_id, timeline_id @@ -1636,6 +1676,19 @@ class NeonCli(AbstractNeonCli): return self.raw_cli(args, check_return_code=True) + def tenant_migrate( + self, tenant_shard_id: TenantShardId, new_pageserver: int, timeout_secs: Optional[int] + ): + args = [ + "tenant", + "migrate", + "--tenant-id", + str(tenant_shard_id), + "--id", + str(new_pageserver), + ] + return self.raw_cli(args, check_return_code=True, timeout=timeout_secs) + def start(self, check_return_code=True) -> "subprocess.CompletedProcess[str]": return self.raw_cli(["start"], check_return_code=check_return_code) @@ -1684,9 +1737,10 @@ class Pagectl(AbstractNeonCli): class NeonAttachmentService: - def __init__(self, env: NeonEnv): + def __init__(self, env: NeonEnv, auth_enabled): self.env = env self.running = False + self.auth_enabled = auth_enabled def start(self): assert not self.running @@ -1700,27 +1754,50 @@ class NeonAttachmentService: self.running = False return self - def attach_hook_issue(self, tenant_id: TenantId, pageserver_id: int) -> int: - response = requests.post( + def request(self, method, *args, **kwargs) -> requests.Response: + kwargs["headers"] = self.headers() + return requests.request(method, *args, **kwargs) + + def headers(self) -> Dict[str, str]: + headers = {} + if self.auth_enabled: + jwt_token = self.env.auth_keys.generate_pageserver_token() + headers["Authorization"] = f"Bearer {jwt_token}" + + return headers + + def attach_hook_issue( + self, tenant_shard_id: Union[TenantId, TenantShardId], pageserver_id: int + ) -> int: + response = self.request( + "POST", f"{self.env.control_plane_api}/attach-hook", - json={"tenant_id": str(tenant_id), "node_id": pageserver_id}, + json={"tenant_shard_id": str(tenant_shard_id), "node_id": pageserver_id}, + headers=self.headers(), ) response.raise_for_status() gen = response.json()["gen"] assert isinstance(gen, int) return gen - def attach_hook_drop(self, tenant_id: TenantId): - response = requests.post( + def attach_hook_drop(self, tenant_shard_id: Union[TenantId, TenantShardId]): + response = self.request( + "POST", f"{self.env.control_plane_api}/attach-hook", - json={"tenant_id": str(tenant_id), "node_id": None}, + json={"tenant_shard_id": str(tenant_shard_id), "node_id": None}, + headers=self.headers(), ) response.raise_for_status() - def inspect(self, tenant_id: TenantId) -> Optional[tuple[int, int]]: - response = requests.post( + def inspect(self, tenant_shard_id: Union[TenantId, TenantShardId]) -> Optional[tuple[int, int]]: + """ + :return: 2-tuple of (generation, pageserver id), or None if unknown + """ + response = self.request( + "POST", f"{self.env.control_plane_api}/inspect", - json={"tenant_id": str(tenant_id)}, + json={"tenant_shard_id": str(tenant_shard_id)}, + headers=self.headers(), ) response.raise_for_status() json = response.json() @@ -1731,6 +1808,79 @@ class NeonAttachmentService: else: return None + def node_register(self, node: NeonPageserver): + body = { + "node_id": int(node.id), + "listen_http_addr": "localhost", + "listen_http_port": node.service_port.http, + } + log.info(f"node_register({body})") + self.request( + "POST", f"{self.env.control_plane_api}/node", json=body, headers=self.headers() + ).raise_for_status() + + def tenant_create( + self, + tenant_id: TenantId, + shard_count: Optional[int] = None, + shard_stripe_size: Optional[int] = None, + tenant_config: Optional[Dict[Any, Any]] = None, + ): + body: Dict[str, Any] = {"new_tenant_id": str(tenant_id)} + + if shard_count is not None: + shard_params = {"count": shard_count} + if shard_stripe_size is not None: + shard_params["stripe_size"] = shard_stripe_size + + body["shard_parameters"] = shard_params + + if tenant_config is not None: + for k, v in tenant_config.items(): + body[k] = v + + response = self.request("POST", f"{self.env.control_plane_api}/tenant", json=body) + response.raise_for_status() + log.info(f"tenant_create success: {response.json()}") + + def tenant_timeline_create(self, tenant_id: TenantId, timeline_id: TimelineId): + body: Dict[str, Any] = {"new_timeline_id": str(timeline_id)} + + response = self.request( + "POST", f"{self.env.control_plane_api}/tenant/{tenant_id}/timeline", json=body + ) + response.raise_for_status() + log.info(f"tenant_timeline_create success: {response.json()}") + + def locate(self, tenant_id: TenantId) -> list[dict[str, Any]]: + response = self.request("GET", f"{self.env.control_plane_api}/tenant/{tenant_id}/locate") + response.raise_for_status() + body = response.json() + shards: list[dict[str, Any]] = body["shards"] + return shards + + def tenant_shard_split(self, tenant_id: TenantId, shard_count: int) -> list[TenantShardId]: + response = self.request( + "PUT", + f"{self.env.control_plane_api}/tenant/{tenant_id}/shard_split", + json={"new_shard_count": shard_count}, + ) + response.raise_for_status() + body = response.json() + log.info(f"tenant_shard_split success: {body}") + shards: list[TenantShardId] = body["new_shards"] + return shards + + def tenant_shard_migrate(self, tenant_shard_id: TenantShardId, dest_ps_id: int): + response = self.request( + "PUT", + f"{self.env.control_plane_api}/tenant/{tenant_shard_id}/migrate", + json={"tenant_shard_id": str(tenant_shard_id), "node_id": dest_ps_id}, + ) + response.raise_for_status() + log.info(f"Migrated tenant {tenant_shard_id} to pageserver {dest_ps_id}") + assert self.env.get_tenant_pageserver(tenant_shard_id).id == dest_ps_id + def __enter__(self) -> "NeonAttachmentService": return self @@ -2831,7 +2981,7 @@ class Endpoint(PgProtocol): hot_standby=hot_standby, lsn=lsn, pageserver_id=pageserver_id, - ).start(remote_ext_config=remote_ext_config) + ).start(remote_ext_config=remote_ext_config, pageserver_id=pageserver_id) log.info(f"Postgres startup took {time.time() - started_at} seconds") @@ -3344,7 +3494,7 @@ def pytest_addoption(parser: Parser): SMALL_DB_FILE_NAME_REGEX: re.Pattern = re.compile( # type: ignore[type-arg] - r"config|config-v1|heatmap-v1|metadata|.+\.(?:toml|pid|json|sql)" + r"config|config-v1|heatmap-v1|metadata|.+\.(?:toml|pid|json|sql|conf)" ) @@ -3481,9 +3631,7 @@ def list_files_to_compare(pgdata_dir: Path) -> List[str]: # pg is the existing and running compute node, that we want to compare with a basebackup -def check_restored_datadir_content( - test_output_dir: Path, env: NeonEnv, endpoint: Endpoint, pageserver_id: Optional[int] = None -): +def check_restored_datadir_content(test_output_dir: Path, env: NeonEnv, endpoint: Endpoint): # Get the timeline ID. We need it for the 'basebackup' command timeline_id = TimelineId(endpoint.safe_psql("SHOW neon.timeline_id")[0][0]) @@ -3504,6 +3652,7 @@ def check_restored_datadir_content( pg_bin = PgBin(test_output_dir, env.pg_distrib_dir, env.pg_version) psql_path = os.path.join(pg_bin.pg_bin_path, "psql") + pageserver_id = env.attachment_service.locate(endpoint.tenant_id)[0]["node_id"] cmd = rf""" {psql_path} \ --no-psqlrc \ @@ -3572,6 +3721,38 @@ def logical_replication_sync(subscriber: VanillaPostgres, publisher: Endpoint) - time.sleep(0.5) +def tenant_get_shards( + env: NeonEnv, tenant_id: TenantId, pageserver_id: Optional[int] +) -> list[tuple[TenantShardId, NeonPageserver]]: + """ + Helper for when you want to talk to one or more pageservers, and the + caller _might_ have specified a pageserver, or they might leave it to + us to figure out the shards for a tenant. + + If the caller provides `pageserver_id`, it will be used for all shards, even + if the shard is indicated by attachment service to be on some other pageserver. + + Caller should over the response to apply their per-pageserver action to + each shard + """ + if pageserver_id is not None: + override_pageserver = [p for p in env.pageservers if p.id == pageserver_id][0] + else: + override_pageserver = None + + if len(env.pageservers) > 1: + return [ + ( + TenantShardId.parse(s["shard_id"]), + override_pageserver or env.get_pageserver(s["node_id"]), + ) + for s in env.attachment_service.locate(tenant_id) + ] + else: + # Assume an unsharded tenant + return [(TenantShardId(tenant_id, 0, 0), override_pageserver or env.pageserver)] + + def wait_for_last_flush_lsn( env: NeonEnv, endpoint: Endpoint, @@ -3581,10 +3762,24 @@ def wait_for_last_flush_lsn( ) -> Lsn: """Wait for pageserver to catch up the latest flush LSN, returns the last observed lsn.""" + shards = tenant_get_shards(env, tenant, pageserver_id) + last_flush_lsn = Lsn(endpoint.safe_psql("SELECT pg_current_wal_flush_lsn()")[0][0]) - return wait_for_last_record_lsn( - env.get_pageserver(pageserver_id).http_client(), tenant, timeline, last_flush_lsn - ) + + results = [] + for tenant_shard_id, pageserver in shards: + log.info( + f"wait_for_last_flush_lsn: waiting for {last_flush_lsn} on shard {tenant_shard_id} on pageserver {pageserver.id})" + ) + waited = wait_for_last_record_lsn( + pageserver.http_client(), tenant_shard_id, timeline, last_flush_lsn + ) + + assert waited >= last_flush_lsn + results.append(waited) + + # Return the lowest LSN that has been ingested by all shards + return min(results) def wait_for_wal_insert_lsn( @@ -3596,9 +3791,16 @@ def wait_for_wal_insert_lsn( ) -> Lsn: """Wait for pageserver to catch up the latest flush LSN, returns the last observed lsn.""" last_flush_lsn = Lsn(endpoint.safe_psql("SELECT pg_current_wal_insert_lsn()")[0][0]) - return wait_for_last_record_lsn( - env.get_pageserver(pageserver_id).http_client(), tenant, timeline, last_flush_lsn - ) + result = None + for tenant_shard_id, pageserver in tenant_get_shards(env, tenant, pageserver_id): + shard_r = wait_for_last_record_lsn( + pageserver.http_client(), tenant_shard_id, timeline, last_flush_lsn + ) + if result is None: + result = shard_r + + assert result is not None + return result def fork_at_current_lsn( @@ -3632,11 +3834,13 @@ def last_flush_lsn_upload( last_flush_lsn = wait_for_last_flush_lsn( env, endpoint, tenant_id, timeline_id, pageserver_id=pageserver_id ) - ps_http = env.get_pageserver(pageserver_id).http_client() - wait_for_last_record_lsn(ps_http, tenant_id, timeline_id, last_flush_lsn) - # force a checkpoint to trigger upload - ps_http.timeline_checkpoint(tenant_id, timeline_id) - wait_for_upload(ps_http, tenant_id, timeline_id, last_flush_lsn) + shards = tenant_get_shards(env, tenant_id, pageserver_id) + for tenant_shard_id, pageserver in shards: + ps_http = pageserver.http_client() + wait_for_last_record_lsn(ps_http, tenant_shard_id, timeline_id, last_flush_lsn) + # force a checkpoint to trigger upload + ps_http.timeline_checkpoint(tenant_shard_id, timeline_id) + wait_for_upload(ps_http, tenant_shard_id, timeline_id, last_flush_lsn) return last_flush_lsn diff --git a/test_runner/fixtures/pageserver/http.py b/test_runner/fixtures/pageserver/http.py index b24de342f8..cfa2a2674d 100644 --- a/test_runner/fixtures/pageserver/http.py +++ b/test_runner/fixtures/pageserver/http.py @@ -4,7 +4,7 @@ import json import time from collections import defaultdict from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Any, Dict, List, Optional, Set, Tuple, Union import requests from requests.adapters import HTTPAdapter @@ -13,7 +13,7 @@ from urllib3.util.retry import Retry from fixtures.log_helper import log from fixtures.metrics import Metrics, parse_metrics from fixtures.pg_version import PgVersion -from fixtures.types import Lsn, TenantId, TimelineId +from fixtures.types import Lsn, TenantId, TenantShardId, TimelineId from fixtures.utils import Fn @@ -211,7 +211,7 @@ class PageserverHttpClient(requests.Session): def tenant_create( self, - new_tenant_id: TenantId, + new_tenant_id: Union[TenantId, TenantShardId], conf: Optional[Dict[str, Any]] = None, generation: Optional[int] = None, ) -> TenantId: @@ -239,7 +239,7 @@ class PageserverHttpClient(requests.Session): def tenant_attach( self, - tenant_id: TenantId, + tenant_id: Union[TenantId, TenantShardId], config: None | Dict[str, Any] = None, config_null: bool = False, generation: Optional[int] = None, @@ -269,7 +269,7 @@ class PageserverHttpClient(requests.Session): res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/detach", params=params) self.verbose_error(res) - def tenant_reset(self, tenant_id: TenantId, drop_cache: bool): + def tenant_reset(self, tenant_id: Union[TenantId, TenantShardId], drop_cache: bool): params = {} if drop_cache: params["drop_cache"] = "true" @@ -278,7 +278,7 @@ class PageserverHttpClient(requests.Session): self.verbose_error(res) def tenant_location_conf( - self, tenant_id: TenantId, location_conf=dict[str, Any], flush_ms=None + self, tenant_id: Union[TenantId, TenantShardId], location_conf=dict[str, Any], flush_ms=None ): body = location_conf.copy() body["tenant_id"] = str(tenant_id) @@ -294,7 +294,7 @@ class PageserverHttpClient(requests.Session): ) self.verbose_error(res) - def tenant_delete(self, tenant_id: TenantId): + def tenant_delete(self, tenant_id: Union[TenantId, TenantShardId]): res = self.delete(f"http://localhost:{self.port}/v1/tenant/{tenant_id}") self.verbose_error(res) return res @@ -310,27 +310,27 @@ class PageserverHttpClient(requests.Session): res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/ignore") self.verbose_error(res) - def tenant_status(self, tenant_id: TenantId) -> Dict[Any, Any]: + def tenant_status(self, tenant_id: Union[TenantId, TenantShardId]) -> Dict[Any, Any]: res = self.get(f"http://localhost:{self.port}/v1/tenant/{tenant_id}") self.verbose_error(res) res_json = res.json() assert isinstance(res_json, dict) return res_json - def tenant_config(self, tenant_id: TenantId) -> TenantConfig: + def tenant_config(self, tenant_id: Union[TenantId, TenantShardId]) -> TenantConfig: res = self.get(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/config") self.verbose_error(res) return TenantConfig.from_json(res.json()) - def tenant_heatmap_upload(self, tenant_id: TenantId): + def tenant_heatmap_upload(self, tenant_id: Union[TenantId, TenantShardId]): res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/heatmap_upload") self.verbose_error(res) - def tenant_secondary_download(self, tenant_id: TenantId): + def tenant_secondary_download(self, tenant_id: Union[TenantId, TenantShardId]): res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/secondary/download") self.verbose_error(res) - def set_tenant_config(self, tenant_id: TenantId, config: dict[str, Any]): + def set_tenant_config(self, tenant_id: Union[TenantId, TenantShardId], config: dict[str, Any]): assert "tenant_id" not in config.keys() res = self.put( f"http://localhost:{self.port}/v1/tenant/config", @@ -352,10 +352,12 @@ class PageserverHttpClient(requests.Session): del current[key] self.set_tenant_config(tenant_id, current) - def tenant_size(self, tenant_id: TenantId) -> int: + def tenant_size(self, tenant_id: Union[TenantId, TenantShardId]) -> int: return self.tenant_size_and_modelinputs(tenant_id)[0] - def tenant_size_and_modelinputs(self, tenant_id: TenantId) -> Tuple[int, Dict[str, Any]]: + def tenant_size_and_modelinputs( + self, tenant_id: Union[TenantId, TenantShardId] + ) -> Tuple[int, Dict[str, Any]]: """ Returns the tenant size, together with the model inputs as the second tuple item. """ @@ -370,7 +372,7 @@ class PageserverHttpClient(requests.Session): assert isinstance(inputs, dict) return (size, inputs) - def tenant_size_debug(self, tenant_id: TenantId) -> str: + def tenant_size_debug(self, tenant_id: Union[TenantId, TenantShardId]) -> str: """ Returns the tenant size debug info, as an HTML string """ @@ -382,7 +384,7 @@ class PageserverHttpClient(requests.Session): def timeline_list( self, - tenant_id: TenantId, + tenant_id: Union[TenantId, TenantShardId], include_non_incremental_logical_size: bool = False, include_timeline_dir_layer_file_size_sum: bool = False, ) -> List[Dict[str, Any]]: @@ -403,7 +405,7 @@ class PageserverHttpClient(requests.Session): def timeline_create( self, pg_version: PgVersion, - tenant_id: TenantId, + tenant_id: Union[TenantId, TenantShardId], new_timeline_id: TimelineId, ancestor_timeline_id: Optional[TimelineId] = None, ancestor_start_lsn: Optional[Lsn] = None, @@ -437,7 +439,7 @@ class PageserverHttpClient(requests.Session): def timeline_detail( self, - tenant_id: TenantId, + tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, include_non_incremental_logical_size: bool = False, include_timeline_dir_layer_file_size_sum: bool = False, @@ -462,7 +464,9 @@ class PageserverHttpClient(requests.Session): assert isinstance(res_json, dict) return res_json - def timeline_delete(self, tenant_id: TenantId, timeline_id: TimelineId, **kwargs): + def timeline_delete( + self, tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, **kwargs + ): """ Note that deletion is not instant, it is scheduled and performed mostly in the background. So if you need to wait for it to complete use `timeline_delete_wait_completed`. @@ -476,7 +480,10 @@ class PageserverHttpClient(requests.Session): assert res_json is None def timeline_gc( - self, tenant_id: TenantId, timeline_id: TimelineId, gc_horizon: Optional[int] + self, + tenant_id: Union[TenantId, TenantShardId], + timeline_id: TimelineId, + gc_horizon: Optional[int], ) -> dict[str, Any]: """ Unlike most handlers, this will wait for the layers to be actually @@ -499,7 +506,10 @@ class PageserverHttpClient(requests.Session): return res_json def timeline_compact( - self, tenant_id: TenantId, timeline_id: TimelineId, force_repartition=False + self, + tenant_id: Union[TenantId, TenantShardId], + timeline_id: TimelineId, + force_repartition=False, ): self.is_testing_enabled_or_skip() query = {} @@ -518,7 +528,7 @@ class PageserverHttpClient(requests.Session): def timeline_get_lsn_by_timestamp( self, - tenant_id: TenantId, + tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, timestamp, version: Optional[int] = None, @@ -537,7 +547,9 @@ class PageserverHttpClient(requests.Session): res_json = res.json() return res_json - def timeline_get_timestamp_of_lsn(self, tenant_id: TenantId, timeline_id: TimelineId, lsn: Lsn): + def timeline_get_timestamp_of_lsn( + self, tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, lsn: Lsn + ): log.info(f"Requesting time range of lsn {lsn}, tenant {tenant_id}, timeline {timeline_id}") res = self.get( f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/get_timestamp_of_lsn?lsn={lsn}", @@ -547,7 +559,10 @@ class PageserverHttpClient(requests.Session): return res_json def timeline_checkpoint( - self, tenant_id: TenantId, timeline_id: TimelineId, force_repartition=False + self, + tenant_id: Union[TenantId, TenantShardId], + timeline_id: TimelineId, + force_repartition=False, ): self.is_testing_enabled_or_skip() query = {} @@ -566,7 +581,7 @@ class PageserverHttpClient(requests.Session): def timeline_spawn_download_remote_layers( self, - tenant_id: TenantId, + tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, max_concurrent_downloads: int, ) -> dict[str, Any]: @@ -585,7 +600,7 @@ class PageserverHttpClient(requests.Session): def timeline_poll_download_remote_layers_status( self, - tenant_id: TenantId, + tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, spawn_response: dict[str, Any], poll_state=None, @@ -607,7 +622,7 @@ class PageserverHttpClient(requests.Session): def timeline_download_remote_layers( self, - tenant_id: TenantId, + tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, max_concurrent_downloads: int, errors_ok=False, @@ -689,9 +704,37 @@ class PageserverHttpClient(requests.Session): assert len(results) == 1, f"metric {name} with given filters is not unique, got: {results}" return results[0].value + def get_metrics_values( + self, names: list[str], filter: Optional[Dict[str, str]] = None + ) -> Dict[str, float]: + """ + When fetching multiple named metrics, it is more efficient to use this + than to call `get_metric_value` repeatedly. + + Throws RuntimeError if no metrics matching `names` are found, or if + not all of `names` are found: this method is intended for loading sets + of metrics whose existence is coupled. + """ + metrics = self.get_metrics() + samples = [] + for name in names: + samples.extend(metrics.query_all(name, filter=filter)) + + result = {} + for sample in samples: + if sample.name in result: + raise RuntimeError(f"Multiple values found for {sample.name}") + result[sample.name] = sample.value + + if len(result) != len(names): + log.info(f"Metrics found: {metrics.metrics}") + raise RuntimeError(f"could not find all metrics {' '.join(names)}") + + return result + def layer_map_info( self, - tenant_id: TenantId, + tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, ) -> LayerMapInfo: res = self.get( @@ -700,7 +743,9 @@ class PageserverHttpClient(requests.Session): self.verbose_error(res) return LayerMapInfo.from_json(res.json()) - def download_layer(self, tenant_id: TenantId, timeline_id: TimelineId, layer_name: str): + def download_layer( + self, tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, layer_name: str + ): res = self.get( f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/layer/{layer_name}", ) @@ -708,14 +753,18 @@ class PageserverHttpClient(requests.Session): assert res.status_code == 200 - def download_all_layers(self, tenant_id: TenantId, timeline_id: TimelineId): + def download_all_layers( + self, tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId + ): info = self.layer_map_info(tenant_id, timeline_id) for layer in info.historic_layers: if not layer.remote: continue self.download_layer(tenant_id, timeline_id, layer.layer_file_name) - def evict_layer(self, tenant_id: TenantId, timeline_id: TimelineId, layer_name: str): + def evict_layer( + self, tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, layer_name: str + ): res = self.delete( f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/layer/{layer_name}", ) @@ -723,7 +772,7 @@ class PageserverHttpClient(requests.Session): assert res.status_code in (200, 304) - def evict_all_layers(self, tenant_id: TenantId, timeline_id: TimelineId): + def evict_all_layers(self, tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId): info = self.layer_map_info(tenant_id, timeline_id) for layer in info.historic_layers: self.evict_layer(tenant_id, timeline_id, layer.layer_file_name) @@ -736,7 +785,7 @@ class PageserverHttpClient(requests.Session): self.verbose_error(res) return res.json() - def tenant_break(self, tenant_id: TenantId): + def tenant_break(self, tenant_id: Union[TenantId, TenantShardId]): res = self.put(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/break") self.verbose_error(res) diff --git a/test_runner/fixtures/pageserver/utils.py b/test_runner/fixtures/pageserver/utils.py index e7b78cfb9a..d0bb566408 100644 --- a/test_runner/fixtures/pageserver/utils.py +++ b/test_runner/fixtures/pageserver/utils.py @@ -1,12 +1,12 @@ import time -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from mypy_boto3_s3.type_defs import ListObjectsV2OutputTypeDef, ObjectTypeDef from fixtures.log_helper import log from fixtures.pageserver.http import PageserverApiException, PageserverHttpClient from fixtures.remote_storage import RemoteStorageKind, S3Storage -from fixtures.types import Lsn, TenantId, TimelineId +from fixtures.types import Lsn, TenantId, TenantShardId, TimelineId from fixtures.utils import wait_until @@ -22,7 +22,9 @@ def assert_tenant_state( def remote_consistent_lsn( - pageserver_http: PageserverHttpClient, tenant: TenantId, timeline: TimelineId + pageserver_http: PageserverHttpClient, + tenant: Union[TenantId, TenantShardId], + timeline: TimelineId, ) -> Lsn: detail = pageserver_http.timeline_detail(tenant, timeline) @@ -39,7 +41,7 @@ def remote_consistent_lsn( def wait_for_upload( pageserver_http: PageserverHttpClient, - tenant: TenantId, + tenant: Union[TenantId, TenantShardId], timeline: TimelineId, lsn: Lsn, ): @@ -92,7 +94,7 @@ def wait_until_tenant_state( def wait_until_timeline_state( pageserver_http: PageserverHttpClient, - tenant_id: TenantId, + tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, expected_state: str, iterations: int, @@ -141,7 +143,9 @@ def wait_until_tenant_active( def last_record_lsn( - pageserver_http_client: PageserverHttpClient, tenant: TenantId, timeline: TimelineId + pageserver_http_client: PageserverHttpClient, + tenant: Union[TenantId, TenantShardId], + timeline: TimelineId, ) -> Lsn: detail = pageserver_http_client.timeline_detail(tenant, timeline) @@ -152,7 +156,7 @@ def last_record_lsn( def wait_for_last_record_lsn( pageserver_http: PageserverHttpClient, - tenant: TenantId, + tenant: Union[TenantId, TenantShardId], timeline: TimelineId, lsn: Lsn, ) -> Lsn: @@ -194,7 +198,7 @@ def wait_for_upload_queue_empty( def wait_timeline_detail_404( pageserver_http: PageserverHttpClient, - tenant_id: TenantId, + tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, iterations: int, interval: Optional[float] = None, @@ -219,7 +223,7 @@ def wait_timeline_detail_404( def timeline_delete_wait_completed( pageserver_http: PageserverHttpClient, - tenant_id: TenantId, + tenant_id: Union[TenantId, TenantShardId], timeline_id: TimelineId, iterations: int = 20, interval: Optional[float] = None, diff --git a/test_runner/fixtures/workload.py b/test_runner/fixtures/workload.py index 241531437c..30def1194d 100644 --- a/test_runner/fixtures/workload.py +++ b/test_runner/fixtures/workload.py @@ -5,6 +5,7 @@ from fixtures.neon_fixtures import ( Endpoint, NeonEnv, last_flush_lsn_upload, + tenant_get_shards, wait_for_last_flush_lsn, ) from fixtures.pageserver.utils import wait_for_last_record_lsn, wait_for_upload @@ -31,7 +32,7 @@ class Workload: self._endpoint: Optional[Endpoint] = None - def endpoint(self, pageserver_id: int) -> Endpoint: + def endpoint(self, pageserver_id: Optional[int] = None) -> Endpoint: if self._endpoint is None: self._endpoint = self.env.endpoints.create( "main", @@ -54,7 +55,7 @@ class Workload: if self._endpoint is not None: self._endpoint.stop() - def init(self, pageserver_id: int): + def init(self, pageserver_id: Optional[int] = None): endpoint = self.endpoint(pageserver_id) endpoint.safe_psql(f"CREATE TABLE {self.table} (id INTEGER PRIMARY KEY, val text);") @@ -63,7 +64,7 @@ class Workload: self.env, endpoint, self.tenant_id, self.timeline_id, pageserver_id=pageserver_id ) - def write_rows(self, n, pageserver_id): + def write_rows(self, n, pageserver_id: Optional[int] = None): endpoint = self.endpoint(pageserver_id) start = self.expect_rows end = start + n - 1 @@ -81,7 +82,7 @@ class Workload: self.env, endpoint, self.tenant_id, self.timeline_id, pageserver_id=pageserver_id ) - def churn_rows(self, n, pageserver_id, upload=True): + def churn_rows(self, n, pageserver_id: Optional[int] = None, upload=True): assert self.expect_rows >= n max_iters = 10 @@ -119,21 +120,24 @@ class Workload: ] ) - last_flush_lsn = wait_for_last_flush_lsn( - self.env, endpoint, self.tenant_id, self.timeline_id, pageserver_id=pageserver_id - ) - ps_http = self.env.get_pageserver(pageserver_id).http_client() - wait_for_last_record_lsn(ps_http, self.tenant_id, self.timeline_id, last_flush_lsn) + for tenant_shard_id, pageserver in tenant_get_shards( + self.env, self.tenant_id, pageserver_id + ): + last_flush_lsn = wait_for_last_flush_lsn( + self.env, endpoint, self.tenant_id, self.timeline_id, pageserver_id=pageserver_id + ) + ps_http = pageserver.http_client() + wait_for_last_record_lsn(ps_http, tenant_shard_id, self.timeline_id, last_flush_lsn) - if upload: - # force a checkpoint to trigger upload - ps_http.timeline_checkpoint(self.tenant_id, self.timeline_id) - wait_for_upload(ps_http, self.tenant_id, self.timeline_id, last_flush_lsn) - log.info(f"Churn: waiting for remote LSN {last_flush_lsn}") - else: - log.info(f"Churn: not waiting for upload, disk LSN {last_flush_lsn}") + if upload: + # force a checkpoint to trigger upload + ps_http.timeline_checkpoint(tenant_shard_id, self.timeline_id) + wait_for_upload(ps_http, tenant_shard_id, self.timeline_id, last_flush_lsn) + log.info(f"Churn: waiting for remote LSN {last_flush_lsn}") + else: + log.info(f"Churn: not waiting for upload, disk LSN {last_flush_lsn}") - def validate(self, pageserver_id): + def validate(self, pageserver_id: Optional[int] = None): endpoint = self.endpoint(pageserver_id) result = endpoint.safe_psql_many( [ diff --git a/test_runner/performance/test_bulk_insert.py b/test_runner/performance/test_bulk_insert.py index edc23b29ba..72173dc2a7 100644 --- a/test_runner/performance/test_bulk_insert.py +++ b/test_runner/performance/test_bulk_insert.py @@ -61,7 +61,7 @@ def measure_recovery_time(env: NeonCompare): # of view, but the same as far as the safekeeper/WAL is concerned. To work around that, # we will explicitly create the tenant in the same generation that it was previously # attached in. - attach_status = env.env.attachment_service.inspect(tenant_id=env.tenant) + attach_status = env.env.attachment_service.inspect(tenant_shard_id=env.tenant) assert attach_status is not None (attach_gen, _) = attach_status diff --git a/test_runner/regress/test_broken_timeline.py b/test_runner/regress/test_broken_timeline.py index 4da0ba7b20..b046ed7f1b 100644 --- a/test_runner/regress/test_broken_timeline.py +++ b/test_runner/regress/test_broken_timeline.py @@ -10,6 +10,7 @@ from fixtures.neon_fixtures import ( NeonEnvBuilder, wait_for_last_flush_lsn, ) +from fixtures.pg_version import PgVersion from fixtures.types import TenantId, TimelineId @@ -126,7 +127,7 @@ def test_timeline_init_break_before_checkpoint(neon_env_builder: NeonEnvBuilder) # Introduce failpoint during timeline init (some intermediate files are on disk), before it's checkpointed. pageserver_http.configure_failpoints(("before-checkpoint-new-timeline", "return")) with pytest.raises(Exception, match="before-checkpoint-new-timeline"): - _ = env.neon_cli.create_timeline("test_timeline_init_break_before_checkpoint", tenant_id) + _ = pageserver_http.timeline_create(PgVersion.NOT_SET, tenant_id, TimelineId.generate()) # Restart the page server env.pageserver.restart(immediate=True) @@ -160,7 +161,7 @@ def test_timeline_init_break_before_checkpoint_recreate( ] ) - env.pageserver.tenant_create(env.initial_tenant) + env.neon_cli.create_tenant(env.initial_tenant) tenant_id = env.initial_tenant timelines_dir = env.pageserver.timeline_dir(tenant_id) @@ -216,7 +217,7 @@ def test_timeline_create_break_after_uninit_mark(neon_env_builder: NeonEnvBuilde # Introduce failpoint when creating a new timeline uninit mark, before any other files were created pageserver_http.configure_failpoints(("after-timeline-uninit-mark-creation", "return")) with pytest.raises(Exception, match="after-timeline-uninit-mark-creation"): - _ = env.neon_cli.create_timeline("test_timeline_create_break_after_uninit_mark", tenant_id) + _ = pageserver_http.timeline_create(PgVersion.NOT_SET, tenant_id, TimelineId.generate()) # Creating the timeline didn't finish. The other timelines on tenant should still be present and work normally. # "New" timeline is not present in the list, allowing pageserver to retry the same request diff --git a/test_runner/regress/test_disk_usage_eviction.py b/test_runner/regress/test_disk_usage_eviction.py index 9fdc4d59f5..64654c41a2 100644 --- a/test_runner/regress/test_disk_usage_eviction.py +++ b/test_runner/regress/test_disk_usage_eviction.py @@ -9,13 +9,14 @@ from fixtures.log_helper import log from fixtures.neon_fixtures import ( NeonEnv, NeonEnvBuilder, + NeonPageserver, PgBin, wait_for_last_flush_lsn, ) from fixtures.pageserver.http import PageserverHttpClient from fixtures.pageserver.utils import wait_for_upload_queue_empty from fixtures.remote_storage import RemoteStorageKind -from fixtures.types import Lsn, TenantId, TimelineId +from fixtures.types import Lsn, TenantId, TenantShardId, TimelineId from fixtures.utils import wait_until GLOBAL_LRU_LOG_LINE = "tenant_min_resident_size-respecting LRU would not relieve pressure, evicting more following global LRU policy" @@ -75,9 +76,15 @@ class EvictionOrder(str, enum.Enum): if self == EvictionOrder.ABSOLUTE_ORDER: return {"type": "AbsoluteAccessed"} elif self == EvictionOrder.RELATIVE_ORDER_EQUAL: - return {"type": "RelativeAccessed", "args": {"highest_layer_count_loses_first": False}} + return { + "type": "RelativeAccessed", + "args": {"highest_layer_count_loses_first": False}, + } elif self == EvictionOrder.RELATIVE_ORDER_SPARE: - return {"type": "RelativeAccessed", "args": {"highest_layer_count_loses_first": True}} + return { + "type": "RelativeAccessed", + "args": {"highest_layer_count_loses_first": True}, + } else: raise RuntimeError(f"not implemented: {self}") @@ -91,14 +98,24 @@ class EvictionEnv: layer_size: int pgbench_init_lsns: Dict[TenantId, Lsn] - def timelines_du(self) -> Tuple[int, int, int]: + @property + def pageserver(self): + """ + Shortcut for tests that only use one pageserver. + """ + return self.neon_env.pageserver + + def timelines_du(self, pageserver: NeonPageserver) -> Tuple[int, int, int]: return poor_mans_du( - self.neon_env, [(tid, tlid) for tid, tlid in self.timelines], verbose=False + self.neon_env, + [(tid, tlid) for tid, tlid in self.timelines], + pageserver, + verbose=False, ) - def du_by_timeline(self) -> Dict[Tuple[TenantId, TimelineId], int]: + def du_by_timeline(self, pageserver: NeonPageserver) -> Dict[Tuple[TenantId, TimelineId], int]: return { - (tid, tlid): poor_mans_du(self.neon_env, [(tid, tlid)], verbose=True)[0] + (tid, tlid): poor_mans_du(self.neon_env, [(tid, tlid)], pageserver, verbose=True)[0] for tid, tlid in self.timelines } @@ -126,7 +143,13 @@ class EvictionEnv: _avg = cur.fetchone() def pageserver_start_with_disk_usage_eviction( - self, period, max_usage_pct, min_avail_bytes, mock_behavior, eviction_order: EvictionOrder + self, + pageserver: NeonPageserver, + period, + max_usage_pct, + min_avail_bytes, + mock_behavior, + eviction_order: EvictionOrder, ): disk_usage_config = { "period": period, @@ -138,7 +161,12 @@ class EvictionEnv: enc = toml.TomlEncoder() - self.neon_env.pageserver.start( + # these can sometimes happen during startup before any tenants have been + # loaded, so nothing can be evicted, we just wait for next iteration which + # is able to evict. + pageserver.allowed_errors.append(".*WARN.* disk usage still high.*") + + pageserver.start( overrides=( "--pageserver-config-override=disk_usage_based_eviction=" + enc.dump_inline_table(disk_usage_config).replace("\n", " "), @@ -152,15 +180,10 @@ class EvictionEnv: ) def statvfs_called(): - assert self.neon_env.pageserver.log_contains(".*running mocked statvfs.*") + assert pageserver.log_contains(".*running mocked statvfs.*") wait_until(10, 1, statvfs_called) - # these can sometimes happen during startup before any tenants have been - # loaded, so nothing can be evicted, we just wait for next iteration which - # is able to evict. - self.neon_env.pageserver.allowed_errors.append(".*WARN.* disk usage still high.*") - def human_bytes(amt: float) -> str: suffixes = ["", "Ki", "Mi", "Gi"] @@ -175,23 +198,25 @@ def human_bytes(amt: float) -> str: raise RuntimeError("unreachable") -@pytest.fixture -def eviction_env(request, neon_env_builder: NeonEnvBuilder, pg_bin: PgBin) -> EvictionEnv: +def _eviction_env( + request, neon_env_builder: NeonEnvBuilder, pg_bin: PgBin, num_pageservers: int +) -> EvictionEnv: """ Creates two tenants, one somewhat larger than the other. """ log.info(f"setting up eviction_env for test {request.node.name}") + neon_env_builder.num_pageservers = num_pageservers neon_env_builder.enable_pageserver_remote_storage(RemoteStorageKind.LOCAL_FS) # initial tenant will not be present on this pageserver env = neon_env_builder.init_configs() env.start() - pageserver_http = env.pageserver.http_client() # allow because we are invoking this manually; we always warn on executing disk based eviction - env.pageserver.allowed_errors.append(r".* running disk usage based eviction due to pressure.*") + for ps in env.pageservers: + ps.allowed_errors.append(r".* running disk usage based eviction due to pressure.*") # Choose small layer_size so that we can use low pgbench_scales and still get a large count of layers. # Large count of layers and small layer size is good for testing because it makes evictions predictable. @@ -227,6 +252,8 @@ def eviction_env(request, neon_env_builder: NeonEnvBuilder, pg_bin: PgBin) -> Ev # after stopping the safekeepers, we know that no new WAL will be coming in for tenant_id, timeline_id in timelines: + pageserver_http = env.get_tenant_pageserver(tenant_id).http_client() + pageserver_http.timeline_checkpoint(tenant_id, timeline_id) wait_for_upload_queue_empty(pageserver_http, tenant_id, timeline_id) tl_info = pageserver_http.timeline_detail(tenant_id, timeline_id) @@ -252,6 +279,20 @@ def eviction_env(request, neon_env_builder: NeonEnvBuilder, pg_bin: PgBin) -> Ev return eviction_env +@pytest.fixture +def eviction_env(request, neon_env_builder: NeonEnvBuilder, pg_bin: PgBin) -> EvictionEnv: + return _eviction_env(request, neon_env_builder, pg_bin, num_pageservers=1) + + +@pytest.fixture +def eviction_env_ha(request, neon_env_builder: NeonEnvBuilder, pg_bin: PgBin) -> EvictionEnv: + """ + Variant of the eviction environment with two pageservers for testing eviction on + HA configurations with a secondary location. + """ + return _eviction_env(request, neon_env_builder, pg_bin, num_pageservers=2) + + def test_broken_tenants_are_skipped(eviction_env: EvictionEnv): env = eviction_env @@ -264,10 +305,16 @@ def test_broken_tenants_are_skipped(eviction_env: EvictionEnv): healthy_tenant_id, healthy_timeline_id = env.timelines[1] broken_size_pre, _, _ = poor_mans_du( - env.neon_env, [(broken_tenant_id, broken_timeline_id)], verbose=True + env.neon_env, + [(broken_tenant_id, broken_timeline_id)], + env.pageserver, + verbose=True, ) healthy_size_pre, _, _ = poor_mans_du( - env.neon_env, [(healthy_tenant_id, healthy_timeline_id)], verbose=True + env.neon_env, + [(healthy_tenant_id, healthy_timeline_id)], + env.pageserver, + verbose=True, ) # try to evict everything, then validate that broken tenant wasn't touched @@ -277,10 +324,16 @@ def test_broken_tenants_are_skipped(eviction_env: EvictionEnv): log.info(f"{response}") broken_size_post, _, _ = poor_mans_du( - env.neon_env, [(broken_tenant_id, broken_timeline_id)], verbose=True + env.neon_env, + [(broken_tenant_id, broken_timeline_id)], + env.pageserver, + verbose=True, ) healthy_size_post, _, _ = poor_mans_du( - env.neon_env, [(healthy_tenant_id, healthy_timeline_id)], verbose=True + env.neon_env, + [(healthy_tenant_id, healthy_timeline_id)], + env.pageserver, + verbose=True, ) assert broken_size_pre == broken_size_post, "broken tenant should not be touched" @@ -302,7 +355,7 @@ def test_pageserver_evicts_until_pressure_is_relieved( env = eviction_env pageserver_http = env.pageserver_http - (total_on_disk, _, _) = env.timelines_du() + (total_on_disk, _, _) = env.timelines_du(env.pageserver) target = total_on_disk // 2 @@ -311,7 +364,7 @@ def test_pageserver_evicts_until_pressure_is_relieved( ) log.info(f"{response}") - (later_total_on_disk, _, _) = env.timelines_du() + (later_total_on_disk, _, _) = env.timelines_du(env.pageserver) actual_change = total_on_disk - later_total_on_disk @@ -336,8 +389,8 @@ def test_pageserver_respects_overridden_resident_size( env = eviction_env ps_http = env.pageserver_http - (total_on_disk, _, _) = env.timelines_du() - du_by_timeline = env.du_by_timeline() + (total_on_disk, _, _) = env.timelines_du(env.pageserver) + du_by_timeline = env.du_by_timeline(env.pageserver) log.info("du_by_timeline: %s", du_by_timeline) assert len(du_by_timeline) == 2, "this test assumes two tenants" @@ -379,8 +432,8 @@ def test_pageserver_respects_overridden_resident_size( GLOBAL_LRU_LOG_LINE, ), "this test is pointless if it fell back to global LRU" - (later_total_on_disk, _, _) = env.timelines_du() - later_du_by_timeline = env.du_by_timeline() + (later_total_on_disk, _, _) = env.timelines_du(env.pageserver) + later_du_by_timeline = env.du_by_timeline(env.pageserver) log.info("later_du_by_timeline: %s", later_du_by_timeline) actual_change = total_on_disk - later_total_on_disk @@ -412,7 +465,7 @@ def test_pageserver_falls_back_to_global_lru(eviction_env: EvictionEnv, order: E env = eviction_env ps_http = env.pageserver_http - (total_on_disk, _, _) = env.timelines_du() + (total_on_disk, _, _) = env.timelines_du(env.pageserver) target = total_on_disk response = ps_http.disk_usage_eviction_run( @@ -420,7 +473,7 @@ def test_pageserver_falls_back_to_global_lru(eviction_env: EvictionEnv, order: E ) log.info(f"{response}") - (later_total_on_disk, _, _) = env.timelines_du() + (later_total_on_disk, _, _) = env.timelines_du(env.pageserver) actual_change = total_on_disk - later_total_on_disk assert 0 <= actual_change, "nothing can load layers during this test" assert actual_change >= target, "eviction must always evict more than target" @@ -448,8 +501,8 @@ def test_partial_evict_tenant(eviction_env: EvictionEnv, order: EvictionOrder): env = eviction_env ps_http = env.pageserver_http - (total_on_disk, _, _) = env.timelines_du() - du_by_timeline = env.du_by_timeline() + (total_on_disk, _, _) = env.timelines_du(env.pageserver) + du_by_timeline = env.du_by_timeline(env.pageserver) # pick smaller or greater (iteration order is insertion order of scale=4 and scale=6) [warm, cold] = list(du_by_timeline.keys()) @@ -467,12 +520,12 @@ def test_partial_evict_tenant(eviction_env: EvictionEnv, order: EvictionOrder): ) log.info(f"{response}") - (later_total_on_disk, _, _) = env.timelines_du() + (later_total_on_disk, _, _) = env.timelines_du(env.pageserver) actual_change = total_on_disk - later_total_on_disk assert 0 <= actual_change, "nothing can load layers during this test" assert actual_change >= target, "eviction must always evict more than target" - later_du_by_timeline = env.du_by_timeline() + later_du_by_timeline = env.du_by_timeline(env.pageserver) for tenant, later_tenant_usage in later_du_by_timeline.items(): assert ( later_tenant_usage < du_by_timeline[tenant] @@ -508,7 +561,10 @@ def test_partial_evict_tenant(eviction_env: EvictionEnv, order: EvictionOrder): def poor_mans_du( - env: NeonEnv, timelines: list[Tuple[TenantId, TimelineId]], verbose: bool = False + env: NeonEnv, + timelines: list[Tuple[TenantId, TimelineId]], + pageserver: NeonPageserver, + verbose: bool = False, ) -> Tuple[int, int, int]: """ Disk usage, largest, smallest layer for layer files over the given (tenant, timeline) tuples; @@ -518,7 +574,7 @@ def poor_mans_du( largest_layer = 0 smallest_layer = None for tenant_id, timeline_id in timelines: - timeline_dir = env.pageserver.timeline_dir(tenant_id, timeline_id) + timeline_dir = pageserver.timeline_dir(tenant_id, timeline_id) assert timeline_dir.exists(), f"timeline dir does not exist: {timeline_dir}" total = 0 for file in timeline_dir.iterdir(): @@ -549,6 +605,7 @@ def test_statvfs_error_handling(eviction_env: EvictionEnv): env = eviction_env env.neon_env.pageserver.stop() env.pageserver_start_with_disk_usage_eviction( + env.pageserver, period="1s", max_usage_pct=90, min_avail_bytes=0, @@ -573,11 +630,12 @@ def test_statvfs_pressure_usage(eviction_env: EvictionEnv): env.neon_env.pageserver.stop() # make it seem like we're at 100% utilization by setting total bytes to the used bytes - total_size, _, _ = env.timelines_du() + total_size, _, _ = env.timelines_du(env.pageserver) blocksize = 512 total_blocks = (total_size + (blocksize - 1)) // blocksize env.pageserver_start_with_disk_usage_eviction( + env.pageserver, period="1s", max_usage_pct=33, min_avail_bytes=0, @@ -597,7 +655,7 @@ def test_statvfs_pressure_usage(eviction_env: EvictionEnv): wait_until(10, 1, relieved_log_message) - post_eviction_total_size, _, _ = env.timelines_du() + post_eviction_total_size, _, _ = env.timelines_du(env.pageserver) assert post_eviction_total_size <= 0.33 * total_size, "we requested max 33% usage" @@ -612,13 +670,14 @@ def test_statvfs_pressure_min_avail_bytes(eviction_env: EvictionEnv): env.neon_env.pageserver.stop() # make it seem like we're at 100% utilization by setting total bytes to the used bytes - total_size, _, _ = env.timelines_du() + total_size, _, _ = env.timelines_du(env.pageserver) blocksize = 512 total_blocks = (total_size + (blocksize - 1)) // blocksize min_avail_bytes = total_size // 3 env.pageserver_start_with_disk_usage_eviction( + env.pageserver, period="1s", max_usage_pct=100, min_avail_bytes=min_avail_bytes, @@ -638,7 +697,77 @@ def test_statvfs_pressure_min_avail_bytes(eviction_env: EvictionEnv): wait_until(10, 1, relieved_log_message) - post_eviction_total_size, _, _ = env.timelines_du() + post_eviction_total_size, _, _ = env.timelines_du(env.pageserver) + + assert ( + total_size - post_eviction_total_size >= min_avail_bytes + ), "we requested at least min_avail_bytes worth of free space" + + +def test_secondary_mode_eviction(eviction_env_ha: EvictionEnv): + env = eviction_env_ha + + tenant_ids = [t[0] for t in env.timelines] + + # Set up a situation where one pageserver _only_ has secondary locations on it, + # so that when we release space we are sure it is via secondary locations. + + log.info("Setting up secondary location...") + ps_attached = env.neon_env.pageservers[0] + ps_secondary = env.neon_env.pageservers[1] + for tenant_id in tenant_ids: + # Migrate all attached tenants to the same pageserver, so that all the secondaries + # will run on the other pageserver. This is necessary because when we create tenants, + # they are spread over pageservers by default. + env.neon_env.attachment_service.tenant_shard_migrate( + TenantShardId(tenant_id, 0, 0), ps_attached.id + ) + + ps_secondary.tenant_location_configure( + tenant_id, + { + "mode": "Secondary", + "secondary_conf": {"warm": True}, + "tenant_conf": {}, + }, + ) + readback_conf = ps_secondary.read_tenant_location_conf(tenant_id) + log.info(f"Read back conf: {readback_conf}") + + # Request secondary location to download all layers that the attached location has + ps_attached.http_client().tenant_heatmap_upload(tenant_id) + ps_secondary.http_client().tenant_secondary_download(tenant_id) + + # Configure the secondary pageserver to have a phony small disk size + ps_secondary.stop() + total_size, _, _ = env.timelines_du(ps_secondary) + blocksize = 512 + total_blocks = (total_size + (blocksize - 1)) // blocksize + + min_avail_bytes = total_size // 3 + + env.pageserver_start_with_disk_usage_eviction( + ps_secondary, + period="1s", + max_usage_pct=100, + min_avail_bytes=min_avail_bytes, + mock_behavior={ + "type": "Success", + "blocksize": blocksize, + "total_blocks": total_blocks, + # Only count layer files towards used bytes in the mock_statvfs. + # This avoids accounting for metadata files & tenant conf in the tests. + "name_filter": ".*__.*", + }, + eviction_order=EvictionOrder.ABSOLUTE_ORDER, + ) + + def relieved_log_message(): + assert ps_secondary.log_contains(".*disk usage pressure relieved") + + wait_until(10, 1, relieved_log_message) + + post_eviction_total_size, _, _ = env.timelines_du(ps_secondary) assert ( total_size - post_eviction_total_size >= min_avail_bytes diff --git a/test_runner/regress/test_pageserver_generations.py b/test_runner/regress/test_pageserver_generations.py index 87a4fa01fc..dd55d737ac 100644 --- a/test_runner/regress/test_pageserver_generations.py +++ b/test_runner/regress/test_pageserver_generations.py @@ -62,6 +62,7 @@ def generate_uploads_and_deletions( tenant_id: Optional[TenantId] = None, timeline_id: Optional[TimelineId] = None, data: Optional[str] = None, + pageserver_id: Optional[int] = None, ): """ Using the environment's default tenant + timeline, generate a load pattern @@ -78,7 +79,9 @@ def generate_uploads_and_deletions( ps_http = env.pageserver.http_client() - with env.endpoints.create_start("main", tenant_id=tenant_id) as endpoint: + with env.endpoints.create_start( + "main", tenant_id=tenant_id, pageserver_id=pageserver_id + ) as endpoint: if init: endpoint.safe_psql("CREATE TABLE foo (id INTEGER PRIMARY KEY, val text)") last_flush_lsn_upload(env, endpoint, tenant_id, timeline_id) @@ -202,7 +205,7 @@ def test_generations_upgrade(neon_env_builder: NeonEnvBuilder): env.neon_cli.create_tenant( tenant_id=env.initial_tenant, conf=TENANT_CONF, timeline_id=env.initial_timeline ) - generate_uploads_and_deletions(env) + generate_uploads_and_deletions(env, pageserver_id=env.pageserver.id) def parse_generation_suffix(key): m = re.match(".+-([0-9a-zA-Z]{8})$", key) @@ -224,7 +227,7 @@ def test_generations_upgrade(neon_env_builder: NeonEnvBuilder): # Starting without the override that disabled control_plane_api env.pageserver.start() - generate_uploads_and_deletions(env, init=False) + generate_uploads_and_deletions(env, pageserver_id=env.pageserver.id, init=False) legacy_objects: list[str] = [] suffixed_objects = [] @@ -268,6 +271,7 @@ def test_deferred_deletion(neon_env_builder: NeonEnvBuilder): env = neon_env_builder.init_start(initial_tenant_conf=TENANT_CONF) some_other_pageserver = 1234 + ps_http = env.pageserver.http_client() generate_uploads_and_deletions(env) @@ -290,7 +294,7 @@ def test_deferred_deletion(neon_env_builder: NeonEnvBuilder): # Now advance the generation in the control plane: subsequent validations # from the running pageserver will fail. No more deletions should happen. env.attachment_service.attach_hook_issue(env.initial_tenant, some_other_pageserver) - generate_uploads_and_deletions(env, init=False) + generate_uploads_and_deletions(env, init=False, pageserver_id=env.pageserver.id) assert_deletion_queue(ps_http, lambda n: n > 0) queue_depth_before = get_deletion_queue_depth(ps_http) @@ -456,7 +460,7 @@ def test_emergency_mode(neon_env_builder: NeonEnvBuilder, pg_bin: PgBin): ps_http = env.pageserver.http_client() - generate_uploads_and_deletions(env) + generate_uploads_and_deletions(env, pageserver_id=env.pageserver.id) env.pageserver.allowed_errors.extend( [ @@ -473,7 +477,7 @@ def test_emergency_mode(neon_env_builder: NeonEnvBuilder, pg_bin: PgBin): # Remember how many validations had happened before the control plane went offline validated = get_deletion_queue_validated(ps_http) - generate_uploads_and_deletions(env, init=False) + generate_uploads_and_deletions(env, init=False, pageserver_id=env.pageserver.id) # The running pageserver should stop progressing deletions time.sleep(10) @@ -488,7 +492,7 @@ def test_emergency_mode(neon_env_builder: NeonEnvBuilder, pg_bin: PgBin): ) # The pageserver should provide service to clients - generate_uploads_and_deletions(env, init=False) + generate_uploads_and_deletions(env, init=False, pageserver_id=env.pageserver.id) # The pageserver should neither validate nor execute any deletions, it should have # loaded the DeletionLists from before though @@ -509,7 +513,7 @@ def test_emergency_mode(neon_env_builder: NeonEnvBuilder, pg_bin: PgBin): env.pageserver.stop() # Non-immediate: implicitly checking that shutdown doesn't hang waiting for CP env.pageserver.start() - generate_uploads_and_deletions(env, init=False) + generate_uploads_and_deletions(env, init=False, pageserver_id=env.pageserver.id) ps_http.deletion_queue_flush(execute=True) assert get_deletion_queue_depth(ps_http) == 0 assert get_deletion_queue_validated(ps_http) > 0 diff --git a/test_runner/regress/test_proxy.py b/test_runner/regress/test_proxy.py index 0f2cd9768f..1d62f09840 100644 --- a/test_runner/regress/test_proxy.py +++ b/test_runner/regress/test_proxy.py @@ -203,6 +203,21 @@ def test_close_on_connections_exit(static_proxy: NeonProxy): static_proxy.wait_for_exit() +def test_sql_over_http_serverless_driver(static_proxy: NeonProxy): + static_proxy.safe_psql("create role http with login password 'http' superuser") + + connstr = f"postgresql://http:http@{static_proxy.domain}:{static_proxy.proxy_port}/postgres" + response = requests.post( + f"https://api.localtest.me:{static_proxy.external_http_port}/sql", + data=json.dumps({"query": "select 42 as answer", "params": []}), + headers={"Content-Type": "application/sql", "Neon-Connection-String": connstr}, + verify=str(static_proxy.test_output_dir / "proxy.crt"), + ) + assert response.status_code == 200, response.text + rows = response.json()["rows"] + assert rows == [{"answer": 42}] + + def test_sql_over_http(static_proxy: NeonProxy): static_proxy.safe_psql("create role http with login password 'http' superuser") diff --git a/test_runner/regress/test_tenant_delete.py b/test_runner/regress/test_tenant_delete.py index 7c52fb3071..46dfa26a3f 100644 --- a/test_runner/regress/test_tenant_delete.py +++ b/test_runner/regress/test_tenant_delete.py @@ -556,6 +556,216 @@ def test_tenant_delete_concurrent( assert ps_http.get_metric_value("pageserver_tenant_manager_slots") == 0 +def test_tenant_delete_races_timeline_creation_01( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_02( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_03( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_04( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_05( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_06( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_07( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_08( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_09( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_10( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_11( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_12( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_13( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_14( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_15( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_16( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_17( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_18( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_19( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_20( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_21( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_22( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_23( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_24( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_25( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_26( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_27( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_28( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_29( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + +def test_tenant_delete_races_timeline_creation_30( + neon_env_builder: NeonEnvBuilder, + pg_bin: PgBin, +): + test_tenant_delete_races_timeline_creation(neon_env_builder, pg_bin) + + def test_tenant_delete_races_timeline_creation( neon_env_builder: NeonEnvBuilder, pg_bin: PgBin, diff --git a/test_runner/regress/test_tenant_relocation.py b/test_runner/regress/test_tenant_relocation.py index 1887bca23b..80b4fab1d3 100644 --- a/test_runner/regress/test_tenant_relocation.py +++ b/test_runner/regress/test_tenant_relocation.py @@ -67,6 +67,7 @@ def load(endpoint: Endpoint, stop_event: threading.Event, load_ok_event: threadi log.info("successfully recovered %s", inserted_ctr) failed = False load_ok_event.set() + log.info("load thread stopped") @@ -144,18 +145,14 @@ def check_timeline_attached( def switch_pg_to_new_pageserver( origin_ps: NeonPageserver, endpoint: Endpoint, - new_pageserver_port: int, + new_pageserver_id: int, tenant_id: TenantId, timeline_id: TimelineId, ) -> Path: + # We could reconfigure online with endpoint.reconfigure(), but this stop/start + # is needed to trigger the logic in load() to set its ok event after restart. endpoint.stop() - - pg_config_file_path = Path(endpoint.config_file_path()) - pg_config_file_path.open("a").write( - f"\nneon.pageserver_connstring = 'postgresql://no_user:@localhost:{new_pageserver_port}'" - ) - - endpoint.start() + endpoint.start(pageserver_id=new_pageserver_id) timeline_to_detach_local_path = origin_ps.timeline_dir(tenant_id, timeline_id) files_before_detach = os.listdir(timeline_to_detach_local_path) @@ -212,7 +209,7 @@ def test_tenant_relocation( env = neon_env_builder.init_start() - tenant_id = TenantId("74ee8b079a0e437eb0afea7d26a07209") + tenant_id = env.initial_tenant env.pageservers[0].allowed_errors.extend( [ @@ -236,8 +233,7 @@ def test_tenant_relocation( origin_http = origin_ps.http_client() destination_http = destination_ps.http_client() - _, initial_timeline_id = env.neon_cli.create_tenant(tenant_id) - log.info("tenant to relocate %s initial_timeline_id %s", tenant_id, initial_timeline_id) + log.info("tenant to relocate %s initial_timeline_id %s", tenant_id, env.initial_timeline) env.neon_cli.create_branch("test_tenant_relocation_main", tenant_id=tenant_id) ep_main = env.endpoints.create_start( @@ -380,7 +376,7 @@ def test_tenant_relocation( old_local_path_main = switch_pg_to_new_pageserver( origin_ps, ep_main, - destination_ps.service_port.pg, + destination_ps.id, tenant_id, timeline_id_main, ) @@ -388,7 +384,7 @@ def test_tenant_relocation( old_local_path_second = switch_pg_to_new_pageserver( origin_ps, ep_second, - destination_ps.service_port.pg, + destination_ps.id, tenant_id, timeline_id_second, ) diff --git a/test_runner/regress/test_tenants.py b/test_runner/regress/test_tenants.py index 5f2c1500d8..2ee2d8125a 100644 --- a/test_runner/regress/test_tenants.py +++ b/test_runner/regress/test_tenants.py @@ -214,14 +214,14 @@ def test_metrics_normal_work(neon_env_builder: NeonEnvBuilder): labels = ",".join([f'{key}="{value}"' for key, value in sample.labels.items()]) log.info(f"{sample.name}{{{labels}}} {sample.value}") - # Test that we gather tenant create metric + # Test that we gather tenant operations metrics storage_operation_metrics = [ "pageserver_storage_operations_seconds_global_bucket", "pageserver_storage_operations_seconds_global_sum", "pageserver_storage_operations_seconds_global_count", ] for metric in storage_operation_metrics: - value = ps_metrics.query_all(metric, filter={"operation": "create tenant"}) + value = ps_metrics.query_all(metric, filter={"operation": "layer flush"}) assert value diff --git a/test_runner/regress/test_timeline_size.py b/test_runner/regress/test_timeline_size.py index 92f2e72378..2e58a413e4 100644 --- a/test_runner/regress/test_timeline_size.py +++ b/test_runner/regress/test_timeline_size.py @@ -759,15 +759,7 @@ def test_ondemand_activation(neon_env_builder: NeonEnvBuilder): tenant_ids = {env.initial_tenant} for _i in range(0, n_tenants - 1): tenant_id = TenantId.generate() - env.pageserver.tenant_create(tenant_id) - - # Empty tenants are not subject to waiting for logical size calculations, because - # those hapen on timeline level - timeline_id = TimelineId.generate() - env.neon_cli.create_timeline( - new_branch_name="main", tenant_id=tenant_id, timeline_id=timeline_id - ) - + env.neon_cli.create_tenant(tenant_id) tenant_ids.add(tenant_id) # Restart pageserver with logical size calculations paused diff --git a/workspace_hack/Cargo.toml b/workspace_hack/Cargo.toml index 57aa1ef0bc..dbd46054a4 100644 --- a/workspace_hack/Cargo.toml +++ b/workspace_hack/Cargo.toml @@ -40,9 +40,11 @@ futures-io = { version = "0.3" } futures-sink = { version = "0.3" } futures-util = { version = "0.3", features = ["channel", "io", "sink"] } getrandom = { version = "0.2", default-features = false, features = ["std"] } +hashbrown = { version = "0.14", default-features = false, features = ["raw"] } hex = { version = "0.4", features = ["serde"] } hmac = { version = "0.12", default-features = false, features = ["reset"] } hyper = { version = "0.14", features = ["full"] } +indexmap = { version = "1", default-features = false, features = ["std"] } itertools = { version = "0.10" } libc = { version = "0.2", features = ["extra_traits"] } log = { version = "0.4", default-features = false, features = ["std"] } @@ -89,6 +91,8 @@ cc = { version = "1", default-features = false, features = ["parallel"] } chrono = { version = "0.4", default-features = false, features = ["clock", "serde", "wasmbind"] } either = { version = "1" } getrandom = { version = "0.2", default-features = false, features = ["std"] } +hashbrown = { version = "0.14", default-features = false, features = ["raw"] } +indexmap = { version = "1", default-features = false, features = ["std"] } itertools = { version = "0.10" } libc = { version = "0.2", features = ["extra_traits"] } log = { version = "0.4", default-features = false, features = ["std"] }