Merge remote-tracking branch 'origin/main' into problame/configurable-one-runtime

control_plane/attachment_service/Cargo.toml

@@ -1,52 +0,0 @@
-[package]
-name = "attachment_service"
-version = "0.1.0"
-edition.workspace = true
-license.workspace = true
-
-[[bin]]
-name = "storage_controller"
-path = "src/main.rs"
-
-[features]
-default = []
-# Enables test-only APIs and behaviors
-testing = []
-
-[dependencies]
-anyhow.workspace = true
-aws-config.workspace = true
-bytes.workspace = true
-camino.workspace = true
-clap.workspace = true
-fail.workspace = true
-futures.workspace = true
-git-version.workspace = true
-hex.workspace = true
-hyper.workspace = true
-humantime.workspace = true
-itertools.workspace = true
-lasso.workspace = true
-once_cell.workspace = true
-pageserver_api.workspace = true
-pageserver_client.workspace = true
-postgres_connection.workspace = true
-reqwest.workspace = true
-routerify.workspace = true
-serde.workspace = true
-serde_json.workspace = true
-thiserror.workspace = true
-tokio.workspace = true
-tokio-util.workspace = true
-tracing.workspace = true
-measured.workspace = true
-
-diesel = { version = "2.1.4", features = ["serde_json", "postgres", "r2d2"] }
-diesel_migrations = { version = "2.1.0" }
-r2d2 = { version = "0.8.10" }
-
-utils = { path = "../../libs/utils/" }
-metrics = { path = "../../libs/metrics/" }
-control_plane = { path = ".." }
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
-

control_plane/attachment_service/migrations/00000000000000_diesel_initial_setup/down.sql

@@ -1,6 +0,0 @@
--- This file was automatically created by Diesel to setup helper functions
--- and other internal bookkeeping. This file is safe to edit, any future
--- changes will be added to existing projects as new migrations.
-
-DROP FUNCTION IF EXISTS diesel_manage_updated_at(_tbl regclass);
-DROP FUNCTION IF EXISTS diesel_set_updated_at();

control_plane/attachment_service/migrations/00000000000000_diesel_initial_setup/up.sql

@@ -1,36 +0,0 @@
--- This file was automatically created by Diesel to setup helper functions
--- and other internal bookkeeping. This file is safe to edit, any future
--- changes will be added to existing projects as new migrations.
-
-
-
-
--- Sets up a trigger for the given table to automatically set a column called
--- `updated_at` whenever the row is modified (unless `updated_at` was included
--- in the modified columns)
---
--- # Example
---
--- ```sql
--- CREATE TABLE users (id SERIAL PRIMARY KEY, updated_at TIMESTAMP NOT NULL DEFAULT NOW());
---
--- SELECT diesel_manage_updated_at('users');
--- ```
-CREATE OR REPLACE FUNCTION diesel_manage_updated_at(_tbl regclass) RETURNS VOID AS $$
-BEGIN
-    EXECUTE format('CREATE TRIGGER set_updated_at BEFORE UPDATE ON %s
-                    FOR EACH ROW EXECUTE PROCEDURE diesel_set_updated_at()', _tbl);
-END;
-$$ LANGUAGE plpgsql;
-
-CREATE OR REPLACE FUNCTION diesel_set_updated_at() RETURNS trigger AS $$
-BEGIN
-    IF (
-        NEW IS DISTINCT FROM OLD AND
-        NEW.updated_at IS NOT DISTINCT FROM OLD.updated_at
-    ) THEN
-        NEW.updated_at := current_timestamp;
-    END IF;
-    RETURN NEW;
-END;
-$$ LANGUAGE plpgsql;

control_plane/attachment_service/migrations/2024-01-07-211257_create_tenant_shards/down.sql

@@ -1 +0,0 @@
-DROP TABLE tenant_shards;

control_plane/attachment_service/migrations/2024-01-07-211257_create_tenant_shards/up.sql

@@ -1,13 +0,0 @@
-CREATE TABLE tenant_shards (
-  tenant_id VARCHAR NOT NULL,
-  shard_number INTEGER NOT NULL,
-  shard_count INTEGER NOT NULL,
-  PRIMARY KEY(tenant_id, shard_number, shard_count),
-  shard_stripe_size INTEGER NOT NULL,
-  generation INTEGER NOT NULL,
-  generation_pageserver BIGINT NOT NULL,
-  placement_policy VARCHAR NOT NULL,
-  splitting SMALLINT NOT NULL,
-  -- config is JSON encoded, opaque to the database.
-  config TEXT NOT NULL
-);

control_plane/attachment_service/migrations/2024-01-07-212945_create_nodes/down.sql

@@ -1 +0,0 @@
-DROP TABLE nodes;

control_plane/attachment_service/migrations/2024-01-07-212945_create_nodes/up.sql

@@ -1,10 +0,0 @@
-CREATE TABLE nodes (
-  node_id BIGINT PRIMARY KEY NOT NULL,
-
-  scheduling_policy VARCHAR NOT NULL,
-
-  listen_http_addr VARCHAR NOT NULL,
-  listen_http_port INTEGER NOT NULL,
-  listen_pg_addr VARCHAR NOT NULL,
-  listen_pg_port INTEGER NOT NULL
-);

control_plane/attachment_service/migrations/2024-02-29-094122_generations_null/down.sql

@@ -1,2 +0,0 @@
-ALTER TABLE tenant_shards ALTER generation SET NOT NULL;
-ALTER TABLE tenant_shards ALTER generation_pageserver SET NOT NULL;

control_plane/attachment_service/migrations/2024-02-29-094122_generations_null/up.sql

@@ -1,4 +0,0 @@
-
-
-ALTER TABLE tenant_shards ALTER generation DROP NOT NULL;
-ALTER TABLE tenant_shards ALTER generation_pageserver DROP NOT NULL;

control_plane/attachment_service/migrations/2024-03-18-184429_rename_policy/down.sql

@@ -1,3 +0,0 @@
-
-UPDATE tenant_shards set placement_policy='{"Double": 1}' where placement_policy='{"Attached": 1}';
-UPDATE tenant_shards set placement_policy='"Single"' where placement_policy='{"Attached": 0}';

control_plane/attachment_service/migrations/2024-03-18-184429_rename_policy/up.sql

@@ -1,3 +0,0 @@
-
-UPDATE tenant_shards set placement_policy='{"Attached": 1}' where placement_policy='{"Double": 1}';
-UPDATE tenant_shards set placement_policy='{"Attached": 0}' where placement_policy='"Single"';

control_plane/attachment_service/migrations/2024-03-27-133204_tenant_policies/down.sql

@@ -1,3 +0,0 @@
--- This file should undo anything in `up.sql`
-
-ALTER TABLE tenant_shards drop scheduling_policy;

control_plane/attachment_service/migrations/2024-03-27-133204_tenant_policies/up.sql

@@ -1,2 +0,0 @@
-
-ALTER TABLE tenant_shards add scheduling_policy VARCHAR NOT NULL DEFAULT '"Active"';

control_plane/attachment_service/src/auth.rs

@@ -1,9 +0,0 @@
-use utils::auth::{AuthError, Claims, Scope};
-
-pub fn check_permission(claims: &Claims, required_scope: Scope) -> Result<(), AuthError> {
-    if claims.scope != required_scope {
-        return Err(AuthError("Scope mismatch. Permission denied".into()));
-    }
-
-    Ok(())
-}

control_plane/attachment_service/src/compute_hook.rs

@@ -1,465 +0,0 @@
-use std::{collections::HashMap, time::Duration};
-
-use control_plane::endpoint::{ComputeControlPlane, EndpointStatus};
-use control_plane::local_env::LocalEnv;
-use hyper::{Method, StatusCode};
-use pageserver_api::shard::{ShardCount, ShardNumber, ShardStripeSize, TenantShardId};
-use postgres_connection::parse_host_port;
-use serde::{Deserialize, Serialize};
-use tokio_util::sync::CancellationToken;
-use utils::{
-    backoff::{self},
-    id::{NodeId, TenantId},
-};
-
-use crate::service::Config;
-
-const SLOWDOWN_DELAY: Duration = Duration::from_secs(5);
-
-pub(crate) const API_CONCURRENCY: usize = 32;
-
-struct ShardedComputeHookTenant {
-    stripe_size: ShardStripeSize,
-    shard_count: ShardCount,
-    shards: Vec<(ShardNumber, NodeId)>,
-}
-
-enum ComputeHookTenant {
-    Unsharded(NodeId),
-    Sharded(ShardedComputeHookTenant),
-}
-
-impl ComputeHookTenant {
-    /// Construct with at least one shard's information
-    fn new(tenant_shard_id: TenantShardId, stripe_size: ShardStripeSize, node_id: NodeId) -> Self {
-        if tenant_shard_id.shard_count.count() > 1 {
-            Self::Sharded(ShardedComputeHookTenant {
-                shards: vec![(tenant_shard_id.shard_number, node_id)],
-                stripe_size,
-                shard_count: tenant_shard_id.shard_count,
-            })
-        } else {
-            Self::Unsharded(node_id)
-        }
-    }
-
-    /// Set one shard's location.  If stripe size or shard count have changed, Self is reset
-    /// and drops existing content.
-    fn update(
-        &mut self,
-        tenant_shard_id: TenantShardId,
-        stripe_size: ShardStripeSize,
-        node_id: NodeId,
-    ) {
-        match self {
-            Self::Unsharded(existing_node_id) if tenant_shard_id.shard_count.count() == 1 => {
-                *existing_node_id = node_id
-            }
-            Self::Sharded(sharded_tenant)
-                if sharded_tenant.stripe_size == stripe_size
-                    && sharded_tenant.shard_count == tenant_shard_id.shard_count =>
-            {
-                if let Some(existing) = sharded_tenant
-                    .shards
-                    .iter()
-                    .position(|s| s.0 == tenant_shard_id.shard_number)
-                {
-                    sharded_tenant.shards.get_mut(existing).unwrap().1 = node_id;
-                } else {
-                    sharded_tenant
-                        .shards
-                        .push((tenant_shard_id.shard_number, node_id));
-                    sharded_tenant.shards.sort_by_key(|s| s.0)
-                }
-            }
-            _ => {
-                // Shard count changed: reset struct.
-                *self = Self::new(tenant_shard_id, stripe_size, node_id);
-            }
-        }
-    }
-}
-
-#[derive(Serialize, Deserialize, Debug)]
-struct ComputeHookNotifyRequestShard {
-    node_id: NodeId,
-    shard_number: ShardNumber,
-}
-
-/// Request body that we send to the control plane to notify it of where a tenant is attached
-#[derive(Serialize, Deserialize, Debug)]
-struct ComputeHookNotifyRequest {
-    tenant_id: TenantId,
-    stripe_size: Option<ShardStripeSize>,
-    shards: Vec<ComputeHookNotifyRequestShard>,
-}
-
-/// Error type for attempts to call into the control plane compute notification hook
-#[derive(thiserror::Error, Debug)]
-pub(crate) enum NotifyError {
-    // Request was not send successfully, e.g. transport error
-    #[error("Sending request: {0}")]
-    Request(#[from] reqwest::Error),
-    // Request could not be serviced right now due to ongoing Operation in control plane, but should be possible soon.
-    #[error("Control plane tenant busy")]
-    Busy,
-    // Explicit 429 response asking us to retry less frequently
-    #[error("Control plane overloaded")]
-    SlowDown,
-    // A 503 response indicates the control plane can't handle the request right now
-    #[error("Control plane unavailable (status {0})")]
-    Unavailable(StatusCode),
-    // API returned unexpected non-success status.  We will retry, but log a warning.
-    #[error("Control plane returned unexpected status {0}")]
-    Unexpected(StatusCode),
-    // We shutdown while sending
-    #[error("Shutting down")]
-    ShuttingDown,
-    // A response indicates we will never succeed, such as 400 or 404
-    #[error("Non-retryable error {0}")]
-    Fatal(StatusCode),
-}
-
-impl ComputeHookTenant {
-    fn maybe_reconfigure(&self, tenant_id: TenantId) -> Option<ComputeHookNotifyRequest> {
-        match self {
-            Self::Unsharded(node_id) => Some(ComputeHookNotifyRequest {
-                tenant_id,
-                shards: vec![ComputeHookNotifyRequestShard {
-                    shard_number: ShardNumber(0),
-                    node_id: *node_id,
-                }],
-                stripe_size: None,
-            }),
-            Self::Sharded(sharded_tenant)
-                if sharded_tenant.shards.len() == sharded_tenant.shard_count.count() as usize =>
-            {
-                Some(ComputeHookNotifyRequest {
-                    tenant_id,
-                    shards: sharded_tenant
-                        .shards
-                        .iter()
-                        .map(|(shard_number, node_id)| ComputeHookNotifyRequestShard {
-                            shard_number: *shard_number,
-                            node_id: *node_id,
-                        })
-                        .collect(),
-                    stripe_size: Some(sharded_tenant.stripe_size),
-                })
-            }
-            Self::Sharded(sharded_tenant) => {
-                // Sharded tenant doesn't yet have information for all its shards
-
-                tracing::info!(
-                    "ComputeHookTenant::maybe_reconfigure: not enough shards ({}/{})",
-                    sharded_tenant.shards.len(),
-                    sharded_tenant.shard_count.count()
-                );
-                None
-            }
-        }
-    }
-}
-
-/// 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 {
-    config: Config,
-    state: tokio::sync::Mutex<HashMap<TenantId, ComputeHookTenant>>,
-    authorization_header: Option<String>,
-}
-
-impl ComputeHook {
-    pub(super) fn new(config: Config) -> Self {
-        let authorization_header = config
-            .control_plane_jwt_token
-            .clone()
-            .map(|jwt| format!("Bearer {}", jwt));
-
-        Self {
-            state: Default::default(),
-            config,
-            authorization_header,
-        }
-    }
-
-    /// For test environments: use neon_local's LocalEnv to update compute
-    async fn do_notify_local(
-        &self,
-        reconfigure_request: ComputeHookNotifyRequest,
-    ) -> anyhow::Result<()> {
-        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 ComputeHookNotifyRequest {
-            tenant_id,
-            shards,
-            stripe_size,
-        } = reconfigure_request;
-
-        let compute_pageservers = shards
-            .into_iter()
-            .map(|shard| {
-                let ps_conf = env
-                    .get_pageserver_conf(shard.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::<Vec<_>>();
-
-        for (endpoint_name, endpoint) in &cplane.endpoints {
-            if endpoint.tenant_id == tenant_id && endpoint.status() == EndpointStatus::Running {
-                tracing::info!("Reconfiguring endpoint {}", endpoint_name,);
-                endpoint
-                    .reconfigure(compute_pageservers.clone(), stripe_size)
-                    .await?;
-            }
-        }
-
-        Ok(())
-    }
-
-    async fn do_notify_iteration(
-        &self,
-        client: &reqwest::Client,
-        url: &String,
-        reconfigure_request: &ComputeHookNotifyRequest,
-        cancel: &CancellationToken,
-    ) -> Result<(), NotifyError> {
-        let req = client.request(Method::PUT, url);
-        let req = if let Some(value) = &self.authorization_header {
-            req.header(reqwest::header::AUTHORIZATION, value)
-        } else {
-            req
-        };
-
-        tracing::info!(
-            "Sending notify request to {} ({:?})",
-            url,
-            reconfigure_request
-        );
-        let send_result = req.json(&reconfigure_request).send().await;
-        let response = match send_result {
-            Ok(r) => r,
-            Err(e) => return Err(e.into()),
-        };
-
-        // Treat all 2xx responses as success
-        if response.status() >= StatusCode::OK && response.status() < StatusCode::MULTIPLE_CHOICES {
-            if response.status() != StatusCode::OK {
-                // Non-200 2xx response: it doesn't make sense to retry, but this is unexpected, so
-                // log a warning.
-                tracing::warn!(
-                    "Unexpected 2xx response code {} from control plane",
-                    response.status()
-                );
-            }
-
-            return Ok(());
-        }
-
-        // Error response codes
-        match response.status() {
-            StatusCode::TOO_MANY_REQUESTS => {
-                // TODO: 429 handling should be global: set some state visible to other requests
-                // so that they will delay before starting, rather than all notifications trying
-                // once before backing off.
-                tokio::time::timeout(SLOWDOWN_DELAY, cancel.cancelled())
-                    .await
-                    .ok();
-                Err(NotifyError::SlowDown)
-            }
-            StatusCode::LOCKED => {
-                // We consider this fatal, because it's possible that the operation blocking the control one is
-                // also the one that is waiting for this reconcile.  We should let the reconciler calling
-                // this hook fail, to give control plane a chance to un-lock.
-                tracing::info!("Control plane reports tenant is locked, dropping out of notify");
-                Err(NotifyError::Busy)
-            }
-            StatusCode::SERVICE_UNAVAILABLE
-            | StatusCode::GATEWAY_TIMEOUT
-            | StatusCode::BAD_GATEWAY => Err(NotifyError::Unavailable(response.status())),
-            StatusCode::BAD_REQUEST | StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
-                Err(NotifyError::Fatal(response.status()))
-            }
-            _ => Err(NotifyError::Unexpected(response.status())),
-        }
-    }
-
-    async fn do_notify(
-        &self,
-        url: &String,
-        reconfigure_request: ComputeHookNotifyRequest,
-        cancel: &CancellationToken,
-    ) -> Result<(), NotifyError> {
-        let client = reqwest::Client::new();
-        backoff::retry(
-            || self.do_notify_iteration(&client, url, &reconfigure_request, cancel),
-            |e| {
-                matches!(
-                    e,
-                    NotifyError::Fatal(_) | NotifyError::Unexpected(_) | NotifyError::Busy
-                )
-            },
-            3,
-            10,
-            "Send compute notification",
-            cancel,
-        )
-        .await
-        .ok_or_else(|| NotifyError::ShuttingDown)
-        .and_then(|x| x)
-    }
-
-    /// Call this to notify the compute (postgres) tier of new pageservers to use
-    /// for a tenant.  notify() is called by each shard individually, and this function
-    /// will decide whether an update to the tenant is sent.  An update is sent on the
-    /// condition that:
-    /// - We know a pageserver for every shard.
-    /// - All the shards have the same shard_count (i.e. we are not mid-split)
-    ///
-    /// Cancellation token enables callers to drop out, e.g. if calling from a Reconciler
-    /// that is cancelled.
-    ///
-    /// This function is fallible, including in the case that the control plane is transiently
-    /// unavailable.  A limited number of retries are done internally to efficiently hide short unavailability
-    /// periods, but we don't retry forever.  The **caller** is responsible for handling failures and
-    /// ensuring that they eventually call again to ensure that the compute is eventually notified of
-    /// the proper pageserver nodes for a tenant.
-    #[tracing::instrument(skip_all, fields(tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug(), node_id))]
-    pub(super) async fn notify(
-        &self,
-        tenant_shard_id: TenantShardId,
-        node_id: NodeId,
-        stripe_size: ShardStripeSize,
-        cancel: &CancellationToken,
-    ) -> Result<(), NotifyError> {
-        let mut locked = self.state.lock().await;
-
-        use std::collections::hash_map::Entry;
-        let tenant = match locked.entry(tenant_shard_id.tenant_id) {
-            Entry::Vacant(e) => e.insert(ComputeHookTenant::new(
-                tenant_shard_id,
-                stripe_size,
-                node_id,
-            )),
-            Entry::Occupied(e) => {
-                let tenant = e.into_mut();
-                tenant.update(tenant_shard_id, stripe_size, node_id);
-                tenant
-            }
-        };
-
-        let reconfigure_request = tenant.maybe_reconfigure(tenant_shard_id.tenant_id);
-        let Some(reconfigure_request) = reconfigure_request else {
-            // The tenant doesn't yet have pageservers for all its shards: we won't notify anything
-            // until it does.
-            tracing::info!("Tenant isn't yet ready to emit a notification");
-            return Ok(());
-        };
-
-        if let Some(notify_url) = &self.config.compute_hook_url {
-            self.do_notify(notify_url, reconfigure_request, cancel)
-                .await
-        } else {
-            self.do_notify_local(reconfigure_request)
-                .await
-                .map_err(|e| {
-                    // This path is for testing only, so munge the error into our prod-style error type.
-                    tracing::error!("Local notification hook failed: {e}");
-                    NotifyError::Fatal(StatusCode::INTERNAL_SERVER_ERROR)
-                })
-        }
-    }
-}
-
-#[cfg(test)]
-pub(crate) mod tests {
-    use pageserver_api::shard::{ShardCount, ShardNumber};
-    use utils::id::TenantId;
-
-    use super::*;
-
-    #[test]
-    fn tenant_updates() -> anyhow::Result<()> {
-        let tenant_id = TenantId::generate();
-        let mut tenant_state = ComputeHookTenant::new(
-            TenantShardId {
-                tenant_id,
-                shard_count: ShardCount::new(0),
-                shard_number: ShardNumber(0),
-            },
-            ShardStripeSize(12345),
-            NodeId(1),
-        );
-
-        // An unsharded tenant is always ready to emit a notification
-        assert!(tenant_state.maybe_reconfigure(tenant_id).is_some());
-        assert_eq!(
-            tenant_state
-                .maybe_reconfigure(tenant_id)
-                .unwrap()
-                .shards
-                .len(),
-            1
-        );
-        assert!(tenant_state
-            .maybe_reconfigure(tenant_id)
-            .unwrap()
-            .stripe_size
-            .is_none());
-
-        // Writing the first shard of a multi-sharded situation (i.e. in a split)
-        // resets the tenant state and puts it in an non-notifying state (need to
-        // see all shards)
-        tenant_state.update(
-            TenantShardId {
-                tenant_id,
-                shard_count: ShardCount::new(2),
-                shard_number: ShardNumber(1),
-            },
-            ShardStripeSize(32768),
-            NodeId(1),
-        );
-        assert!(tenant_state.maybe_reconfigure(tenant_id).is_none());
-
-        // Writing the second shard makes it ready to notify
-        tenant_state.update(
-            TenantShardId {
-                tenant_id,
-                shard_count: ShardCount::new(2),
-                shard_number: ShardNumber(0),
-            },
-            ShardStripeSize(32768),
-            NodeId(1),
-        );
-
-        assert!(tenant_state.maybe_reconfigure(tenant_id).is_some());
-        assert_eq!(
-            tenant_state
-                .maybe_reconfigure(tenant_id)
-                .unwrap()
-                .shards
-                .len(),
-            2
-        );
-        assert_eq!(
-            tenant_state
-                .maybe_reconfigure(tenant_id)
-                .unwrap()
-                .stripe_size,
-            Some(ShardStripeSize(32768))
-        );
-
-        Ok(())
-    }
-}

control_plane/attachment_service/src/heartbeater.rs

@@ -1,227 +0,0 @@
-use futures::{stream::FuturesUnordered, StreamExt};
-use std::{
-    collections::HashMap,
-    sync::Arc,
-    time::{Duration, Instant},
-};
-use tokio_util::sync::CancellationToken;
-
-use pageserver_api::{
-    controller_api::{NodeAvailability, UtilizationScore},
-    models::PageserverUtilization,
-};
-
-use thiserror::Error;
-use utils::id::NodeId;
-
-use crate::node::Node;
-
-struct HeartbeaterTask {
-    receiver: tokio::sync::mpsc::UnboundedReceiver<HeartbeatRequest>,
-    cancel: CancellationToken,
-
-    state: HashMap<NodeId, PageserverState>,
-
-    max_unavailable_interval: Duration,
-    jwt_token: Option<String>,
-}
-
-#[derive(Debug, Clone)]
-pub(crate) enum PageserverState {
-    Available {
-        last_seen_at: Instant,
-        utilization: PageserverUtilization,
-    },
-    Offline,
-}
-
-#[derive(Debug)]
-pub(crate) struct AvailablityDeltas(pub Vec<(NodeId, PageserverState)>);
-
-#[derive(Debug, Error)]
-pub(crate) enum HeartbeaterError {
-    #[error("Cancelled")]
-    Cancel,
-}
-
-struct HeartbeatRequest {
-    pageservers: Arc<HashMap<NodeId, Node>>,
-    reply: tokio::sync::oneshot::Sender<Result<AvailablityDeltas, HeartbeaterError>>,
-}
-
-pub(crate) struct Heartbeater {
-    sender: tokio::sync::mpsc::UnboundedSender<HeartbeatRequest>,
-}
-
-impl Heartbeater {
-    pub(crate) fn new(
-        jwt_token: Option<String>,
-        max_unavailable_interval: Duration,
-        cancel: CancellationToken,
-    ) -> Self {
-        let (sender, receiver) = tokio::sync::mpsc::unbounded_channel::<HeartbeatRequest>();
-        let mut heartbeater =
-            HeartbeaterTask::new(receiver, jwt_token, max_unavailable_interval, cancel);
-        tokio::task::spawn(async move { heartbeater.run().await });
-
-        Self { sender }
-    }
-
-    pub(crate) async fn heartbeat(
-        &self,
-        pageservers: Arc<HashMap<NodeId, Node>>,
-    ) -> Result<AvailablityDeltas, HeartbeaterError> {
-        let (sender, receiver) = tokio::sync::oneshot::channel();
-        self.sender
-            .send(HeartbeatRequest {
-                pageservers,
-                reply: sender,
-            })
-            .unwrap();
-
-        receiver.await.unwrap()
-    }
-}
-
-impl HeartbeaterTask {
-    fn new(
-        receiver: tokio::sync::mpsc::UnboundedReceiver<HeartbeatRequest>,
-        jwt_token: Option<String>,
-        max_unavailable_interval: Duration,
-        cancel: CancellationToken,
-    ) -> Self {
-        Self {
-            receiver,
-            cancel,
-            state: HashMap::new(),
-            max_unavailable_interval,
-            jwt_token,
-        }
-    }
-
-    async fn run(&mut self) {
-        loop {
-            tokio::select! {
-                request = self.receiver.recv() => {
-                    match request {
-                        Some(req) => {
-                            let res = self.heartbeat(req.pageservers).await;
-                            req.reply.send(res).unwrap();
-                        },
-                        None => { return; }
-                    }
-                },
-                _ = self.cancel.cancelled() => return
-            }
-        }
-    }
-
-    async fn heartbeat(
-        &mut self,
-        pageservers: Arc<HashMap<NodeId, Node>>,
-    ) -> Result<AvailablityDeltas, HeartbeaterError> {
-        let mut new_state = HashMap::new();
-
-        let mut heartbeat_futs = FuturesUnordered::new();
-        for (node_id, node) in &*pageservers {
-            heartbeat_futs.push({
-                let jwt_token = self.jwt_token.clone();
-                let cancel = self.cancel.clone();
-
-                // Clone the node and mark it as available such that the request
-                // goes through to the pageserver even when the node is marked offline.
-                // This doesn't impact the availability observed by [`crate::service::Service`].
-                let mut node = node.clone();
-                node.set_availability(NodeAvailability::Active(UtilizationScore::worst()));
-
-                async move {
-                    let response = node
-                        .with_client_retries(
-                            |client| async move { client.get_utilization().await },
-                            &jwt_token,
-                            3,
-                            3,
-                            Duration::from_secs(1),
-                            &cancel,
-                        )
-                        .await;
-
-                    let response = match response {
-                        Some(r) => r,
-                        None => {
-                            // This indicates cancellation of the request.
-                            // We ignore the node in this case.
-                            return None;
-                        }
-                    };
-
-                    let status = if let Ok(utilization) = response {
-                        PageserverState::Available {
-                            last_seen_at: Instant::now(),
-                            utilization,
-                        }
-                    } else {
-                        PageserverState::Offline
-                    };
-
-                    Some((*node_id, status))
-                }
-            });
-
-            loop {
-                let maybe_status = tokio::select! {
-                    next = heartbeat_futs.next() => {
-                        match next {
-                            Some(result) => result,
-                            None => { break; }
-                        }
-                    },
-                    _ = self.cancel.cancelled() => { return Err(HeartbeaterError::Cancel); }
-                };
-
-                if let Some((node_id, status)) = maybe_status {
-                    new_state.insert(node_id, status);
-                }
-            }
-        }
-
-        let mut deltas = Vec::new();
-        let now = Instant::now();
-        for (node_id, ps_state) in new_state {
-            use std::collections::hash_map::Entry::*;
-            let entry = self.state.entry(node_id);
-
-            let mut needs_update = false;
-            match entry {
-                Occupied(ref occ) => match (occ.get(), &ps_state) {
-                    (PageserverState::Offline, PageserverState::Offline) => {}
-                    (PageserverState::Available { last_seen_at, .. }, PageserverState::Offline) => {
-                        if now - *last_seen_at >= self.max_unavailable_interval {
-                            deltas.push((node_id, ps_state.clone()));
-                            needs_update = true;
-                        }
-                    }
-                    _ => {
-                        deltas.push((node_id, ps_state.clone()));
-                        needs_update = true;
-                    }
-                },
-                Vacant(_) => {
-                    deltas.push((node_id, ps_state.clone()));
-                }
-            }
-
-            match entry {
-                Occupied(mut occ) if needs_update => {
-                    (*occ.get_mut()) = ps_state;
-                }
-                Vacant(vac) => {
-                    vac.insert(ps_state);
-                }
-                _ => {}
-            }
-        }
-
-        Ok(AvailablityDeltas(deltas))
-    }
-}

control_plane/attachment_service/src/http.rs

@@ -1,894 +0,0 @@
-use crate::metrics::{
-    HttpRequestLatencyLabelGroup, HttpRequestStatusLabelGroup, PageserverRequestLabelGroup,
-    METRICS_REGISTRY,
-};
-use crate::reconciler::ReconcileError;
-use crate::service::{Service, STARTUP_RECONCILE_TIMEOUT};
-use futures::Future;
-use hyper::header::CONTENT_TYPE;
-use hyper::{Body, Request, Response};
-use hyper::{StatusCode, Uri};
-use pageserver_api::models::{
-    TenantConfigRequest, TenantCreateRequest, TenantLocationConfigRequest, TenantShardSplitRequest,
-    TenantTimeTravelRequest, TimelineCreateRequest,
-};
-use pageserver_api::shard::TenantShardId;
-use pageserver_client::mgmt_api;
-use std::sync::Arc;
-use std::time::{Duration, Instant};
-use tokio_util::sync::CancellationToken;
-use utils::auth::{Scope, SwappableJwtAuth};
-use utils::failpoint_support::failpoints_handler;
-use utils::http::endpoint::{auth_middleware, check_permission_with, request_span};
-use utils::http::request::{must_get_query_param, parse_query_param, parse_request_param};
-use utils::id::{TenantId, TimelineId};
-
-use utils::{
-    http::{
-        endpoint::{self},
-        error::ApiError,
-        json::{json_request, json_response},
-        RequestExt, RouterBuilder,
-    },
-    id::NodeId,
-};
-
-use pageserver_api::controller_api::{
-    NodeAvailability, NodeConfigureRequest, NodeRegisterRequest, TenantPolicyRequest,
-    TenantShardMigrateRequest,
-};
-use pageserver_api::upcall_api::{ReAttachRequest, ValidateRequest};
-
-use control_plane::storage_controller::{AttachHookRequest, InspectRequest};
-
-use routerify::Middleware;
-
-/// State available to HTTP request handlers
-#[derive(Clone)]
-pub struct HttpState {
-    service: Arc<crate::service::Service>,
-    auth: Option<Arc<SwappableJwtAuth>>,
-    allowlist_routes: Vec<Uri>,
-}
-
-impl HttpState {
-    pub fn new(service: Arc<crate::service::Service>, auth: Option<Arc<SwappableJwtAuth>>) -> Self {
-        let allowlist_routes = ["/status", "/ready", "/metrics"]
-            .iter()
-            .map(|v| v.parse().unwrap())
-            .collect::<Vec<_>>();
-        Self {
-            service,
-            auth,
-            allowlist_routes,
-        }
-    }
-}
-
-#[inline(always)]
-fn get_state(request: &Request<Body>) -> &HttpState {
-    request
-        .data::<Arc<HttpState>>()
-        .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<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::GenerationsApi)?;
-
-    let reattach_req = json_request::<ReAttachRequest>(&mut req).await?;
-    let state = get_state(&req);
-    json_response(StatusCode::OK, state.service.re_attach(reattach_req).await?)
-}
-
-/// 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<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::GenerationsApi)?;
-
-    let validate_req = json_request::<ValidateRequest>(&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<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let attach_req = json_request::<AttachHookRequest>(&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<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let inspect_req = json_request::<InspectRequest>(&mut req).await?;
-
-    let state = get_state(&req);
-
-    json_response(StatusCode::OK, state.service.inspect(inspect_req))
-}
-
-async fn handle_tenant_create(
-    service: Arc<Service>,
-    mut req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::PageServerApi)?;
-
-    let create_req = json_request::<TenantCreateRequest>(&mut req).await?;
-
-    json_response(
-        StatusCode::CREATED,
-        service.tenant_create(create_req).await?,
-    )
-}
-
-// For tenant and timeline deletions, which both implement an "initially return 202, then 404 once
-// we're done" semantic, we wrap with a retry loop to expose a simpler API upstream.  This avoids
-// needing to track a "deleting" state for tenants.
-async fn deletion_wrapper<R, F>(service: Arc<Service>, f: F) -> Result<Response<Body>, ApiError>
-where
-    R: std::future::Future<Output = Result<StatusCode, ApiError>> + Send + 'static,
-    F: Fn(Arc<Service>) -> R + Send + Sync + 'static,
-{
-    let started_at = Instant::now();
-    // To keep deletion reasonably snappy for small tenants, initially check after 1 second if deletion
-    // completed.
-    let mut retry_period = Duration::from_secs(1);
-    // On subsequent retries, wait longer.
-    let max_retry_period = Duration::from_secs(5);
-    // Enable callers with a 30 second request timeout to reliably get a response
-    let max_wait = Duration::from_secs(25);
-
-    loop {
-        let status = f(service.clone()).await?;
-        match status {
-            StatusCode::ACCEPTED => {
-                tracing::info!("Deletion accepted, waiting to try again...");
-                tokio::time::sleep(retry_period).await;
-                retry_period = max_retry_period;
-            }
-            StatusCode::NOT_FOUND => {
-                tracing::info!("Deletion complete");
-                return json_response(StatusCode::OK, ());
-            }
-            _ => {
-                tracing::warn!("Unexpected status {status}");
-                return json_response(status, ());
-            }
-        }
-
-        let now = Instant::now();
-        if now + retry_period > started_at + max_wait {
-            tracing::info!("Deletion timed out waiting for 404");
-            // REQUEST_TIMEOUT would be more appropriate, but CONFLICT is already part of
-            // the pageserver's swagger definition for this endpoint, and has the same desired
-            // effect of causing the control plane to retry later.
-            return json_response(StatusCode::CONFLICT, ());
-        }
-    }
-}
-
-async fn handle_tenant_location_config(
-    service: Arc<Service>,
-    mut req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    let tenant_shard_id: TenantShardId = parse_request_param(&req, "tenant_shard_id")?;
-    check_permissions(&req, Scope::PageServerApi)?;
-
-    let config_req = json_request::<TenantLocationConfigRequest>(&mut req).await?;
-    json_response(
-        StatusCode::OK,
-        service
-            .tenant_location_config(tenant_shard_id, config_req)
-            .await?,
-    )
-}
-
-async fn handle_tenant_config_set(
-    service: Arc<Service>,
-    mut req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::PageServerApi)?;
-
-    let config_req = json_request::<TenantConfigRequest>(&mut req).await?;
-
-    json_response(StatusCode::OK, service.tenant_config_set(config_req).await?)
-}
-
-async fn handle_tenant_config_get(
-    service: Arc<Service>,
-    req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
-    check_permissions(&req, Scope::PageServerApi)?;
-
-    json_response(StatusCode::OK, service.tenant_config_get(tenant_id)?)
-}
-
-async fn handle_tenant_time_travel_remote_storage(
-    service: Arc<Service>,
-    mut req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
-    check_permissions(&req, Scope::PageServerApi)?;
-
-    let time_travel_req = json_request::<TenantTimeTravelRequest>(&mut req).await?;
-
-    let timestamp_raw = must_get_query_param(&req, "travel_to")?;
-    let _timestamp = humantime::parse_rfc3339(&timestamp_raw).map_err(|_e| {
-        ApiError::BadRequest(anyhow::anyhow!(
-            "Invalid time for travel_to: {timestamp_raw:?}"
-        ))
-    })?;
-
-    let done_if_after_raw = must_get_query_param(&req, "done_if_after")?;
-    let _done_if_after = humantime::parse_rfc3339(&done_if_after_raw).map_err(|_e| {
-        ApiError::BadRequest(anyhow::anyhow!(
-            "Invalid time for done_if_after: {done_if_after_raw:?}"
-        ))
-    })?;
-
-    service
-        .tenant_time_travel_remote_storage(
-            &time_travel_req,
-            tenant_id,
-            timestamp_raw,
-            done_if_after_raw,
-        )
-        .await?;
-    json_response(StatusCode::OK, ())
-}
-
-async fn handle_tenant_secondary_download(
-    service: Arc<Service>,
-    req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
-    let wait = parse_query_param(&req, "wait_ms")?.map(Duration::from_millis);
-
-    let (status, progress) = service.tenant_secondary_download(tenant_id, wait).await?;
-    json_response(status, progress)
-}
-
-async fn handle_tenant_delete(
-    service: Arc<Service>,
-    req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
-    check_permissions(&req, Scope::PageServerApi)?;
-
-    deletion_wrapper(service, move |service| async move {
-        service.tenant_delete(tenant_id).await
-    })
-    .await
-}
-
-async fn handle_tenant_timeline_create(
-    service: Arc<Service>,
-    mut req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
-    check_permissions(&req, Scope::PageServerApi)?;
-
-    let create_req = json_request::<TimelineCreateRequest>(&mut req).await?;
-    json_response(
-        StatusCode::CREATED,
-        service
-            .tenant_timeline_create(tenant_id, create_req)
-            .await?,
-    )
-}
-
-async fn handle_tenant_timeline_delete(
-    service: Arc<Service>,
-    req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
-    check_permissions(&req, Scope::PageServerApi)?;
-
-    let timeline_id: TimelineId = parse_request_param(&req, "timeline_id")?;
-
-    deletion_wrapper(service, move |service| async move {
-        service.tenant_timeline_delete(tenant_id, timeline_id).await
-    })
-    .await
-}
-
-async fn handle_tenant_timeline_passthrough(
-    service: Arc<Service>,
-    req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
-    check_permissions(&req, Scope::PageServerApi)?;
-
-    let Some(path) = req.uri().path_and_query() else {
-        // This should never happen, our request router only calls us if there is a path
-        return Err(ApiError::BadRequest(anyhow::anyhow!("Missing path")));
-    };
-
-    tracing::info!("Proxying request for tenant {} ({})", tenant_id, path);
-
-    // Find the node that holds shard zero
-    let (node, tenant_shard_id) = service.tenant_shard0_node(tenant_id)?;
-
-    // Callers will always pass an unsharded tenant ID.  Before proxying, we must
-    // rewrite this to a shard-aware shard zero ID.
-    let path = format!("{}", path);
-    let tenant_str = tenant_id.to_string();
-    let tenant_shard_str = format!("{}", tenant_shard_id);
-    let path = path.replace(&tenant_str, &tenant_shard_str);
-
-    let latency = &METRICS_REGISTRY
-        .metrics_group
-        .storage_controller_passthrough_request_latency;
-
-    // This is a bit awkward. We remove the param from the request
-    // and join the words by '_' to get a label for the request.
-    let just_path = path.replace(&tenant_shard_str, "");
-    let path_label = just_path
-        .split('/')
-        .filter(|token| !token.is_empty())
-        .collect::<Vec<_>>()
-        .join("_");
-    let labels = PageserverRequestLabelGroup {
-        pageserver_id: &node.get_id().to_string(),
-        path: &path_label,
-        method: crate::metrics::Method::Get,
-    };
-
-    let _timer = latency.start_timer(labels.clone());
-
-    let client = mgmt_api::Client::new(node.base_url(), service.get_config().jwt_token.as_deref());
-    let resp = client.get_raw(path).await.map_err(|_e|
-        // FIXME: give APiError a proper Unavailable variant.  We return 503 here because
-        // if we can't successfully send a request to the pageserver, we aren't available.
-        ApiError::ShuttingDown)?;
-
-    if !resp.status().is_success() {
-        let error_counter = &METRICS_REGISTRY
-            .metrics_group
-            .storage_controller_passthrough_request_error;
-        error_counter.inc(labels);
-    }
-
-    // We have a reqest::Response, would like a http::Response
-    let mut builder = hyper::Response::builder()
-        .status(resp.status())
-        .version(resp.version());
-    for (k, v) in resp.headers() {
-        builder = builder.header(k, v);
-    }
-
-    let response = builder
-        .body(Body::wrap_stream(resp.bytes_stream()))
-        .map_err(|e| ApiError::InternalServerError(e.into()))?;
-
-    Ok(response)
-}
-
-async fn handle_tenant_locate(
-    service: Arc<Service>,
-    req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
-    json_response(StatusCode::OK, service.tenant_locate(tenant_id)?)
-}
-
-async fn handle_tenant_describe(
-    service: Arc<Service>,
-    req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
-    json_response(StatusCode::OK, service.tenant_describe(tenant_id)?)
-}
-
-async fn handle_tenant_list(
-    service: Arc<Service>,
-    req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    json_response(StatusCode::OK, service.tenant_list())
-}
-
-async fn handle_node_register(mut req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let register_req = json_request::<NodeRegisterRequest>(&mut req).await?;
-    let state = get_state(&req);
-    state.service.node_register(register_req).await?;
-    json_response(StatusCode::OK, ())
-}
-
-async fn handle_node_list(req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let state = get_state(&req);
-    let nodes = state.service.node_list().await?;
-    let api_nodes = nodes.into_iter().map(|n| n.describe()).collect::<Vec<_>>();
-
-    json_response(StatusCode::OK, api_nodes)
-}
-
-async fn handle_node_drop(req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let state = get_state(&req);
-    let node_id: NodeId = parse_request_param(&req, "node_id")?;
-    json_response(StatusCode::OK, state.service.node_drop(node_id).await?)
-}
-
-async fn handle_node_configure(mut req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let node_id: NodeId = parse_request_param(&req, "node_id")?;
-    let config_req = json_request::<NodeConfigureRequest>(&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.node_id,
-                config_req.availability.map(NodeAvailability::from),
-                config_req.scheduling,
-            )
-            .await?,
-    )
-}
-
-async fn handle_tenant_shard_split(
-    service: Arc<Service>,
-    mut req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
-    let split_req = json_request::<TenantShardSplitRequest>(&mut req).await?;
-
-    json_response(
-        StatusCode::OK,
-        service.tenant_shard_split(tenant_id, split_req).await?,
-    )
-}
-
-async fn handle_tenant_shard_migrate(
-    service: Arc<Service>,
-    mut req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let tenant_shard_id: TenantShardId = parse_request_param(&req, "tenant_shard_id")?;
-    let migrate_req = json_request::<TenantShardMigrateRequest>(&mut req).await?;
-    json_response(
-        StatusCode::OK,
-        service
-            .tenant_shard_migrate(tenant_shard_id, migrate_req)
-            .await?,
-    )
-}
-
-async fn handle_tenant_update_policy(mut req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
-    let update_req = json_request::<TenantPolicyRequest>(&mut req).await?;
-    let state = get_state(&req);
-
-    json_response(
-        StatusCode::OK,
-        state
-            .service
-            .tenant_update_policy(tenant_id, update_req)
-            .await?,
-    )
-}
-
-async fn handle_tenant_drop(req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
-    check_permissions(&req, Scope::PageServerApi)?;
-
-    let state = get_state(&req);
-
-    json_response(StatusCode::OK, state.service.tenant_drop(tenant_id).await?)
-}
-
-async fn handle_tenants_dump(req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let state = get_state(&req);
-    state.service.tenants_dump()
-}
-
-async fn handle_scheduler_dump(req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let state = get_state(&req);
-    state.service.scheduler_dump()
-}
-
-async fn handle_consistency_check(req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let state = get_state(&req);
-
-    json_response(StatusCode::OK, state.service.consistency_check().await?)
-}
-
-async fn handle_reconcile_all(req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let state = get_state(&req);
-
-    json_response(StatusCode::OK, state.service.reconcile_all_now().await?)
-}
-
-/// Status endpoint is just used for checking that our HTTP listener is up
-async fn handle_status(_req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    json_response(StatusCode::OK, ())
-}
-
-/// Readiness endpoint indicates when we're done doing startup I/O (e.g. reconciling
-/// with remote pageserver nodes).  This is intended for use as a kubernetes readiness probe.
-async fn handle_ready(req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    let state = get_state(&req);
-    if state.service.startup_complete.is_ready() {
-        json_response(StatusCode::OK, ())
-    } else {
-        json_response(StatusCode::SERVICE_UNAVAILABLE, ())
-    }
-}
-
-impl From<ReconcileError> for ApiError {
-    fn from(value: ReconcileError) -> Self {
-        ApiError::Conflict(format!("Reconciliation error: {}", value))
-    }
-}
-
-/// Common wrapper for request handlers that call into Service and will operate on tenants: they must only
-/// be allowed to run if Service has finished its initial reconciliation.
-async fn tenant_service_handler<R, H>(
-    request: Request<Body>,
-    handler: H,
-    request_name: RequestName,
-) -> R::Output
-where
-    R: std::future::Future<Output = Result<Response<Body>, ApiError>> + Send + 'static,
-    H: FnOnce(Arc<Service>, Request<Body>) -> R + Send + Sync + 'static,
-{
-    let state = get_state(&request);
-    let service = state.service.clone();
-
-    let startup_complete = service.startup_complete.clone();
-    if tokio::time::timeout(STARTUP_RECONCILE_TIMEOUT, startup_complete.wait())
-        .await
-        .is_err()
-    {
-        // This shouldn't happen: it is the responsibilty of [`Service::startup_reconcile`] to use appropriate
-        // timeouts around its remote calls, to bound its runtime.
-        return Err(ApiError::Timeout(
-            "Timed out waiting for service readiness".into(),
-        ));
-    }
-
-    named_request_span(
-        request,
-        |request| async move { handler(service, request).await },
-        request_name,
-    )
-    .await
-}
-
-/// Check if the required scope is held in the request's token, or if the request has
-/// a token with 'admin' scope then always permit it.
-fn check_permissions(request: &Request<Body>, required_scope: Scope) -> Result<(), ApiError> {
-    check_permission_with(request, |claims| {
-        match crate::auth::check_permission(claims, required_scope) {
-            Err(e) => match crate::auth::check_permission(claims, Scope::Admin) {
-                Ok(()) => Ok(()),
-                Err(_) => Err(e),
-            },
-            Ok(()) => Ok(()),
-        }
-    })
-}
-
-#[derive(Clone, Debug)]
-struct RequestMeta {
-    method: hyper::http::Method,
-    at: Instant,
-}
-
-fn prologue_metrics_middleware<B: hyper::body::HttpBody + Send + Sync + 'static>(
-) -> Middleware<B, ApiError> {
-    Middleware::pre(move |req| async move {
-        let meta = RequestMeta {
-            method: req.method().clone(),
-            at: Instant::now(),
-        };
-
-        req.set_context(meta);
-
-        Ok(req)
-    })
-}
-
-fn epilogue_metrics_middleware<B: hyper::body::HttpBody + Send + Sync + 'static>(
-) -> Middleware<B, ApiError> {
-    Middleware::post_with_info(move |resp, req_info| async move {
-        let request_name = match req_info.context::<RequestName>() {
-            Some(name) => name,
-            None => {
-                return Ok(resp);
-            }
-        };
-
-        if let Some(meta) = req_info.context::<RequestMeta>() {
-            let status = &crate::metrics::METRICS_REGISTRY
-                .metrics_group
-                .storage_controller_http_request_status;
-            let latency = &crate::metrics::METRICS_REGISTRY
-                .metrics_group
-                .storage_controller_http_request_latency;
-
-            status.inc(HttpRequestStatusLabelGroup {
-                path: request_name.0,
-                method: meta.method.clone().into(),
-                status: crate::metrics::StatusCode(resp.status()),
-            });
-
-            latency.observe(
-                HttpRequestLatencyLabelGroup {
-                    path: request_name.0,
-                    method: meta.method.into(),
-                },
-                meta.at.elapsed().as_secs_f64(),
-            );
-        }
-        Ok(resp)
-    })
-}
-
-pub async fn measured_metrics_handler(_req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    pub const TEXT_FORMAT: &str = "text/plain; version=0.0.4";
-
-    let payload = crate::metrics::METRICS_REGISTRY.encode();
-    let response = Response::builder()
-        .status(200)
-        .header(CONTENT_TYPE, TEXT_FORMAT)
-        .body(payload.into())
-        .unwrap();
-
-    Ok(response)
-}
-
-#[derive(Clone)]
-struct RequestName(&'static str);
-
-async fn named_request_span<R, H>(
-    request: Request<Body>,
-    handler: H,
-    name: RequestName,
-) -> R::Output
-where
-    R: Future<Output = Result<Response<Body>, ApiError>> + Send + 'static,
-    H: FnOnce(Request<Body>) -> R + Send + Sync + 'static,
-{
-    request.set_context(name);
-    request_span(request, handler).await
-}
-
-pub fn make_router(
-    service: Arc<Service>,
-    auth: Option<Arc<SwappableJwtAuth>>,
-) -> RouterBuilder<hyper::Body, ApiError> {
-    let mut router = endpoint::make_router()
-        .middleware(prologue_metrics_middleware())
-        .middleware(epilogue_metrics_middleware());
-    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("/metrics", |r| {
-            named_request_span(r, measured_metrics_handler, RequestName("metrics"))
-        })
-        // Non-prefixed generic endpoints (status, metrics)
-        .get("/status", |r| {
-            named_request_span(r, handle_status, RequestName("status"))
-        })
-        .get("/ready", |r| {
-            named_request_span(r, handle_ready, RequestName("ready"))
-        })
-        // Upcalls for the pageserver: point the pageserver's `control_plane_api` config to this prefix
-        .post("/upcall/v1/re-attach", |r| {
-            named_request_span(r, handle_re_attach, RequestName("upcall_v1_reattach"))
-        })
-        .post("/upcall/v1/validate", |r| {
-            named_request_span(r, handle_validate, RequestName("upcall_v1_validate"))
-        })
-        // Test/dev/debug endpoints
-        .post("/debug/v1/attach-hook", |r| {
-            named_request_span(r, handle_attach_hook, RequestName("debug_v1_attach_hook"))
-        })
-        .post("/debug/v1/inspect", |r| {
-            named_request_span(r, handle_inspect, RequestName("debug_v1_inspect"))
-        })
-        .post("/debug/v1/tenant/:tenant_id/drop", |r| {
-            named_request_span(r, handle_tenant_drop, RequestName("debug_v1_tenant_drop"))
-        })
-        .post("/debug/v1/node/:node_id/drop", |r| {
-            named_request_span(r, handle_node_drop, RequestName("debug_v1_node_drop"))
-        })
-        .get("/debug/v1/tenant", |r| {
-            named_request_span(r, handle_tenants_dump, RequestName("debug_v1_tenant"))
-        })
-        .get("/debug/v1/tenant/:tenant_id/locate", |r| {
-            tenant_service_handler(
-                r,
-                handle_tenant_locate,
-                RequestName("debug_v1_tenant_locate"),
-            )
-        })
-        .get("/debug/v1/scheduler", |r| {
-            named_request_span(r, handle_scheduler_dump, RequestName("debug_v1_scheduler"))
-        })
-        .post("/debug/v1/consistency_check", |r| {
-            named_request_span(
-                r,
-                handle_consistency_check,
-                RequestName("debug_v1_consistency_check"),
-            )
-        })
-        .post("/debug/v1/reconcile_all", |r| {
-            request_span(r, handle_reconcile_all)
-        })
-        .put("/debug/v1/failpoints", |r| {
-            request_span(r, |r| failpoints_handler(r, CancellationToken::new()))
-        })
-        // Node operations
-        .post("/control/v1/node", |r| {
-            named_request_span(r, handle_node_register, RequestName("control_v1_node"))
-        })
-        .get("/control/v1/node", |r| {
-            named_request_span(r, handle_node_list, RequestName("control_v1_node"))
-        })
-        .put("/control/v1/node/:node_id/config", |r| {
-            named_request_span(
-                r,
-                handle_node_configure,
-                RequestName("control_v1_node_config"),
-            )
-        })
-        // Tenant Shard operations
-        .put("/control/v1/tenant/:tenant_shard_id/migrate", |r| {
-            tenant_service_handler(
-                r,
-                handle_tenant_shard_migrate,
-                RequestName("control_v1_tenant_migrate"),
-            )
-        })
-        .put("/control/v1/tenant/:tenant_id/shard_split", |r| {
-            tenant_service_handler(
-                r,
-                handle_tenant_shard_split,
-                RequestName("control_v1_tenant_shard_split"),
-            )
-        })
-        .get("/control/v1/tenant/:tenant_id", |r| {
-            tenant_service_handler(
-                r,
-                handle_tenant_describe,
-                RequestName("control_v1_tenant_describe"),
-            )
-        })
-        .get("/control/v1/tenant", |r| {
-            tenant_service_handler(r, handle_tenant_list, RequestName("control_v1_tenant_list"))
-        })
-        .put("/control/v1/tenant/:tenant_id/policy", |r| {
-            named_request_span(
-                r,
-                handle_tenant_update_policy,
-                RequestName("control_v1_tenant_policy"),
-            )
-        })
-        // Tenant operations
-        // The ^/v1/ endpoints act as a "Virtual Pageserver", enabling shard-naive clients to call into
-        // this service to manage tenants that actually consist of many tenant shards, as if they are a single entity.
-        .post("/v1/tenant", |r| {
-            tenant_service_handler(r, handle_tenant_create, RequestName("v1_tenant"))
-        })
-        .delete("/v1/tenant/:tenant_id", |r| {
-            tenant_service_handler(r, handle_tenant_delete, RequestName("v1_tenant"))
-        })
-        .put("/v1/tenant/config", |r| {
-            tenant_service_handler(r, handle_tenant_config_set, RequestName("v1_tenant_config"))
-        })
-        .get("/v1/tenant/:tenant_id/config", |r| {
-            tenant_service_handler(r, handle_tenant_config_get, RequestName("v1_tenant_config"))
-        })
-        .put("/v1/tenant/:tenant_shard_id/location_config", |r| {
-            tenant_service_handler(
-                r,
-                handle_tenant_location_config,
-                RequestName("v1_tenant_location_config"),
-            )
-        })
-        .put("/v1/tenant/:tenant_id/time_travel_remote_storage", |r| {
-            tenant_service_handler(
-                r,
-                handle_tenant_time_travel_remote_storage,
-                RequestName("v1_tenant_time_travel_remote_storage"),
-            )
-        })
-        .post("/v1/tenant/:tenant_id/secondary/download", |r| {
-            tenant_service_handler(
-                r,
-                handle_tenant_secondary_download,
-                RequestName("v1_tenant_secondary_download"),
-            )
-        })
-        // Timeline operations
-        .delete("/v1/tenant/:tenant_id/timeline/:timeline_id", |r| {
-            tenant_service_handler(
-                r,
-                handle_tenant_timeline_delete,
-                RequestName("v1_tenant_timeline"),
-            )
-        })
-        .post("/v1/tenant/:tenant_id/timeline", |r| {
-            tenant_service_handler(
-                r,
-                handle_tenant_timeline_create,
-                RequestName("v1_tenant_timeline"),
-            )
-        })
-        // Tenant detail GET passthrough to shard zero
-        .get("/v1/tenant/:tenant_id", |r| {
-            tenant_service_handler(
-                r,
-                handle_tenant_timeline_passthrough,
-                RequestName("v1_tenant_passthrough"),
-            )
-        })
-        // Timeline GET passthrough to shard zero.  Note that the `*` in the URL is a wildcard: any future
-        // timeline GET APIs will be implicitly included.
-        .get("/v1/tenant/:tenant_id/timeline*", |r| {
-            tenant_service_handler(
-                r,
-                handle_tenant_timeline_passthrough,
-                RequestName("v1_tenant_timeline_passthrough"),
-            )
-        })
-}

control_plane/attachment_service/src/id_lock_map.rs

@@ -1,54 +0,0 @@
-use std::{collections::HashMap, sync::Arc};
-
-/// A map of locks covering some arbitrary identifiers. Useful if you have a collection of objects but don't
-/// want to embed a lock in each one, or if your locking granularity is different to your object granularity.
-/// For example, used in the storage controller where the objects are tenant shards, but sometimes locking
-/// is needed at a tenant-wide granularity.
-pub(crate) struct IdLockMap<T>
-where
-    T: Eq + PartialEq + std::hash::Hash,
-{
-    /// A synchronous lock for getting/setting the async locks that our callers will wait on.
-    entities: std::sync::Mutex<std::collections::HashMap<T, Arc<tokio::sync::RwLock<()>>>>,
-}
-
-impl<T> IdLockMap<T>
-where
-    T: Eq + PartialEq + std::hash::Hash,
-{
-    pub(crate) fn shared(
-        &self,
-        key: T,
-    ) -> impl std::future::Future<Output = tokio::sync::OwnedRwLockReadGuard<()>> {
-        let mut locked = self.entities.lock().unwrap();
-        let entry = locked.entry(key).or_default();
-        entry.clone().read_owned()
-    }
-
-    pub(crate) fn exclusive(
-        &self,
-        key: T,
-    ) -> impl std::future::Future<Output = tokio::sync::OwnedRwLockWriteGuard<()>> {
-        let mut locked = self.entities.lock().unwrap();
-        let entry = locked.entry(key).or_default();
-        entry.clone().write_owned()
-    }
-
-    /// Rather than building a lock guard that re-takes the [`Self::entities`] lock, we just do
-    /// periodic housekeeping to avoid the map growing indefinitely
-    pub(crate) fn housekeeping(&self) {
-        let mut locked = self.entities.lock().unwrap();
-        locked.retain(|_k, lock| lock.try_write().is_err())
-    }
-}
-
-impl<T> Default for IdLockMap<T>
-where
-    T: Eq + PartialEq + std::hash::Hash,
-{
-    fn default() -> Self {
-        Self {
-            entities: std::sync::Mutex::new(HashMap::new()),
-        }
-    }
-}

control_plane/attachment_service/src/lib.rs

@@ -1,54 +0,0 @@
-use serde::Serialize;
-use utils::seqwait::MonotonicCounter;
-
-mod auth;
-mod compute_hook;
-mod heartbeater;
-pub mod http;
-mod id_lock_map;
-pub mod metrics;
-mod node;
-mod pageserver_client;
-pub mod persistence;
-mod reconciler;
-mod scheduler;
-mod schema;
-pub mod service;
-mod tenant_state;
-
-#[derive(Ord, PartialOrd, Eq, PartialEq, Copy, Clone, Serialize)]
-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 std::fmt::Debug for Sequence {
-    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
-        write!(f, "{}", self.0)
-    }
-}
-
-impl MonotonicCounter<Sequence> 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)
-    }
-}

control_plane/attachment_service/src/main.rs

@@ -1,305 +0,0 @@
-use anyhow::{anyhow, Context};
-use attachment_service::http::make_router;
-use attachment_service::metrics::preinitialize_metrics;
-use attachment_service::persistence::Persistence;
-use attachment_service::service::{Config, Service, MAX_UNAVAILABLE_INTERVAL_DEFAULT};
-use camino::Utf8PathBuf;
-use clap::Parser;
-use diesel::Connection;
-use metrics::launch_timestamp::LaunchTimestamp;
-use std::sync::Arc;
-use tokio::signal::unix::SignalKind;
-use tokio_util::sync::CancellationToken;
-use utils::auth::{JwtAuth, SwappableJwtAuth};
-use utils::logging::{self, LogFormat};
-
-use utils::sentry_init::init_sentry;
-use utils::{project_build_tag, project_git_version, tcp_listener};
-
-project_git_version!(GIT_VERSION);
-project_build_tag!(BUILD_TAG);
-
-use diesel_migrations::{embed_migrations, EmbeddedMigrations};
-pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations");
-
-#[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,
-
-    /// Public key for JWT authentication of clients
-    #[arg(long)]
-    public_key: Option<String>,
-
-    /// Token for authenticating this service with the pageservers it controls
-    #[arg(long)]
-    jwt_token: Option<String>,
-
-    /// Token for authenticating this service with the control plane, when calling
-    /// the compute notification endpoint
-    #[arg(long)]
-    control_plane_jwt_token: Option<String>,
-
-    /// URL to control plane compute notification endpoint
-    #[arg(long)]
-    compute_hook_url: Option<String>,
-
-    /// Path to the .json file to store state (will be created if it doesn't exist)
-    #[arg(short, long)]
-    path: Option<Utf8PathBuf>,
-
-    /// URL to connect to postgres, like postgresql://localhost:1234/attachment_service
-    #[arg(long)]
-    database_url: Option<String>,
-
-    /// Flag to enable dev mode, which permits running without auth
-    #[arg(long, default_value = "false")]
-    dev: bool,
-
-    /// Grace period before marking unresponsive pageserver offline
-    #[arg(long)]
-    max_unavailable_interval: Option<humantime::Duration>,
-}
-
-enum StrictMode {
-    /// In strict mode, we will require that all secrets are loaded, i.e. security features
-    /// may not be implicitly turned off by omitting secrets in the environment.
-    Strict,
-    /// In dev mode, secrets are optional, and omitting a particular secret will implicitly
-    /// disable the auth related to it (e.g. no pageserver jwt key -> send unauthenticated
-    /// requests, no public key -> don't authenticate incoming requests).
-    Dev,
-}
-
-impl Default for StrictMode {
-    fn default() -> Self {
-        Self::Strict
-    }
-}
-
-/// Secrets may either be provided on the command line (for testing), or loaded from AWS SecretManager: this
-/// type encapsulates the logic to decide which and do the loading.
-struct Secrets {
-    database_url: String,
-    public_key: Option<JwtAuth>,
-    jwt_token: Option<String>,
-    control_plane_jwt_token: Option<String>,
-}
-
-impl Secrets {
-    const DATABASE_URL_ENV: &'static str = "DATABASE_URL";
-    const PAGESERVER_JWT_TOKEN_ENV: &'static str = "PAGESERVER_JWT_TOKEN";
-    const CONTROL_PLANE_JWT_TOKEN_ENV: &'static str = "CONTROL_PLANE_JWT_TOKEN";
-    const PUBLIC_KEY_ENV: &'static str = "PUBLIC_KEY";
-
-    /// Load secrets from, in order of preference:
-    /// - CLI args if database URL is provided on the CLI
-    /// - Environment variables if DATABASE_URL is set.
-    /// - AWS Secrets Manager secrets
-    async fn load(args: &Cli) -> anyhow::Result<Self> {
-        let Some(database_url) =
-            Self::load_secret(&args.database_url, Self::DATABASE_URL_ENV).await
-        else {
-            anyhow::bail!(
-                "Database URL is not set (set `--database-url`, or `DATABASE_URL` environment)"
-            )
-        };
-
-        let public_key = match Self::load_secret(&args.public_key, Self::PUBLIC_KEY_ENV).await {
-            Some(v) => Some(JwtAuth::from_key(v).context("Loading public key")?),
-            None => None,
-        };
-
-        let this = Self {
-            database_url,
-            public_key,
-            jwt_token: Self::load_secret(&args.jwt_token, Self::PAGESERVER_JWT_TOKEN_ENV).await,
-            control_plane_jwt_token: Self::load_secret(
-                &args.control_plane_jwt_token,
-                Self::CONTROL_PLANE_JWT_TOKEN_ENV,
-            )
-            .await,
-        };
-
-        Ok(this)
-    }
-
-    async fn load_secret(cli: &Option<String>, env_name: &str) -> Option<String> {
-        if let Some(v) = cli {
-            Some(v.clone())
-        } else if let Ok(v) = std::env::var(env_name) {
-            Some(v)
-        } else {
-            None
-        }
-    }
-}
-
-/// Execute the diesel migrations that are built into this binary
-async fn migration_run(database_url: &str) -> anyhow::Result<()> {
-    use diesel::PgConnection;
-    use diesel_migrations::{HarnessWithOutput, MigrationHarness};
-    let mut conn = PgConnection::establish(database_url)?;
-
-    HarnessWithOutput::write_to_stdout(&mut conn)
-        .run_pending_migrations(MIGRATIONS)
-        .map(|_| ())
-        .map_err(|e| anyhow::anyhow!(e))?;
-
-    Ok(())
-}
-
-fn main() -> anyhow::Result<()> {
-    let default_panic = std::panic::take_hook();
-    std::panic::set_hook(Box::new(move |info| {
-        default_panic(info);
-        std::process::exit(1);
-    }));
-
-    let _sentry_guard = init_sentry(Some(GIT_VERSION.into()), &[]);
-
-    tokio::runtime::Builder::new_current_thread()
-        // We use spawn_blocking for database operations, so require approximately
-        // as many blocking threads as we will open database connections.
-        .max_blocking_threads(Persistence::MAX_CONNECTIONS as usize)
-        .enable_all()
-        .build()
-        .unwrap()
-        .block_on(async_main())
-}
-
-async fn async_main() -> anyhow::Result<()> {
-    let launch_ts = Box::leak(Box::new(LaunchTimestamp::generate()));
-
-    logging::init(
-        LogFormat::Plain,
-        logging::TracingErrorLayerEnablement::Disabled,
-        logging::Output::Stdout,
-    )?;
-
-    preinitialize_metrics();
-
-    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.as_ref().unwrap_or(&Utf8PathBuf::from("<none>")),
-        args.listen
-    );
-
-    let strict_mode = if args.dev {
-        StrictMode::Dev
-    } else {
-        StrictMode::Strict
-    };
-
-    let secrets = Secrets::load(&args).await?;
-
-    // Validate required secrets and arguments are provided in strict mode
-    match strict_mode {
-        StrictMode::Strict
-            if (secrets.public_key.is_none()
-                || secrets.jwt_token.is_none()
-                || secrets.control_plane_jwt_token.is_none()) =>
-        {
-            // Production systems should always have secrets configured: if public_key was not set
-            // then we would implicitly disable auth.
-            anyhow::bail!(
-                    "Insecure config!  One or more secrets is not set.  This is only permitted in `--dev` mode"
-                );
-        }
-        StrictMode::Strict if args.compute_hook_url.is_none() => {
-            // Production systems should always have a compute hook set, to prevent falling
-            // back to trying to use neon_local.
-            anyhow::bail!(
-                "`--compute-hook-url` is not set: this is only permitted in `--dev` mode"
-            );
-        }
-        StrictMode::Strict => {
-            tracing::info!("Starting in strict mode: configuration is OK.")
-        }
-        StrictMode::Dev => {
-            tracing::warn!("Starting in dev mode: this may be an insecure configuration.")
-        }
-    }
-
-    let config = Config {
-        jwt_token: secrets.jwt_token,
-        control_plane_jwt_token: secrets.control_plane_jwt_token,
-        compute_hook_url: args.compute_hook_url,
-        max_unavailable_interval: args
-            .max_unavailable_interval
-            .map(humantime::Duration::into)
-            .unwrap_or(MAX_UNAVAILABLE_INTERVAL_DEFAULT),
-    };
-
-    // After loading secrets & config, but before starting anything else, apply database migrations
-    migration_run(&secrets.database_url)
-        .await
-        .context("Running database migrations")?;
-
-    let json_path = args.path;
-    let persistence = Arc::new(Persistence::new(secrets.database_url, json_path.clone()));
-
-    let service = Service::spawn(config, persistence.clone()).await?;
-
-    let http_listener = tcp_listener::bind(args.listen)?;
-
-    let auth = secrets
-        .public_key
-        .map(|jwt_auth| Arc::new(SwappableJwtAuth::new(jwt_auth)));
-    let router = make_router(service.clone(), auth)
-        .build()
-        .map_err(|err| anyhow!(err))?;
-    let router_service = utils::http::RouterService::new(router).unwrap();
-
-    // Start HTTP server
-    let server_shutdown = CancellationToken::new();
-    let server = hyper::Server::from_tcp(http_listener)?
-        .serve(router_service)
-        .with_graceful_shutdown({
-            let server_shutdown = server_shutdown.clone();
-            async move {
-                server_shutdown.cancelled().await;
-            }
-        });
-    tracing::info!("Serving on {0}", args.listen);
-    let server_task = tokio::task::spawn(server);
-
-    // Wait until we receive a signal
-    let mut sigint = tokio::signal::unix::signal(SignalKind::interrupt())?;
-    let mut sigquit = tokio::signal::unix::signal(SignalKind::quit())?;
-    let mut sigterm = tokio::signal::unix::signal(SignalKind::terminate())?;
-    tokio::select! {
-        _ = sigint.recv() => {},
-        _ = sigterm.recv() => {},
-        _ = sigquit.recv() => {},
-    }
-    tracing::info!("Terminating on signal");
-
-    if json_path.is_some() {
-        // Write out a JSON dump on shutdown: this is used in compat tests to avoid passing
-        // full postgres dumps around.
-        if let Err(e) = persistence.write_tenants_json().await {
-            tracing::error!("Failed to write JSON on shutdown: {e}")
-        }
-    }
-
-    // Stop HTTP server first, so that we don't have to service requests
-    // while shutting down Service
-    server_shutdown.cancel();
-    if let Err(e) = server_task.await {
-        tracing::error!("Error joining HTTP server task: {e}")
-    }
-    tracing::info!("Joined HTTP server task");
-
-    service.shutdown().await;
-    tracing::info!("Service shutdown complete");
-
-    std::process::exit(0);
-}

control_plane/attachment_service/src/metrics.rs

@@ -1,288 +0,0 @@
-//!
-//! This module provides metric definitions for the storage controller.
-//!
-//! All metrics are grouped in [`StorageControllerMetricGroup`]. [`StorageControllerMetrics`] holds
-//! the mentioned metrics and their encoder. It's globally available via the [`METRICS_REGISTRY`]
-//! constant.
-//!
-//! The rest of the code defines label group types and deals with converting outer types to labels.
-//!
-use bytes::Bytes;
-use measured::{
-    label::{LabelValue, StaticLabelSet},
-    FixedCardinalityLabel, MetricGroup,
-};
-use once_cell::sync::Lazy;
-use std::sync::Mutex;
-
-use crate::persistence::{DatabaseError, DatabaseOperation};
-
-pub(crate) static METRICS_REGISTRY: Lazy<StorageControllerMetrics> =
-    Lazy::new(StorageControllerMetrics::default);
-
-pub fn preinitialize_metrics() {
-    Lazy::force(&METRICS_REGISTRY);
-}
-
-pub(crate) struct StorageControllerMetrics {
-    pub(crate) metrics_group: StorageControllerMetricGroup,
-    encoder: Mutex<measured::text::TextEncoder>,
-}
-
-#[derive(measured::MetricGroup)]
-pub(crate) struct StorageControllerMetricGroup {
-    /// Count of how many times we spawn a reconcile task
-    pub(crate) storage_controller_reconcile_spawn: measured::Counter,
-    /// Reconciler tasks completed, broken down by success/failure/cancelled
-    pub(crate) storage_controller_reconcile_complete:
-        measured::CounterVec<ReconcileCompleteLabelGroupSet>,
-
-    /// Count of how many times we make an optimization change to a tenant's scheduling
-    pub(crate) storage_controller_schedule_optimization: measured::Counter,
-
-    /// HTTP request status counters for handled requests
-    pub(crate) storage_controller_http_request_status:
-        measured::CounterVec<HttpRequestStatusLabelGroupSet>,
-    /// HTTP request handler latency across all status codes
-    pub(crate) storage_controller_http_request_latency:
-        measured::HistogramVec<HttpRequestLatencyLabelGroupSet, 5>,
-
-    /// Count of HTTP requests to the pageserver that resulted in an error,
-    /// broken down by the pageserver node id, request name and method
-    pub(crate) storage_controller_pageserver_request_error:
-        measured::CounterVec<PageserverRequestLabelGroupSet>,
-
-    /// Latency of HTTP requests to the pageserver, broken down by pageserver
-    /// node id, request name and method. This include both successful and unsuccessful
-    /// requests.
-    pub(crate) storage_controller_pageserver_request_latency:
-        measured::HistogramVec<PageserverRequestLabelGroupSet, 5>,
-
-    /// Count of pass-through HTTP requests to the pageserver that resulted in an error,
-    /// broken down by the pageserver node id, request name and method
-    pub(crate) storage_controller_passthrough_request_error:
-        measured::CounterVec<PageserverRequestLabelGroupSet>,
-
-    /// Latency of pass-through HTTP requests to the pageserver, broken down by pageserver
-    /// node id, request name and method. This include both successful and unsuccessful
-    /// requests.
-    pub(crate) storage_controller_passthrough_request_latency:
-        measured::HistogramVec<PageserverRequestLabelGroupSet, 5>,
-
-    /// Count of errors in database queries, broken down by error type and operation.
-    pub(crate) storage_controller_database_query_error:
-        measured::CounterVec<DatabaseQueryErrorLabelGroupSet>,
-
-    /// Latency of database queries, broken down by operation.
-    pub(crate) storage_controller_database_query_latency:
-        measured::HistogramVec<DatabaseQueryLatencyLabelGroupSet, 5>,
-}
-
-impl StorageControllerMetrics {
-    pub(crate) fn encode(&self) -> Bytes {
-        let mut encoder = self.encoder.lock().unwrap();
-        self.metrics_group.collect_into(&mut *encoder);
-        encoder.finish()
-    }
-}
-
-impl Default for StorageControllerMetrics {
-    fn default() -> Self {
-        Self {
-            metrics_group: StorageControllerMetricGroup::new(),
-            encoder: Mutex::new(measured::text::TextEncoder::new()),
-        }
-    }
-}
-
-impl StorageControllerMetricGroup {
-    pub(crate) fn new() -> Self {
-        Self {
-            storage_controller_reconcile_spawn: measured::Counter::new(),
-            storage_controller_reconcile_complete: measured::CounterVec::new(
-                ReconcileCompleteLabelGroupSet {
-                    status: StaticLabelSet::new(),
-                },
-            ),
-            storage_controller_schedule_optimization: measured::Counter::new(),
-            storage_controller_http_request_status: measured::CounterVec::new(
-                HttpRequestStatusLabelGroupSet {
-                    path: lasso::ThreadedRodeo::new(),
-                    method: StaticLabelSet::new(),
-                    status: StaticLabelSet::new(),
-                },
-            ),
-            storage_controller_http_request_latency: measured::HistogramVec::new(
-                measured::metric::histogram::Thresholds::exponential_buckets(0.1, 2.0),
-            ),
-            storage_controller_pageserver_request_error: measured::CounterVec::new(
-                PageserverRequestLabelGroupSet {
-                    pageserver_id: lasso::ThreadedRodeo::new(),
-                    path: lasso::ThreadedRodeo::new(),
-                    method: StaticLabelSet::new(),
-                },
-            ),
-            storage_controller_pageserver_request_latency: measured::HistogramVec::new(
-                measured::metric::histogram::Thresholds::exponential_buckets(0.1, 2.0),
-            ),
-            storage_controller_passthrough_request_error: measured::CounterVec::new(
-                PageserverRequestLabelGroupSet {
-                    pageserver_id: lasso::ThreadedRodeo::new(),
-                    path: lasso::ThreadedRodeo::new(),
-                    method: StaticLabelSet::new(),
-                },
-            ),
-            storage_controller_passthrough_request_latency: measured::HistogramVec::new(
-                measured::metric::histogram::Thresholds::exponential_buckets(0.1, 2.0),
-            ),
-            storage_controller_database_query_error: measured::CounterVec::new(
-                DatabaseQueryErrorLabelGroupSet {
-                    operation: StaticLabelSet::new(),
-                    error_type: StaticLabelSet::new(),
-                },
-            ),
-            storage_controller_database_query_latency: measured::HistogramVec::new(
-                measured::metric::histogram::Thresholds::exponential_buckets(0.1, 2.0),
-            ),
-        }
-    }
-}
-
-#[derive(measured::LabelGroup)]
-#[label(set = ReconcileCompleteLabelGroupSet)]
-pub(crate) struct ReconcileCompleteLabelGroup {
-    pub(crate) status: ReconcileOutcome,
-}
-
-#[derive(measured::LabelGroup)]
-#[label(set = HttpRequestStatusLabelGroupSet)]
-pub(crate) struct HttpRequestStatusLabelGroup<'a> {
-    #[label(dynamic_with = lasso::ThreadedRodeo)]
-    pub(crate) path: &'a str,
-    pub(crate) method: Method,
-    pub(crate) status: StatusCode,
-}
-
-#[derive(measured::LabelGroup)]
-#[label(set = HttpRequestLatencyLabelGroupSet)]
-pub(crate) struct HttpRequestLatencyLabelGroup<'a> {
-    #[label(dynamic_with = lasso::ThreadedRodeo)]
-    pub(crate) path: &'a str,
-    pub(crate) method: Method,
-}
-
-impl Default for HttpRequestLatencyLabelGroupSet {
-    fn default() -> Self {
-        Self {
-            path: lasso::ThreadedRodeo::new(),
-            method: StaticLabelSet::new(),
-        }
-    }
-}
-
-#[derive(measured::LabelGroup, Clone)]
-#[label(set = PageserverRequestLabelGroupSet)]
-pub(crate) struct PageserverRequestLabelGroup<'a> {
-    #[label(dynamic_with = lasso::ThreadedRodeo)]
-    pub(crate) pageserver_id: &'a str,
-    #[label(dynamic_with = lasso::ThreadedRodeo)]
-    pub(crate) path: &'a str,
-    pub(crate) method: Method,
-}
-
-impl Default for PageserverRequestLabelGroupSet {
-    fn default() -> Self {
-        Self {
-            pageserver_id: lasso::ThreadedRodeo::new(),
-            path: lasso::ThreadedRodeo::new(),
-            method: StaticLabelSet::new(),
-        }
-    }
-}
-
-#[derive(measured::LabelGroup)]
-#[label(set = DatabaseQueryErrorLabelGroupSet)]
-pub(crate) struct DatabaseQueryErrorLabelGroup {
-    pub(crate) error_type: DatabaseErrorLabel,
-    pub(crate) operation: DatabaseOperation,
-}
-
-#[derive(measured::LabelGroup)]
-#[label(set = DatabaseQueryLatencyLabelGroupSet)]
-pub(crate) struct DatabaseQueryLatencyLabelGroup {
-    pub(crate) operation: DatabaseOperation,
-}
-
-#[derive(FixedCardinalityLabel)]
-pub(crate) enum ReconcileOutcome {
-    #[label(rename = "ok")]
-    Success,
-    Error,
-    Cancel,
-}
-
-#[derive(FixedCardinalityLabel, Clone)]
-pub(crate) enum Method {
-    Get,
-    Put,
-    Post,
-    Delete,
-    Other,
-}
-
-impl From<hyper::Method> for Method {
-    fn from(value: hyper::Method) -> Self {
-        if value == hyper::Method::GET {
-            Method::Get
-        } else if value == hyper::Method::PUT {
-            Method::Put
-        } else if value == hyper::Method::POST {
-            Method::Post
-        } else if value == hyper::Method::DELETE {
-            Method::Delete
-        } else {
-            Method::Other
-        }
-    }
-}
-
-pub(crate) struct StatusCode(pub(crate) hyper::http::StatusCode);
-
-impl LabelValue for StatusCode {
-    fn visit<V: measured::label::LabelVisitor>(&self, v: V) -> V::Output {
-        v.write_int(self.0.as_u16() as u64)
-    }
-}
-
-impl FixedCardinalityLabel for StatusCode {
-    fn cardinality() -> usize {
-        (100..1000).len()
-    }
-
-    fn encode(&self) -> usize {
-        self.0.as_u16() as usize
-    }
-
-    fn decode(value: usize) -> Self {
-        Self(hyper::http::StatusCode::from_u16(u16::try_from(value).unwrap()).unwrap())
-    }
-}
-
-#[derive(FixedCardinalityLabel)]
-pub(crate) enum DatabaseErrorLabel {
-    Query,
-    Connection,
-    ConnectionPool,
-    Logical,
-}
-
-impl DatabaseError {
-    pub(crate) fn error_label(&self) -> DatabaseErrorLabel {
-        match self {
-            Self::Query(_) => DatabaseErrorLabel::Query,
-            Self::Connection(_) => DatabaseErrorLabel::Connection,
-            Self::ConnectionPool(_) => DatabaseErrorLabel::ConnectionPool,
-            Self::Logical(_) => DatabaseErrorLabel::Logical,
-        }
-    }
-}

control_plane/attachment_service/src/node.rs

@@ -1,285 +0,0 @@
-use std::{str::FromStr, time::Duration};
-
-use hyper::StatusCode;
-use pageserver_api::{
-    controller_api::{
-        NodeAvailability, NodeDescribeResponse, NodeRegisterRequest, NodeSchedulingPolicy,
-        TenantLocateResponseShard,
-    },
-    shard::TenantShardId,
-};
-use pageserver_client::mgmt_api;
-use serde::Serialize;
-use tokio_util::sync::CancellationToken;
-use utils::{backoff, id::NodeId};
-
-use crate::{
-    pageserver_client::PageserverClient, persistence::NodePersistence, scheduler::MaySchedule,
-};
-
-/// Represents the in-memory description of a Node.
-///
-/// Scheduling statistics are maintened separately in [`crate::scheduler`].
-///
-/// The persistent subset of the Node is defined in [`crate::persistence::NodePersistence`]: the
-/// implementation of serialization on this type is only for debug dumps.
-#[derive(Clone, Serialize)]
-pub(crate) struct Node {
-    id: NodeId,
-
-    availability: NodeAvailability,
-    scheduling: NodeSchedulingPolicy,
-
-    listen_http_addr: String,
-    listen_http_port: u16,
-
-    listen_pg_addr: String,
-    listen_pg_port: u16,
-
-    // This cancellation token means "stop any RPCs in flight to this node, and don't start
-    // any more". It is not related to process shutdown.
-    #[serde(skip)]
-    cancel: CancellationToken,
-}
-
-/// When updating [`Node::availability`] we use this type to indicate to the caller
-/// whether/how they changed it.
-pub(crate) enum AvailabilityTransition {
-    ToActive,
-    ToOffline,
-    Unchanged,
-}
-
-impl Node {
-    pub(crate) fn base_url(&self) -> String {
-        format!("http://{}:{}", self.listen_http_addr, self.listen_http_port)
-    }
-
-    pub(crate) fn get_id(&self) -> NodeId {
-        self.id
-    }
-
-    pub(crate) fn set_scheduling(&mut self, scheduling: NodeSchedulingPolicy) {
-        self.scheduling = scheduling
-    }
-
-    /// Does this registration request match `self`?  This is used when deciding whether a registration
-    /// request should be allowed to update an existing record with the same node ID.
-    pub(crate) fn registration_match(&self, register_req: &NodeRegisterRequest) -> bool {
-        self.id == register_req.node_id
-            && self.listen_http_addr == register_req.listen_http_addr
-            && self.listen_http_port == register_req.listen_http_port
-            && self.listen_pg_addr == register_req.listen_pg_addr
-            && self.listen_pg_port == register_req.listen_pg_port
-    }
-
-    /// For a shard located on this node, populate a response object
-    /// with this node's address information.
-    pub(crate) fn shard_location(&self, shard_id: TenantShardId) -> TenantLocateResponseShard {
-        TenantLocateResponseShard {
-            shard_id,
-            node_id: self.id,
-            listen_http_addr: self.listen_http_addr.clone(),
-            listen_http_port: self.listen_http_port,
-            listen_pg_addr: self.listen_pg_addr.clone(),
-            listen_pg_port: self.listen_pg_port,
-        }
-    }
-
-    pub(crate) fn set_availability(&mut self, availability: NodeAvailability) {
-        match self.get_availability_transition(availability) {
-            AvailabilityTransition::ToActive => {
-                // Give the node a new cancellation token, effectively resetting it to un-cancelled.  Any
-                // users of previously-cloned copies of the node will still see the old cancellation
-                // state.  For example, Reconcilers in flight will have to complete and be spawned
-                // again to realize that the node has become available.
-                self.cancel = CancellationToken::new();
-            }
-            AvailabilityTransition::ToOffline => {
-                // Fire the node's cancellation token to cancel any in-flight API requests to it
-                self.cancel.cancel();
-            }
-            AvailabilityTransition::Unchanged => {}
-        }
-        self.availability = availability;
-    }
-
-    /// Without modifying the availability of the node, convert the intended availability
-    /// into a description of the transition.
-    pub(crate) fn get_availability_transition(
-        &self,
-        availability: NodeAvailability,
-    ) -> AvailabilityTransition {
-        use AvailabilityTransition::*;
-        use NodeAvailability::*;
-
-        match (self.availability, availability) {
-            (Offline, Active(_)) => ToActive,
-            (Active(_), Offline) => ToOffline,
-            _ => Unchanged,
-        }
-    }
-
-    /// Whether we may send API requests to this node.
-    pub(crate) fn is_available(&self) -> bool {
-        // When we clone a node, [`Self::availability`] is a snapshot, but [`Self::cancel`] holds
-        // a reference to the original Node's cancellation status.  Checking both of these results
-        // in a "pessimistic" check where we will consider a Node instance unavailable if it was unavailable
-        // when we cloned it, or if the original Node instance's cancellation token was fired.
-        matches!(self.availability, NodeAvailability::Active(_)) && !self.cancel.is_cancelled()
-    }
-
-    /// Is this node elegible to have work scheduled onto it?
-    pub(crate) fn may_schedule(&self) -> MaySchedule {
-        let score = match self.availability {
-            NodeAvailability::Active(score) => score,
-            NodeAvailability::Offline => return MaySchedule::No,
-        };
-
-        match self.scheduling {
-            NodeSchedulingPolicy::Active => MaySchedule::Yes(score),
-            NodeSchedulingPolicy::Draining => MaySchedule::No,
-            NodeSchedulingPolicy::Filling => MaySchedule::Yes(score),
-            NodeSchedulingPolicy::Pause => MaySchedule::No,
-        }
-    }
-
-    pub(crate) fn new(
-        id: NodeId,
-        listen_http_addr: String,
-        listen_http_port: u16,
-        listen_pg_addr: String,
-        listen_pg_port: u16,
-    ) -> Self {
-        Self {
-            id,
-            listen_http_addr,
-            listen_http_port,
-            listen_pg_addr,
-            listen_pg_port,
-            scheduling: NodeSchedulingPolicy::Filling,
-            availability: NodeAvailability::Offline,
-            cancel: CancellationToken::new(),
-        }
-    }
-
-    pub(crate) fn to_persistent(&self) -> NodePersistence {
-        NodePersistence {
-            node_id: self.id.0 as i64,
-            scheduling_policy: self.scheduling.into(),
-            listen_http_addr: self.listen_http_addr.clone(),
-            listen_http_port: self.listen_http_port as i32,
-            listen_pg_addr: self.listen_pg_addr.clone(),
-            listen_pg_port: self.listen_pg_port as i32,
-        }
-    }
-
-    pub(crate) fn from_persistent(np: NodePersistence) -> Self {
-        Self {
-            id: NodeId(np.node_id as u64),
-            // At startup we consider a node offline until proven otherwise.
-            availability: NodeAvailability::Offline,
-            scheduling: NodeSchedulingPolicy::from_str(&np.scheduling_policy)
-                .expect("Bad scheduling policy in DB"),
-            listen_http_addr: np.listen_http_addr,
-            listen_http_port: np.listen_http_port as u16,
-            listen_pg_addr: np.listen_pg_addr,
-            listen_pg_port: np.listen_pg_port as u16,
-            cancel: CancellationToken::new(),
-        }
-    }
-
-    /// Wrapper for issuing requests to pageserver management API: takes care of generic
-    /// retry/backoff for retryable HTTP status codes.
-    ///
-    /// This will return None to indicate cancellation.  Cancellation may happen from
-    /// the cancellation token passed in, or from Self's cancellation token (i.e. node
-    /// going offline).
-    pub(crate) async fn with_client_retries<T, O, F>(
-        &self,
-        mut op: O,
-        jwt: &Option<String>,
-        warn_threshold: u32,
-        max_retries: u32,
-        timeout: Duration,
-        cancel: &CancellationToken,
-    ) -> Option<mgmt_api::Result<T>>
-    where
-        O: FnMut(PageserverClient) -> F,
-        F: std::future::Future<Output = mgmt_api::Result<T>>,
-    {
-        fn is_fatal(e: &mgmt_api::Error) -> bool {
-            use mgmt_api::Error::*;
-            match e {
-                ReceiveBody(_) | ReceiveErrorBody(_) => false,
-                ApiError(StatusCode::SERVICE_UNAVAILABLE, _)
-                | ApiError(StatusCode::GATEWAY_TIMEOUT, _)
-                | ApiError(StatusCode::REQUEST_TIMEOUT, _) => false,
-                ApiError(_, _) => true,
-                Cancelled => true,
-            }
-        }
-
-        backoff::retry(
-            || {
-                let http_client = reqwest::ClientBuilder::new()
-                    .timeout(timeout)
-                    .build()
-                    .expect("Failed to construct HTTP client");
-
-                let client = PageserverClient::from_client(
-                    self.get_id(),
-                    http_client,
-                    self.base_url(),
-                    jwt.as_deref(),
-                );
-
-                let node_cancel_fut = self.cancel.cancelled();
-
-                let op_fut = op(client);
-
-                async {
-                    tokio::select! {
-                        r = op_fut=> {r},
-                        _ = node_cancel_fut => {
-                        Err(mgmt_api::Error::Cancelled)
-                    }}
-                }
-            },
-            is_fatal,
-            warn_threshold,
-            max_retries,
-            &format!(
-                "Call to node {} ({}:{}) management API",
-                self.id, self.listen_http_addr, self.listen_http_port
-            ),
-            cancel,
-        )
-        .await
-    }
-
-    /// Generate the simplified API-friendly description of a node's state
-    pub(crate) fn describe(&self) -> NodeDescribeResponse {
-        NodeDescribeResponse {
-            id: self.id,
-            availability: self.availability.into(),
-            scheduling: self.scheduling,
-            listen_http_addr: self.listen_http_addr.clone(),
-            listen_http_port: self.listen_http_port,
-            listen_pg_addr: self.listen_pg_addr.clone(),
-            listen_pg_port: self.listen_pg_port,
-        }
-    }
-}
-
-impl std::fmt::Display for Node {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{} ({})", self.id, self.listen_http_addr)
-    }
-}
-
-impl std::fmt::Debug for Node {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{} ({})", self.id, self.listen_http_addr)
-    }
-}

control_plane/attachment_service/src/pageserver_client.rs

@@ -1,203 +0,0 @@
-use pageserver_api::{
-    models::{
-        LocationConfig, LocationConfigListResponse, PageserverUtilization, SecondaryProgress,
-        TenantShardSplitRequest, TenantShardSplitResponse, TimelineCreateRequest, TimelineInfo,
-    },
-    shard::TenantShardId,
-};
-use pageserver_client::mgmt_api::{Client, Result};
-use reqwest::StatusCode;
-use utils::id::{NodeId, TimelineId};
-
-/// Thin wrapper around [`pageserver_client::mgmt_api::Client`]. It allows the storage
-/// controller to collect metrics in a non-intrusive manner.
-#[derive(Debug, Clone)]
-pub(crate) struct PageserverClient {
-    inner: Client,
-    node_id_label: String,
-}
-
-macro_rules! measured_request {
-    ($name:literal, $method:expr, $node_id: expr, $invoke:expr) => {{
-        let labels = crate::metrics::PageserverRequestLabelGroup {
-            pageserver_id: $node_id,
-            path: $name,
-            method: $method,
-        };
-
-        let latency = &crate::metrics::METRICS_REGISTRY
-            .metrics_group
-            .storage_controller_pageserver_request_latency;
-        let _timer_guard = latency.start_timer(labels.clone());
-
-        let res = $invoke;
-
-        if res.is_err() {
-            let error_counters = &crate::metrics::METRICS_REGISTRY
-                .metrics_group
-                .storage_controller_pageserver_request_error;
-            error_counters.inc(labels)
-        }
-
-        res
-    }};
-}
-
-impl PageserverClient {
-    pub(crate) fn new(node_id: NodeId, mgmt_api_endpoint: String, jwt: Option<&str>) -> Self {
-        Self {
-            inner: Client::from_client(reqwest::Client::new(), mgmt_api_endpoint, jwt),
-            node_id_label: node_id.0.to_string(),
-        }
-    }
-
-    pub(crate) fn from_client(
-        node_id: NodeId,
-        raw_client: reqwest::Client,
-        mgmt_api_endpoint: String,
-        jwt: Option<&str>,
-    ) -> Self {
-        Self {
-            inner: Client::from_client(raw_client, mgmt_api_endpoint, jwt),
-            node_id_label: node_id.0.to_string(),
-        }
-    }
-
-    pub(crate) async fn tenant_delete(&self, tenant_shard_id: TenantShardId) -> Result<StatusCode> {
-        measured_request!(
-            "tenant",
-            crate::metrics::Method::Delete,
-            &self.node_id_label,
-            self.inner.tenant_delete(tenant_shard_id).await
-        )
-    }
-
-    pub(crate) async fn tenant_time_travel_remote_storage(
-        &self,
-        tenant_shard_id: TenantShardId,
-        timestamp: &str,
-        done_if_after: &str,
-    ) -> Result<()> {
-        measured_request!(
-            "tenant_time_travel_remote_storage",
-            crate::metrics::Method::Put,
-            &self.node_id_label,
-            self.inner
-                .tenant_time_travel_remote_storage(tenant_shard_id, timestamp, done_if_after)
-                .await
-        )
-    }
-
-    pub(crate) async fn tenant_secondary_download(
-        &self,
-        tenant_id: TenantShardId,
-        wait: Option<std::time::Duration>,
-    ) -> Result<(StatusCode, SecondaryProgress)> {
-        measured_request!(
-            "tenant_secondary_download",
-            crate::metrics::Method::Post,
-            &self.node_id_label,
-            self.inner.tenant_secondary_download(tenant_id, wait).await
-        )
-    }
-
-    pub(crate) async fn location_config(
-        &self,
-        tenant_shard_id: TenantShardId,
-        config: LocationConfig,
-        flush_ms: Option<std::time::Duration>,
-        lazy: bool,
-    ) -> Result<()> {
-        measured_request!(
-            "location_config",
-            crate::metrics::Method::Put,
-            &self.node_id_label,
-            self.inner
-                .location_config(tenant_shard_id, config, flush_ms, lazy)
-                .await
-        )
-    }
-
-    pub(crate) async fn list_location_config(&self) -> Result<LocationConfigListResponse> {
-        measured_request!(
-            "location_configs",
-            crate::metrics::Method::Get,
-            &self.node_id_label,
-            self.inner.list_location_config().await
-        )
-    }
-
-    pub(crate) async fn get_location_config(
-        &self,
-        tenant_shard_id: TenantShardId,
-    ) -> Result<Option<LocationConfig>> {
-        measured_request!(
-            "location_config",
-            crate::metrics::Method::Get,
-            &self.node_id_label,
-            self.inner.get_location_config(tenant_shard_id).await
-        )
-    }
-
-    pub(crate) async fn timeline_create(
-        &self,
-        tenant_shard_id: TenantShardId,
-        req: &TimelineCreateRequest,
-    ) -> Result<TimelineInfo> {
-        measured_request!(
-            "timeline",
-            crate::metrics::Method::Post,
-            &self.node_id_label,
-            self.inner.timeline_create(tenant_shard_id, req).await
-        )
-    }
-
-    pub(crate) async fn timeline_delete(
-        &self,
-        tenant_shard_id: TenantShardId,
-        timeline_id: TimelineId,
-    ) -> Result<StatusCode> {
-        measured_request!(
-            "timeline",
-            crate::metrics::Method::Delete,
-            &self.node_id_label,
-            self.inner
-                .timeline_delete(tenant_shard_id, timeline_id)
-                .await
-        )
-    }
-
-    pub(crate) async fn tenant_shard_split(
-        &self,
-        tenant_shard_id: TenantShardId,
-        req: TenantShardSplitRequest,
-    ) -> Result<TenantShardSplitResponse> {
-        measured_request!(
-            "tenant_shard_split",
-            crate::metrics::Method::Put,
-            &self.node_id_label,
-            self.inner.tenant_shard_split(tenant_shard_id, req).await
-        )
-    }
-
-    pub(crate) async fn timeline_list(
-        &self,
-        tenant_shard_id: &TenantShardId,
-    ) -> Result<Vec<TimelineInfo>> {
-        measured_request!(
-            "timelines",
-            crate::metrics::Method::Get,
-            &self.node_id_label,
-            self.inner.timeline_list(tenant_shard_id).await
-        )
-    }
-
-    pub(crate) async fn get_utilization(&self) -> Result<PageserverUtilization> {
-        measured_request!(
-            "utilization",
-            crate::metrics::Method::Get,
-            &self.node_id_label,
-            self.inner.get_utilization().await
-        )
-    }
-}

control_plane/attachment_service/src/persistence.rs

@@ -1,765 +0,0 @@
-pub(crate) mod split_state;
-use std::collections::HashMap;
-use std::str::FromStr;
-use std::time::Duration;
-
-use self::split_state::SplitState;
-use camino::Utf8Path;
-use camino::Utf8PathBuf;
-use diesel::pg::PgConnection;
-use diesel::prelude::*;
-use diesel::Connection;
-use pageserver_api::controller_api::ShardSchedulingPolicy;
-use pageserver_api::controller_api::{NodeSchedulingPolicy, PlacementPolicy};
-use pageserver_api::models::TenantConfig;
-use pageserver_api::shard::ShardConfigError;
-use pageserver_api::shard::ShardIdentity;
-use pageserver_api::shard::ShardStripeSize;
-use pageserver_api::shard::{ShardCount, ShardNumber, TenantShardId};
-use serde::{Deserialize, Serialize};
-use utils::generation::Generation;
-use utils::id::{NodeId, TenantId};
-
-use crate::metrics::{
-    DatabaseQueryErrorLabelGroup, DatabaseQueryLatencyLabelGroup, METRICS_REGISTRY,
-};
-use crate::node::Node;
-
-/// ## What do we store?
-///
-/// The storage controller service does not store most of its state durably.
-///
-/// The essential things to store durably are:
-/// - generation numbers, as these must always advance monotonically to ensure data safety.
-/// - Tenant's PlacementPolicy and TenantConfig, as the source of truth for these is something external.
-/// - Node's scheduling policies, as the source of truth for these is something external.
-///
-/// Other things we store durably as an implementation detail:
-/// - Node's host/port: this could be avoided it we made nodes emit a self-registering heartbeat,
-///   but it is operationally simpler to make this service the authority for which nodes
-///   it talks to.
-///
-/// ## Performance/efficiency
-///
-/// The storage controller service does not go via the database for most things: there are
-/// a couple of places where we must, and where efficiency matters:
-/// - Incrementing generation numbers: the Reconciler has to wait for this to complete
-///   before it can attach a tenant, so this acts as a bound on how fast things like
-///   failover can happen.
-/// - Pageserver re-attach: we will increment many shards' generations when this happens,
-///   so it is important to avoid e.g. issuing O(N) queries.
-///
-/// Database calls relating to nodes have low performance requirements, as they are very rarely
-/// updated, and reads of nodes are always from memory, not the database.  We only require that
-/// we can UPDATE a node's scheduling mode reasonably quickly to mark a bad node offline.
-pub struct Persistence {
-    connection_pool: diesel::r2d2::Pool<diesel::r2d2::ConnectionManager<PgConnection>>,
-
-    // In test environments, we support loading+saving a JSON file.  This is temporary, for the benefit of
-    // test_compatibility.py, so that we don't have to commit to making the database contents fully backward/forward
-    // compatible just yet.
-    json_path: Option<Utf8PathBuf>,
-}
-
-/// Legacy format, for use in JSON compat objects in test environment
-#[derive(Serialize, Deserialize)]
-struct JsonPersistence {
-    tenants: HashMap<TenantShardId, TenantShardPersistence>,
-}
-
-#[derive(thiserror::Error, Debug)]
-pub(crate) enum DatabaseError {
-    #[error(transparent)]
-    Query(#[from] diesel::result::Error),
-    #[error(transparent)]
-    Connection(#[from] diesel::result::ConnectionError),
-    #[error(transparent)]
-    ConnectionPool(#[from] r2d2::Error),
-    #[error("Logical error: {0}")]
-    Logical(String),
-}
-
-#[derive(measured::FixedCardinalityLabel, Clone)]
-pub(crate) enum DatabaseOperation {
-    InsertNode,
-    UpdateNode,
-    DeleteNode,
-    ListNodes,
-    BeginShardSplit,
-    CompleteShardSplit,
-    AbortShardSplit,
-    Detach,
-    ReAttach,
-    IncrementGeneration,
-    ListTenantShards,
-    InsertTenantShards,
-    UpdateTenantShard,
-    DeleteTenant,
-    UpdateTenantConfig,
-}
-
-#[must_use]
-pub(crate) enum AbortShardSplitStatus {
-    /// We aborted the split in the database by reverting to the parent shards
-    Aborted,
-    /// The split had already been persisted.
-    Complete,
-}
-
-pub(crate) type DatabaseResult<T> = Result<T, DatabaseError>;
-
-/// Some methods can operate on either a whole tenant or a single shard
-pub(crate) enum TenantFilter {
-    Tenant(TenantId),
-    Shard(TenantShardId),
-}
-
-impl Persistence {
-    // The default postgres connection limit is 100.  We use up to 99, to leave one free for a human admin under
-    // normal circumstances.  This assumes we have exclusive use of the database cluster to which we connect.
-    pub const MAX_CONNECTIONS: u32 = 99;
-
-    // We don't want to keep a lot of connections alive: close them down promptly if they aren't being used.
-    const IDLE_CONNECTION_TIMEOUT: Duration = Duration::from_secs(10);
-    const MAX_CONNECTION_LIFETIME: Duration = Duration::from_secs(60);
-
-    pub fn new(database_url: String, json_path: Option<Utf8PathBuf>) -> Self {
-        let manager = diesel::r2d2::ConnectionManager::<PgConnection>::new(database_url);
-
-        // We will use a connection pool: this is primarily to _limit_ our connection count, rather than to optimize time
-        // to execute queries (database queries are not generally on latency-sensitive paths).
-        let connection_pool = diesel::r2d2::Pool::builder()
-            .max_size(Self::MAX_CONNECTIONS)
-            .max_lifetime(Some(Self::MAX_CONNECTION_LIFETIME))
-            .idle_timeout(Some(Self::IDLE_CONNECTION_TIMEOUT))
-            // Always keep at least one connection ready to go
-            .min_idle(Some(1))
-            .test_on_check_out(true)
-            .build(manager)
-            .expect("Could not build connection pool");
-
-        Self {
-            connection_pool,
-            json_path,
-        }
-    }
-
-    /// Wraps `with_conn` in order to collect latency and error metrics
-    async fn with_measured_conn<F, R>(&self, op: DatabaseOperation, func: F) -> DatabaseResult<R>
-    where
-        F: FnOnce(&mut PgConnection) -> DatabaseResult<R> + Send + 'static,
-        R: Send + 'static,
-    {
-        let latency = &METRICS_REGISTRY
-            .metrics_group
-            .storage_controller_database_query_latency;
-        let _timer = latency.start_timer(DatabaseQueryLatencyLabelGroup {
-            operation: op.clone(),
-        });
-
-        let res = self.with_conn(func).await;
-
-        if let Err(err) = &res {
-            let error_counter = &METRICS_REGISTRY
-                .metrics_group
-                .storage_controller_database_query_error;
-            error_counter.inc(DatabaseQueryErrorLabelGroup {
-                error_type: err.error_label(),
-                operation: op,
-            })
-        }
-
-        res
-    }
-
-    /// Call the provided function in a tokio blocking thread, with a Diesel database connection.
-    async fn with_conn<F, R>(&self, func: F) -> DatabaseResult<R>
-    where
-        F: FnOnce(&mut PgConnection) -> DatabaseResult<R> + Send + 'static,
-        R: Send + 'static,
-    {
-        let mut conn = self.connection_pool.get()?;
-        tokio::task::spawn_blocking(move || -> DatabaseResult<R> { func(&mut conn) })
-            .await
-            .expect("Task panic")
-    }
-
-    /// When a node is first registered, persist it before using it for anything
-    pub(crate) async fn insert_node(&self, node: &Node) -> DatabaseResult<()> {
-        let np = node.to_persistent();
-        self.with_measured_conn(
-            DatabaseOperation::InsertNode,
-            move |conn| -> DatabaseResult<()> {
-                diesel::insert_into(crate::schema::nodes::table)
-                    .values(&np)
-                    .execute(conn)?;
-                Ok(())
-            },
-        )
-        .await
-    }
-
-    /// At startup, populate the list of nodes which our shards may be placed on
-    pub(crate) async fn list_nodes(&self) -> DatabaseResult<Vec<NodePersistence>> {
-        let nodes: Vec<NodePersistence> = self
-            .with_measured_conn(
-                DatabaseOperation::ListNodes,
-                move |conn| -> DatabaseResult<_> {
-                    Ok(crate::schema::nodes::table.load::<NodePersistence>(conn)?)
-                },
-            )
-            .await?;
-
-        tracing::info!("list_nodes: loaded {} nodes", nodes.len());
-
-        Ok(nodes)
-    }
-
-    pub(crate) async fn update_node(
-        &self,
-        input_node_id: NodeId,
-        input_scheduling: NodeSchedulingPolicy,
-    ) -> DatabaseResult<()> {
-        use crate::schema::nodes::dsl::*;
-        let updated = self
-            .with_measured_conn(DatabaseOperation::UpdateNode, move |conn| {
-                let updated = diesel::update(nodes)
-                    .filter(node_id.eq(input_node_id.0 as i64))
-                    .set((scheduling_policy.eq(String::from(input_scheduling)),))
-                    .execute(conn)?;
-                Ok(updated)
-            })
-            .await?;
-
-        if updated != 1 {
-            Err(DatabaseError::Logical(format!(
-                "Node {node_id:?} not found for update",
-            )))
-        } else {
-            Ok(())
-        }
-    }
-
-    /// At startup, load the high level state for shards, such as their config + policy.  This will
-    /// be enriched at runtime with state discovered on pageservers.
-    pub(crate) async fn list_tenant_shards(&self) -> DatabaseResult<Vec<TenantShardPersistence>> {
-        let loaded = self
-            .with_measured_conn(
-                DatabaseOperation::ListTenantShards,
-                move |conn| -> DatabaseResult<_> {
-                    Ok(crate::schema::tenant_shards::table.load::<TenantShardPersistence>(conn)?)
-                },
-            )
-            .await?;
-
-        if loaded.is_empty() {
-            if let Some(path) = &self.json_path {
-                if tokio::fs::try_exists(path)
-                    .await
-                    .map_err(|e| DatabaseError::Logical(format!("Error stat'ing JSON file: {e}")))?
-                {
-                    tracing::info!("Importing from legacy JSON format at {path}");
-                    return self.list_tenant_shards_json(path).await;
-                }
-            }
-        }
-        Ok(loaded)
-    }
-
-    /// Shim for automated compatibility tests: load tenants from a JSON file instead of database
-    pub(crate) async fn list_tenant_shards_json(
-        &self,
-        path: &Utf8Path,
-    ) -> DatabaseResult<Vec<TenantShardPersistence>> {
-        let bytes = tokio::fs::read(path)
-            .await
-            .map_err(|e| DatabaseError::Logical(format!("Failed to load JSON: {e}")))?;
-
-        let mut decoded = serde_json::from_slice::<JsonPersistence>(&bytes)
-            .map_err(|e| DatabaseError::Logical(format!("Deserialization error: {e}")))?;
-        for shard in decoded.tenants.values_mut() {
-            if shard.placement_policy == "\"Single\"" {
-                // Backward compat for test data after PR https://github.com/neondatabase/neon/pull/7165
-                shard.placement_policy = "{\"Attached\":0}".to_string();
-            }
-
-            if shard.scheduling_policy.is_empty() {
-                shard.scheduling_policy =
-                    serde_json::to_string(&ShardSchedulingPolicy::default()).unwrap();
-            }
-        }
-
-        let tenants: Vec<TenantShardPersistence> = decoded.tenants.into_values().collect();
-
-        // Synchronize database with what is in the JSON file
-        self.insert_tenant_shards(tenants.clone()).await?;
-
-        Ok(tenants)
-    }
-
-    /// For use in testing environments, where we dump out JSON on shutdown.
-    pub async fn write_tenants_json(&self) -> anyhow::Result<()> {
-        let Some(path) = &self.json_path else {
-            anyhow::bail!("Cannot write JSON if path isn't set (test environment bug)");
-        };
-        tracing::info!("Writing state to {path}...");
-        let tenants = self.list_tenant_shards().await?;
-        let mut tenants_map = HashMap::new();
-        for tsp in tenants {
-            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::new(tsp.shard_count as u8),
-            };
-
-            tenants_map.insert(tenant_shard_id, tsp);
-        }
-        let json = serde_json::to_string(&JsonPersistence {
-            tenants: tenants_map,
-        })?;
-
-        tokio::fs::write(path, &json).await?;
-        tracing::info!("Wrote {} bytes to {path}...", json.len());
-
-        Ok(())
-    }
-
-    /// 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<TenantShardPersistence>,
-    ) -> DatabaseResult<()> {
-        use crate::schema::tenant_shards::dsl::*;
-        self.with_measured_conn(
-            DatabaseOperation::InsertTenantShards,
-            move |conn| -> DatabaseResult<()> {
-                conn.transaction(|conn| -> QueryResult<()> {
-                    for tenant in &shards {
-                        diesel::insert_into(tenant_shards)
-                            .values(tenant)
-                            .execute(conn)?;
-                    }
-                    Ok(())
-                })?;
-                Ok(())
-            },
-        )
-        .await
-    }
-
-    /// Ordering: call this _after_ deleting the tenant on pageservers, but _before_ dropping state for
-    /// the tenant from memory on this server.
-    pub(crate) async fn delete_tenant(&self, del_tenant_id: TenantId) -> DatabaseResult<()> {
-        use crate::schema::tenant_shards::dsl::*;
-        self.with_measured_conn(
-            DatabaseOperation::DeleteTenant,
-            move |conn| -> DatabaseResult<()> {
-                diesel::delete(tenant_shards)
-                    .filter(tenant_id.eq(del_tenant_id.to_string()))
-                    .execute(conn)?;
-
-                Ok(())
-            },
-        )
-        .await
-    }
-
-    pub(crate) async fn delete_node(&self, del_node_id: NodeId) -> DatabaseResult<()> {
-        use crate::schema::nodes::dsl::*;
-        self.with_measured_conn(
-            DatabaseOperation::DeleteNode,
-            move |conn| -> DatabaseResult<()> {
-                diesel::delete(nodes)
-                    .filter(node_id.eq(del_node_id.0 as i64))
-                    .execute(conn)?;
-
-                Ok(())
-            },
-        )
-        .await
-    }
-
-    /// When a tenant invokes the /re-attach API, this function is responsible for doing an efficient
-    /// batched increment of the generations of all tenants whose generation_pageserver is equal to
-    /// the node that called /re-attach.
-    #[tracing::instrument(skip_all, fields(node_id))]
-    pub(crate) async fn re_attach(
-        &self,
-        node_id: NodeId,
-    ) -> DatabaseResult<HashMap<TenantShardId, Generation>> {
-        use crate::schema::tenant_shards::dsl::*;
-        let updated = self
-            .with_measured_conn(DatabaseOperation::ReAttach, move |conn| {
-                let rows_updated = diesel::update(tenant_shards)
-                    .filter(generation_pageserver.eq(node_id.0 as i64))
-                    .set(generation.eq(generation + 1))
-                    .execute(conn)?;
-
-                tracing::info!("Incremented {} tenants' generations", rows_updated);
-
-                // TODO: UPDATE+SELECT in one query
-
-                let updated = tenant_shards
-                    .filter(generation_pageserver.eq(node_id.0 as i64))
-                    .select(TenantShardPersistence::as_select())
-                    .load(conn)?;
-                Ok(updated)
-            })
-            .await?;
-
-        let mut result = HashMap::new();
-        for tsp in updated {
-            let tenant_shard_id = TenantShardId {
-                tenant_id: TenantId::from_str(tsp.tenant_id.as_str())
-                    .map_err(|e| DatabaseError::Logical(format!("Malformed tenant id: {e}")))?,
-                shard_number: ShardNumber(tsp.shard_number as u8),
-                shard_count: ShardCount::new(tsp.shard_count as u8),
-            };
-
-            let Some(g) = tsp.generation else {
-                // If the generation_pageserver column was non-NULL, then the generation column should also be non-NULL:
-                // we only set generation_pageserver when setting generation.
-                return Err(DatabaseError::Logical(
-                    "Generation should always be set after incrementing".to_string(),
-                ));
-            };
-            result.insert(tenant_shard_id, Generation::new(g as u32));
-        }
-
-        Ok(result)
-    }
-
-    /// 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: NodeId,
-    ) -> anyhow::Result<Generation> {
-        use crate::schema::tenant_shards::dsl::*;
-        let updated = self
-            .with_measured_conn(DatabaseOperation::IncrementGeneration, move |conn| {
-                let updated = diesel::update(tenant_shards)
-                    .filter(tenant_id.eq(tenant_shard_id.tenant_id.to_string()))
-                    .filter(shard_number.eq(tenant_shard_id.shard_number.0 as i32))
-                    .filter(shard_count.eq(tenant_shard_id.shard_count.literal() as i32))
-                    .set((
-                        generation.eq(generation + 1),
-                        generation_pageserver.eq(node_id.0 as i64),
-                    ))
-                    // TODO: only returning() the generation column
-                    .returning(TenantShardPersistence::as_returning())
-                    .get_result(conn)?;
-
-                Ok(updated)
-            })
-            .await?;
-
-        // Generation is always non-null in the rseult: if the generation column had been NULL, then we
-        // should have experienced an SQL Confilict error while executing a query that tries to increment it.
-        debug_assert!(updated.generation.is_some());
-        let Some(g) = updated.generation else {
-            return Err(DatabaseError::Logical(
-                "Generation should always be set after incrementing".to_string(),
-            )
-            .into());
-        };
-
-        Ok(Generation::new(g as u32))
-    }
-
-    /// For use when updating a persistent property of a tenant, such as its config or placement_policy.
-    ///
-    /// Do not use this for settting generation, unless in the special onboarding code path (/location_config)
-    /// API: use [`Self::increment_generation`] instead.  Setting the generation via this route is a one-time thing
-    /// that we only do the first time a tenant is set to an attached policy via /location_config.
-    pub(crate) async fn update_tenant_shard(
-        &self,
-        tenant: TenantFilter,
-        input_placement_policy: Option<PlacementPolicy>,
-        input_config: Option<TenantConfig>,
-        input_generation: Option<Generation>,
-        input_scheduling_policy: Option<ShardSchedulingPolicy>,
-    ) -> DatabaseResult<()> {
-        use crate::schema::tenant_shards::dsl::*;
-
-        self.with_measured_conn(DatabaseOperation::UpdateTenantShard, move |conn| {
-            let query = match tenant {
-                TenantFilter::Shard(tenant_shard_id) => diesel::update(tenant_shards)
-                    .filter(tenant_id.eq(tenant_shard_id.tenant_id.to_string()))
-                    .filter(shard_number.eq(tenant_shard_id.shard_number.0 as i32))
-                    .filter(shard_count.eq(tenant_shard_id.shard_count.literal() as i32))
-                    .into_boxed(),
-                TenantFilter::Tenant(input_tenant_id) => diesel::update(tenant_shards)
-                    .filter(tenant_id.eq(input_tenant_id.to_string()))
-                    .into_boxed(),
-            };
-
-            #[derive(AsChangeset)]
-            #[diesel(table_name = crate::schema::tenant_shards)]
-            struct ShardUpdate {
-                generation: Option<i32>,
-                placement_policy: Option<String>,
-                config: Option<String>,
-                scheduling_policy: Option<String>,
-            }
-
-            let update = ShardUpdate {
-                generation: input_generation.map(|g| g.into().unwrap() as i32),
-                placement_policy: input_placement_policy
-                    .map(|p| serde_json::to_string(&p).unwrap()),
-                config: input_config.map(|c| serde_json::to_string(&c).unwrap()),
-                scheduling_policy: input_scheduling_policy
-                    .map(|p| serde_json::to_string(&p).unwrap()),
-            };
-
-            query.set(update).execute(conn)?;
-
-            Ok(())
-        })
-        .await?;
-
-        Ok(())
-    }
-
-    pub(crate) async fn detach(&self, tenant_shard_id: TenantShardId) -> anyhow::Result<()> {
-        use crate::schema::tenant_shards::dsl::*;
-        self.with_measured_conn(DatabaseOperation::Detach, move |conn| {
-            let updated = diesel::update(tenant_shards)
-                .filter(tenant_id.eq(tenant_shard_id.tenant_id.to_string()))
-                .filter(shard_number.eq(tenant_shard_id.shard_number.0 as i32))
-                .filter(shard_count.eq(tenant_shard_id.shard_count.literal() as i32))
-                .set((
-                    generation_pageserver.eq(Option::<i64>::None),
-                    placement_policy.eq(serde_json::to_string(&PlacementPolicy::Detached).unwrap()),
-                ))
-                .execute(conn)?;
-
-            Ok(updated)
-        })
-        .await?;
-
-        Ok(())
-    }
-
-    // When we start shard splitting, we must durably mark the tenant so that
-    // on restart, we know that we must go through recovery.
-    //
-    // We create the child shards here, so that they will be available for increment_generation calls
-    // if some pageserver holding a child shard needs to restart before the overall tenant split is complete.
-    pub(crate) async fn begin_shard_split(
-        &self,
-        old_shard_count: ShardCount,
-        split_tenant_id: TenantId,
-        parent_to_children: Vec<(TenantShardId, Vec<TenantShardPersistence>)>,
-    ) -> DatabaseResult<()> {
-        use crate::schema::tenant_shards::dsl::*;
-        self.with_measured_conn(DatabaseOperation::BeginShardSplit, move |conn| -> DatabaseResult<()> {
-            conn.transaction(|conn| -> DatabaseResult<()> {
-                // Mark parent shards as splitting
-
-                let updated = diesel::update(tenant_shards)
-                    .filter(tenant_id.eq(split_tenant_id.to_string()))
-                    .filter(shard_count.eq(old_shard_count.literal() as i32))
-                    .set((splitting.eq(1),))
-                    .execute(conn)?;
-                if u8::try_from(updated)
-                    .map_err(|_| DatabaseError::Logical(
-                        format!("Overflow existing shard count {} while splitting", updated))
-                    )? != old_shard_count.count() {
-                    // Perhaps a deletion or another split raced with this attempt to split, mutating
-                    // the parent shards that we intend to split. In this case the split request should fail.
-                    return Err(DatabaseError::Logical(
-                        format!("Unexpected existing shard count {updated} when preparing tenant for split (expected {})", old_shard_count.count())
-                    ));
-                }
-
-                // FIXME: spurious clone to sidestep closure move rules
-                let parent_to_children = parent_to_children.clone();
-
-                // Insert child shards
-                for (parent_shard_id, children) in parent_to_children {
-                    let mut parent = crate::schema::tenant_shards::table
-                        .filter(tenant_id.eq(parent_shard_id.tenant_id.to_string()))
-                        .filter(shard_number.eq(parent_shard_id.shard_number.0 as i32))
-                        .filter(shard_count.eq(parent_shard_id.shard_count.literal() as i32))
-                        .load::<TenantShardPersistence>(conn)?;
-                    let parent = if parent.len() != 1 {
-                        return Err(DatabaseError::Logical(format!(
-                            "Parent shard {parent_shard_id} not found"
-                        )));
-                    } else {
-                        parent.pop().unwrap()
-                    };
-                    for mut shard in children {
-                        // Carry the parent's generation into the child
-                        shard.generation = parent.generation;
-
-                        debug_assert!(shard.splitting == SplitState::Splitting);
-                        diesel::insert_into(tenant_shards)
-                            .values(shard)
-                            .execute(conn)?;
-                    }
-                }
-
-                Ok(())
-            })?;
-
-            Ok(())
-        })
-        .await
-    }
-
-    // When we finish shard splitting, we must atomically clean up the old shards
-    // and insert the new shards, and clear the splitting marker.
-    pub(crate) async fn complete_shard_split(
-        &self,
-        split_tenant_id: TenantId,
-        old_shard_count: ShardCount,
-    ) -> DatabaseResult<()> {
-        use crate::schema::tenant_shards::dsl::*;
-        self.with_measured_conn(
-            DatabaseOperation::CompleteShardSplit,
-            move |conn| -> DatabaseResult<()> {
-                conn.transaction(|conn| -> QueryResult<()> {
-                    // Drop parent shards
-                    diesel::delete(tenant_shards)
-                        .filter(tenant_id.eq(split_tenant_id.to_string()))
-                        .filter(shard_count.eq(old_shard_count.literal() as i32))
-                        .execute(conn)?;
-
-                    // Clear sharding flag
-                    let updated = diesel::update(tenant_shards)
-                        .filter(tenant_id.eq(split_tenant_id.to_string()))
-                        .set((splitting.eq(0),))
-                        .execute(conn)?;
-                    debug_assert!(updated > 0);
-
-                    Ok(())
-                })?;
-
-                Ok(())
-            },
-        )
-        .await
-    }
-
-    /// Used when the remote part of a shard split failed: we will revert the database state to have only
-    /// the parent shards, with SplitState::Idle.
-    pub(crate) async fn abort_shard_split(
-        &self,
-        split_tenant_id: TenantId,
-        new_shard_count: ShardCount,
-    ) -> DatabaseResult<AbortShardSplitStatus> {
-        use crate::schema::tenant_shards::dsl::*;
-        self.with_measured_conn(
-            DatabaseOperation::AbortShardSplit,
-            move |conn| -> DatabaseResult<AbortShardSplitStatus> {
-                let aborted =
-                    conn.transaction(|conn| -> DatabaseResult<AbortShardSplitStatus> {
-                        // Clear the splitting state on parent shards
-                        let updated = diesel::update(tenant_shards)
-                            .filter(tenant_id.eq(split_tenant_id.to_string()))
-                            .filter(shard_count.ne(new_shard_count.literal() as i32))
-                            .set((splitting.eq(0),))
-                            .execute(conn)?;
-
-                        // Parent shards are already gone: we cannot abort.
-                        if updated == 0 {
-                            return Ok(AbortShardSplitStatus::Complete);
-                        }
-
-                        // Sanity check: if parent shards were present, their cardinality should
-                        // be less than the number of child shards.
-                        if updated >= new_shard_count.count() as usize {
-                            return Err(DatabaseError::Logical(format!(
-                                "Unexpected parent shard count {updated} while aborting split to \
-                            count {new_shard_count:?} on tenant {split_tenant_id}"
-                            )));
-                        }
-
-                        // Erase child shards
-                        diesel::delete(tenant_shards)
-                            .filter(tenant_id.eq(split_tenant_id.to_string()))
-                            .filter(shard_count.eq(new_shard_count.literal() as i32))
-                            .execute(conn)?;
-
-                        Ok(AbortShardSplitStatus::Aborted)
-                    })?;
-
-                Ok(aborted)
-            },
-        )
-        .await
-    }
-}
-
-/// Parts of [`crate::tenant_state::TenantState`] that are stored durably
-#[derive(Queryable, Selectable, Insertable, Serialize, Deserialize, Clone, Eq, PartialEq)]
-#[diesel(table_name = crate::schema::tenant_shards)]
-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,
-
-    // Latest generation number: next time we attach, increment this
-    // and use the incremented number when attaching.
-    //
-    // Generation is only None when first onboarding a tenant, where it may
-    // be in PlacementPolicy::Secondary and therefore have no valid generation state.
-    pub(crate) generation: Option<i32>,
-
-    // Currently attached pageserver
-    #[serde(rename = "pageserver")]
-    pub(crate) generation_pageserver: Option<i64>,
-
-    #[serde(default)]
-    pub(crate) placement_policy: String,
-    #[serde(default)]
-    pub(crate) splitting: SplitState,
-    #[serde(default)]
-    pub(crate) config: String,
-    #[serde(default)]
-    pub(crate) scheduling_policy: String,
-}
-
-impl TenantShardPersistence {
-    pub(crate) fn get_shard_identity(&self) -> Result<ShardIdentity, ShardConfigError> {
-        if self.shard_count == 0 {
-            Ok(ShardIdentity::unsharded())
-        } else {
-            Ok(ShardIdentity::new(
-                ShardNumber(self.shard_number as u8),
-                ShardCount::new(self.shard_count as u8),
-                ShardStripeSize(self.shard_stripe_size as u32),
-            )?)
-        }
-    }
-
-    pub(crate) fn get_tenant_shard_id(&self) -> Result<TenantShardId, hex::FromHexError> {
-        Ok(TenantShardId {
-            tenant_id: TenantId::from_str(self.tenant_id.as_str())?,
-            shard_number: ShardNumber(self.shard_number as u8),
-            shard_count: ShardCount::new(self.shard_count as u8),
-        })
-    }
-}
-
-/// Parts of [`crate::node::Node`] that are stored durably
-#[derive(Serialize, Deserialize, Queryable, Selectable, Insertable, Eq, PartialEq)]
-#[diesel(table_name = crate::schema::nodes)]
-pub(crate) struct NodePersistence {
-    pub(crate) node_id: i64,
-    pub(crate) scheduling_policy: String,
-    pub(crate) listen_http_addr: String,
-    pub(crate) listen_http_port: i32,
-    pub(crate) listen_pg_addr: String,
-    pub(crate) listen_pg_port: i32,
-}

control_plane/attachment_service/src/persistence/split_state.rs

@@ -1,46 +0,0 @@
-use diesel::pg::{Pg, PgValue};
-use diesel::{
-    deserialize::FromSql, deserialize::FromSqlRow, expression::AsExpression, serialize::ToSql,
-    sql_types::Int2,
-};
-use serde::{Deserialize, Serialize};
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, FromSqlRow, AsExpression)]
-#[diesel(sql_type = SplitStateSQLRepr)]
-#[derive(Deserialize, Serialize)]
-pub enum SplitState {
-    Idle = 0,
-    Splitting = 1,
-}
-
-impl Default for SplitState {
-    fn default() -> Self {
-        Self::Idle
-    }
-}
-
-type SplitStateSQLRepr = Int2;
-
-impl ToSql<SplitStateSQLRepr, Pg> for SplitState {
-    fn to_sql<'a>(
-        &'a self,
-        out: &'a mut diesel::serialize::Output<Pg>,
-    ) -> diesel::serialize::Result {
-        let raw_value: i16 = *self as i16;
-        let mut new_out = out.reborrow();
-        ToSql::<SplitStateSQLRepr, Pg>::to_sql(&raw_value, &mut new_out)
-    }
-}
-
-impl FromSql<SplitStateSQLRepr, Pg> for SplitState {
-    fn from_sql(pg_value: PgValue) -> diesel::deserialize::Result<Self> {
-        match FromSql::<SplitStateSQLRepr, Pg>::from_sql(pg_value).map(|v| match v {
-            0 => Some(Self::Idle),
-            1 => Some(Self::Splitting),
-            _ => None,
-        })? {
-            Some(v) => Ok(v),
-            None => Err(format!("Invalid SplitState value, was: {:?}", pg_value.as_bytes()).into()),
-        }
-    }
-}

control_plane/attachment_service/src/reconciler.rs

@@ -1,810 +0,0 @@
-use crate::pageserver_client::PageserverClient;
-use crate::persistence::Persistence;
-use crate::service;
-use hyper::StatusCode;
-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, Instant};
-use tokio_util::sync::CancellationToken;
-use utils::generation::Generation;
-use utils::id::{NodeId, TimelineId};
-use utils::lsn::Lsn;
-use utils::sync::gate::GateGuard;
-
-use crate::compute_hook::{ComputeHook, NotifyError};
-use crate::node::Node;
-use crate::tenant_state::{IntentState, ObservedState, ObservedStateLocation};
-
-const DEFAULT_HEATMAP_PERIOD: &str = "60s";
-
-/// 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: Option<Generation>,
-    pub(crate) intent: TargetState,
-
-    /// Nodes not referenced by [`Self::intent`], from which we should try
-    /// to detach this tenant shard.
-    pub(crate) detach: Vec<Node>,
-
-    pub(crate) config: TenantConfig,
-    pub(crate) observed: ObservedState,
-
-    pub(crate) service_config: service::Config,
-
-    /// A hook to notify the running postgres instances when we change the location
-    /// of a tenant.  Use this via [`Self::compute_notify`] to update our failure flag
-    /// and guarantee eventual retries.
-    pub(crate) compute_hook: Arc<ComputeHook>,
-
-    /// To avoid stalling if the cloud control plane is unavailable, we may proceed
-    /// past failures in [`ComputeHook::notify`], but we _must_ remember that we failed
-    /// so that we can set [`crate::tenant_state::TenantState::pending_compute_notification`] to ensure a later retry.
-    pub(crate) compute_notify_failure: bool,
-
-    /// 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,
-
-    /// Reconcilers are registered with a Gate so that during a graceful shutdown we
-    /// can wait for all the reconcilers to respond to their cancellation tokens.
-    pub(crate) _gate_guard: GateGuard,
-
-    /// Access to persistent storage for updating generation numbers
-    pub(crate) persistence: Arc<Persistence>,
-}
-
-/// This is a snapshot of [`crate::tenant_state::IntentState`], but it does not do any
-/// reference counting for Scheduler.  The IntentState is what the scheduler works with,
-/// and the TargetState is just the instruction for a particular Reconciler run.
-#[derive(Debug)]
-pub(crate) struct TargetState {
-    pub(crate) attached: Option<Node>,
-    pub(crate) secondary: Vec<Node>,
-}
-
-impl TargetState {
-    pub(crate) fn from_intent(nodes: &HashMap<NodeId, Node>, intent: &IntentState) -> Self {
-        Self {
-            attached: intent.get_attached().map(|n| {
-                nodes
-                    .get(&n)
-                    .expect("Intent attached referenced non-existent node")
-                    .clone()
-            }),
-            secondary: intent
-                .get_secondary()
-                .iter()
-                .map(|n| {
-                    nodes
-                        .get(n)
-                        .expect("Intent secondary referenced non-existent node")
-                        .clone()
-                })
-                .collect(),
-        }
-    }
-}
-
-#[derive(thiserror::Error, Debug)]
-pub(crate) enum ReconcileError {
-    #[error(transparent)]
-    Remote(#[from] mgmt_api::Error),
-    #[error(transparent)]
-    Notify(#[from] NotifyError),
-    #[error("Cancelled")]
-    Cancel,
-    #[error(transparent)]
-    Other(#[from] anyhow::Error),
-}
-
-impl Reconciler {
-    async fn location_config(
-        &mut self,
-        node: &Node,
-        config: LocationConfig,
-        flush_ms: Option<Duration>,
-        lazy: bool,
-    ) -> Result<(), ReconcileError> {
-        if !node.is_available() && config.mode == LocationConfigMode::Detached {
-            // Attempts to detach from offline nodes may be imitated without doing I/O: a node which is offline
-            // will get fully reconciled wrt the shard's intent state when it is reactivated, irrespective of
-            // what we put into `observed`, in [`crate::service::Service::node_activate_reconcile`]
-            tracing::info!("Node {node} is unavailable during detach: proceeding anyway, it will be detached on next activation");
-            self.observed.locations.remove(&node.get_id());
-            return Ok(());
-        }
-
-        self.observed
-            .locations
-            .insert(node.get_id(), ObservedStateLocation { conf: None });
-
-        // TODO: amend locations that use long-polling: they will hit this timeout.
-        let timeout = Duration::from_secs(25);
-
-        tracing::info!("location_config({node}) calling: {:?}", config);
-        let tenant_shard_id = self.tenant_shard_id;
-        let config_ref = &config;
-        match node
-            .with_client_retries(
-                |client| async move {
-                    let config = config_ref.clone();
-                    client
-                        .location_config(tenant_shard_id, config.clone(), flush_ms, lazy)
-                        .await
-                },
-                &self.service_config.jwt_token,
-                1,
-                3,
-                timeout,
-                &self.cancel,
-            )
-            .await
-        {
-            Some(Ok(_)) => {}
-            Some(Err(e)) => return Err(e.into()),
-            None => return Err(ReconcileError::Cancel),
-        };
-        tracing::info!("location_config({node}) complete: {:?}", config);
-
-        match config.mode {
-            LocationConfigMode::Detached => {
-                self.observed.locations.remove(&node.get_id());
-            }
-            _ => {
-                self.observed
-                    .locations
-                    .insert(node.get_id(), ObservedStateLocation { conf: Some(config) });
-            }
-        }
-
-        Ok(())
-    }
-
-    fn get_node(&self, node_id: &NodeId) -> Option<&Node> {
-        if let Some(node) = self.intent.attached.as_ref() {
-            if node.get_id() == *node_id {
-                return Some(node);
-            }
-        }
-
-        if let Some(node) = self
-            .intent
-            .secondary
-            .iter()
-            .find(|n| n.get_id() == *node_id)
-        {
-            return Some(node);
-        }
-
-        if let Some(node) = self.detach.iter().find(|n| n.get_id() == *node_id) {
-            return Some(node);
-        }
-
-        None
-    }
-
-    async fn maybe_live_migrate(&mut self) -> Result<(), ReconcileError> {
-        let destination = if let Some(node) = &self.intent.attached {
-            match self.observed.locations.get(&node.get_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
-                        }
-                        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
-                }
-            }
-        } 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 {
-                    // 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 let Some(node) = self.get_node(node_id) {
-                        if node.is_available() {
-                            origin = Some(node.clone());
-                            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.clone()).await?;
-
-        Ok(())
-    }
-
-    async fn get_lsns(
-        &self,
-        tenant_shard_id: TenantShardId,
-        node: &Node,
-    ) -> anyhow::Result<HashMap<TimelineId, Lsn>> {
-        let client = PageserverClient::new(
-            node.get_id(),
-            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: &Node,
-    ) -> Result<(), ReconcileError> {
-        // This is not the timeout for a request, but the total amount of time we're willing to wait
-        // for a secondary location to get up to date before
-        const TOTAL_DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(300);
-
-        // This the long-polling interval for the secondary download requests we send to destination pageserver
-        // during a migration.
-        const REQUEST_DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(20);
-
-        let started_at = Instant::now();
-
-        loop {
-            let (status, progress) = match node
-                .with_client_retries(
-                    |client| async move {
-                        client
-                            .tenant_secondary_download(
-                                tenant_shard_id,
-                                Some(REQUEST_DOWNLOAD_TIMEOUT),
-                            )
-                            .await
-                    },
-                    &self.service_config.jwt_token,
-                    1,
-                    3,
-                    REQUEST_DOWNLOAD_TIMEOUT * 2,
-                    &self.cancel,
-                )
-                .await
-            {
-                None => Err(ReconcileError::Cancel),
-                Some(Ok(v)) => Ok(v),
-                Some(Err(e)) => {
-                    // Give up, but proceed: it's unfortunate if we couldn't freshen the destination before
-                    // attaching, but we should not let an issue with a secondary location stop us proceeding
-                    // with a live migration.
-                    tracing::warn!("Failed to prepare by downloading layers on node {node}: {e})");
-                    return Ok(());
-                }
-            }?;
-
-            if status == StatusCode::OK {
-                tracing::info!(
-                    "Downloads to {} complete: {}/{} layers, {}/{} bytes",
-                    node,
-                    progress.layers_downloaded,
-                    progress.layers_total,
-                    progress.bytes_downloaded,
-                    progress.bytes_total
-                );
-                return Ok(());
-            } else if status == StatusCode::ACCEPTED {
-                let total_runtime = started_at.elapsed();
-                if total_runtime > TOTAL_DOWNLOAD_TIMEOUT {
-                    tracing::warn!("Timed out after {}ms downloading layers to {node}.  Progress so far: {}/{} layers, {}/{} bytes",
-                        total_runtime.as_millis(),
-                        progress.layers_downloaded,
-                        progress.layers_total,
-                        progress.bytes_downloaded,
-                        progress.bytes_total
-                    );
-                    // Give up, but proceed: an incompletely warmed destination doesn't prevent migration working,
-                    // it just makes the I/O performance for users less good.
-                    return Ok(());
-                }
-
-                // Log and proceed around the loop to retry.  We don't sleep between requests, because our HTTP call
-                // to the pageserver is a long-poll.
-                tracing::info!(
-                    "Downloads to {} not yet complete: {}/{} layers, {}/{} bytes",
-                    node,
-                    progress.layers_downloaded,
-                    progress.layers_total,
-                    progress.bytes_downloaded,
-                    progress.bytes_total
-                );
-            }
-        }
-    }
-
-    async fn await_lsn(
-        &self,
-        tenant_shard_id: TenantShardId,
-        node: &Node,
-        baseline: HashMap<TimelineId, Lsn>,
-    ) -> anyhow::Result<()> {
-        loop {
-            let latest = match self.get_lsns(tenant_shard_id, node).await {
-                Ok(l) => l,
-                Err(e) => {
-                    tracing::info!("🕑 Can't get LSNs on node {node} yet, waiting ({e})",);
-                    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) => {
-                        tracing::info!("🕑 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 {
-                tracing::info!("✅ LSN caught up.  Proceeding...");
-                break;
-            } else {
-                std::thread::sleep(Duration::from_millis(500));
-            }
-        }
-
-        Ok(())
-    }
-
-    pub async fn live_migrate(
-        &mut self,
-        origin_ps: Node,
-        dest_ps: Node,
-    ) -> Result<(), ReconcileError> {
-        // `maybe_live_migrate` is responsibble for sanity of inputs
-        assert!(origin_ps.get_id() != dest_ps.get_id());
-
-        fn build_location_config(
-            shard: &ShardIdentity,
-            config: &TenantConfig,
-            mode: LocationConfigMode,
-            generation: Option<Generation>,
-            secondary_conf: Option<LocationConfigSecondary>,
-        ) -> 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.literal(),
-                shard_stripe_size: shard.stripe_size.0,
-            }
-        }
-
-        tracing::info!("🔁 Switching origin node {origin_ps} to stale mode",);
-
-        // 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,
-            self.generation,
-            None,
-        );
-        self.location_config(&origin_ps, stale_conf, Some(Duration::from_secs(10)), false)
-            .await?;
-
-        let baseline_lsns = Some(self.get_lsns(self.tenant_shard_id, &origin_ps).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.get_id()) {
-            if let Some(destination_conf) = &destination_conf.conf {
-                if destination_conf.mode == LocationConfigMode::Secondary {
-                    tracing::info!("🔁 Downloading latest layers to destination node {dest_ps}",);
-                    self.secondary_download(self.tenant_shard_id, &dest_ps)
-                        .await?;
-                }
-            }
-        }
-
-        // Increment generation before attaching to new pageserver
-        self.generation = Some(
-            self.persistence
-                .increment_generation(self.tenant_shard_id, dest_ps.get_id())
-                .await?,
-        );
-
-        let dest_conf = build_location_config(
-            &self.shard,
-            &self.config,
-            LocationConfigMode::AttachedMulti,
-            self.generation,
-            None,
-        );
-
-        tracing::info!("🔁 Attaching to pageserver {dest_ps}");
-        self.location_config(&dest_ps, dest_conf, None, false)
-            .await?;
-
-        if let Some(baseline) = baseline_lsns {
-            tracing::info!("🕑 Waiting for LSN to catch up...");
-            self.await_lsn(self.tenant_shard_id, &dest_ps, baseline)
-                .await?;
-        }
-
-        tracing::info!("🔁 Notifying compute to use pageserver {dest_ps}");
-
-        // During a live migration it is unhelpful to proceed if we couldn't notify compute: if we detach
-        // the origin without notifying compute, we will render the tenant unavailable.
-        while let Err(e) = self.compute_notify().await {
-            match e {
-                NotifyError::Fatal(_) => return Err(ReconcileError::Notify(e)),
-                NotifyError::ShuttingDown => return Err(ReconcileError::Cancel),
-                _ => {
-                    tracing::warn!(
-                        "Live migration blocked by compute notification error, retrying: {e}"
-                    );
-                }
-            }
-        }
-
-        // Downgrade the origin to secondary.  If the tenant's policy is PlacementPolicy::Attached(0), 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, origin_secondary_conf.clone(), None, false)
-            .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.get_id(),
-            ObservedStateLocation {
-                conf: Some(origin_secondary_conf),
-            },
-        );
-
-        tracing::info!("🔁 Switching to AttachedSingle mode on node {dest_ps}",);
-        let dest_final_conf = build_location_config(
-            &self.shard,
-            &self.config,
-            LocationConfigMode::AttachedSingle,
-            self.generation,
-            None,
-        );
-        self.location_config(&dest_ps, dest_final_conf.clone(), None, false)
-            .await?;
-        self.observed.locations.insert(
-            dest_ps.get_id(),
-            ObservedStateLocation {
-                conf: Some(dest_final_conf),
-            },
-        );
-
-        tracing::info!("✅ Migration complete");
-
-        Ok(())
-    }
-
-    async fn maybe_refresh_observed(&mut self) -> Result<(), ReconcileError> {
-        // If the attached node has uncertain state, read it from the pageserver before proceeding: this
-        // is important to avoid spurious generation increments.
-        //
-        // We don't need to do this for secondary/detach locations because it's harmless to just PUT their
-        // location conf, whereas for attached locations it can interrupt clients if we spuriously destroy/recreate
-        // the `Timeline` object in the pageserver.
-
-        let Some(attached_node) = self.intent.attached.as_ref() else {
-            // Nothing to do
-            return Ok(());
-        };
-
-        if matches!(
-            self.observed.locations.get(&attached_node.get_id()),
-            Some(ObservedStateLocation { conf: None })
-        ) {
-            let tenant_shard_id = self.tenant_shard_id;
-            let observed_conf = match attached_node
-                .with_client_retries(
-                    |client| async move { client.get_location_config(tenant_shard_id).await },
-                    &self.service_config.jwt_token,
-                    1,
-                    1,
-                    Duration::from_secs(5),
-                    &self.cancel,
-                )
-                .await
-            {
-                Some(Ok(observed)) => Some(observed),
-                Some(Err(mgmt_api::Error::ApiError(status, _msg)))
-                    if status == StatusCode::NOT_FOUND =>
-                {
-                    None
-                }
-                Some(Err(e)) => return Err(e.into()),
-                None => return Err(ReconcileError::Cancel),
-            };
-            tracing::info!("Scanned location configuration on {attached_node}: {observed_conf:?}");
-            match observed_conf {
-                Some(conf) => {
-                    // Pageserver returned a state: update it in observed.  This may still be an indeterminate (None) state,
-                    // if internally the pageserver's TenantSlot was being mutated (e.g. some long running API call is still running)
-                    self.observed
-                        .locations
-                        .insert(attached_node.get_id(), ObservedStateLocation { conf });
-                }
-                None => {
-                    // Pageserver returned 404: we have confirmation that there is no state for this shard on that pageserver.
-                    self.observed.locations.remove(&attached_node.get_id());
-                }
-            }
-        }
-
-        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> {
-        // Prepare: if we have uncertain `observed` state for our would-be attachement location, then refresh it
-        self.maybe_refresh_observed().await?;
-
-        // Special case: live migration
-        self.maybe_live_migrate().await?;
-
-        // If the attached pageserver is not attached, do so now.
-        if let Some(node) = self.intent.attached.as_ref() {
-            // If we are in an attached policy, then generation must have been set (null generations
-            // are only present when a tenant is initially loaded with a secondary policy)
-            debug_assert!(self.generation.is_some());
-            let Some(generation) = self.generation else {
-                return Err(ReconcileError::Other(anyhow::anyhow!(
-                    "Attempted to attach with NULL generation"
-                )));
-            };
-
-            let mut wanted_conf = attached_location_conf(
-                generation,
-                &self.shard,
-                &self.config,
-                !self.intent.secondary.is_empty(),
-            );
-            match self.observed.locations.get(&node.get_id()) {
-                Some(conf) if conf.conf.as_ref() == Some(&wanted_conf) => {
-                    // Nothing to do
-                    tracing::info!(node_id=%node.get_id(), "Observed configuration already correct.")
-                }
-                observed => {
-                    // 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.
-
-                    // The general case is to increment the generation.  However, there are cases
-                    // where this is not necessary:
-                    // - if we are only updating the TenantConf part of the location
-                    // - if we are only changing the attachment mode (e.g. going to attachedmulti or attachedstale)
-                    //   and the location was already in the correct generation
-                    let increment_generation = match observed {
-                        None => true,
-                        Some(ObservedStateLocation { conf: None }) => true,
-                        Some(ObservedStateLocation {
-                            conf: Some(observed),
-                        }) => {
-                            let generations_match = observed.generation == wanted_conf.generation;
-
-                            use LocationConfigMode::*;
-                            let mode_transition_requires_gen_inc =
-                                match (observed.mode, wanted_conf.mode) {
-                                    // Usually the short-lived attachment modes (multi and stale) are only used
-                                    // in the case of [`Self::live_migrate`], but it is simple to handle them correctly
-                                    // here too.  Locations are allowed to go Single->Stale and Multi->Single within the same generation.
-                                    (AttachedSingle, AttachedStale) => false,
-                                    (AttachedMulti, AttachedSingle) => false,
-                                    (lhs, rhs) => lhs != rhs,
-                                };
-
-                            !generations_match || mode_transition_requires_gen_inc
-                        }
-                    };
-
-                    if increment_generation {
-                        let generation = self
-                            .persistence
-                            .increment_generation(self.tenant_shard_id, node.get_id())
-                            .await?;
-                        self.generation = Some(generation);
-                        wanted_conf.generation = generation.into();
-                    }
-                    tracing::info!(node_id=%node.get_id(), "Observed configuration requires update.");
-
-                    // Because `node` comes from a ref to &self, clone it before calling into a &mut self
-                    // function: this could be avoided by refactoring the state mutated by location_config into
-                    // a separate type to Self.
-                    let node = node.clone();
-
-                    // Use lazy=true, because we may run many of Self concurrently, and do not want to
-                    // overload the pageserver with logical size calculations.
-                    self.location_config(&node, wanted_conf, None, true).await?;
-                    self.compute_notify().await?;
-                }
-            }
-        }
-
-        // Configure secondary locations: if these were previously attached this
-        // implicitly downgrades them from attached to secondary.
-        let mut changes = Vec::new();
-        for node in &self.intent.secondary {
-            let wanted_conf = secondary_location_conf(&self.shard, &self.config);
-            match self.observed.locations.get(&node.get_id()) {
-                Some(conf) if conf.conf.as_ref() == Some(&wanted_conf) => {
-                    // Nothing to do
-                    tracing::info!(node_id=%node.get_id(), "Observed configuration already correct.")
-                }
-                _ => {
-                    // In all cases other than a matching observed configuration, we will
-                    // reconcile this location.
-                    tracing::info!(node_id=%node.get_id(), "Observed configuration requires update.");
-                    changes.push((node.clone(), wanted_conf))
-                }
-            }
-        }
-
-        // Detach any extraneous pageservers that are no longer referenced
-        // by our intent.
-        for node in &self.detach {
-            changes.push((
-                node.clone(),
-                LocationConfig {
-                    mode: LocationConfigMode::Detached,
-                    generation: None,
-                    secondary_conf: None,
-                    shard_number: self.shard.number.0,
-                    shard_count: self.shard.count.literal(),
-                    shard_stripe_size: self.shard.stripe_size.0,
-                    tenant_conf: self.config.clone(),
-                },
-            ));
-        }
-
-        for (node, conf) in changes {
-            if self.cancel.is_cancelled() {
-                return Err(ReconcileError::Cancel);
-            }
-            self.location_config(&node, conf, None, false).await?;
-        }
-
-        Ok(())
-    }
-
-    pub(crate) async fn compute_notify(&mut self) -> Result<(), NotifyError> {
-        // Whenever a particular Reconciler emits a notification, it is always notifying for the intended
-        // destination.
-        if let Some(node) = &self.intent.attached {
-            let result = self
-                .compute_hook
-                .notify(
-                    self.tenant_shard_id,
-                    node.get_id(),
-                    self.shard.stripe_size,
-                    &self.cancel,
-                )
-                .await;
-            if let Err(e) = &result {
-                // It is up to the caller whether they want to drop out on this error, but they don't have to:
-                // in general we should avoid letting unavailability of the cloud control plane stop us from
-                // making progress.
-                tracing::warn!("Failed to notify compute of attached pageserver {node}: {e}");
-                // Set this flag so that in our ReconcileResult we will set the flag on the shard that it
-                // needs to retry at some point.
-                self.compute_notify_failure = true;
-            }
-            result
-        } else {
-            Ok(())
-        }
-    }
-}
-
-/// We tweak the externally-set TenantConfig while configuring
-/// locations, using our awareness of whether secondary locations
-/// are in use to automatically enable/disable heatmap uploads.
-fn ha_aware_config(config: &TenantConfig, has_secondaries: bool) -> TenantConfig {
-    let mut config = config.clone();
-    if has_secondaries {
-        if config.heatmap_period.is_none() {
-            config.heatmap_period = Some(DEFAULT_HEATMAP_PERIOD.to_string());
-        }
-    } else {
-        config.heatmap_period = None;
-    }
-    config
-}
-
-pub(crate) fn attached_location_conf(
-    generation: Generation,
-    shard: &ShardIdentity,
-    config: &TenantConfig,
-    has_secondaries: bool,
-) -> LocationConfig {
-    LocationConfig {
-        mode: LocationConfigMode::AttachedSingle,
-        generation: generation.into(),
-        secondary_conf: None,
-        shard_number: shard.number.0,
-        shard_count: shard.count.literal(),
-        shard_stripe_size: shard.stripe_size.0,
-        tenant_conf: ha_aware_config(config, has_secondaries),
-    }
-}
-
-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.literal(),
-        shard_stripe_size: shard.stripe_size.0,
-        tenant_conf: ha_aware_config(config, true),
-    }
-}

control_plane/attachment_service/src/scheduler.rs

@@ -1,445 +0,0 @@
-use crate::{node::Node, tenant_state::TenantState};
-use pageserver_api::controller_api::UtilizationScore;
-use serde::Serialize;
-use std::collections::HashMap;
-use utils::{http::error::ApiError, id::NodeId};
-
-/// 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<ScheduleError> for ApiError {
-    fn from(value: ScheduleError) -> Self {
-        ApiError::Conflict(format!("Scheduling error: {}", value))
-    }
-}
-
-#[derive(Serialize, Eq, PartialEq)]
-pub enum MaySchedule {
-    Yes(UtilizationScore),
-    No,
-}
-
-#[derive(Serialize)]
-struct SchedulerNode {
-    /// How many shards are currently scheduled on this node, via their [`crate::tenant_state::IntentState`].
-    shard_count: usize,
-
-    /// Whether this node is currently elegible to have new shards scheduled (this is derived
-    /// from a node's availability state and scheduling policy).
-    may_schedule: MaySchedule,
-}
-
-impl PartialEq for SchedulerNode {
-    fn eq(&self, other: &Self) -> bool {
-        let may_schedule_matches = matches!(
-            (&self.may_schedule, &other.may_schedule),
-            (MaySchedule::Yes(_), MaySchedule::Yes(_)) | (MaySchedule::No, MaySchedule::No)
-        );
-
-        may_schedule_matches && self.shard_count == other.shard_count
-    }
-}
-
-impl Eq for SchedulerNode {}
-
-/// This type is responsible for selecting which node is used when a tenant shard needs to choose a pageserver
-/// on which to run.
-///
-/// The type has no persistent state of its own: this is all populated at startup.  The Serialize
-/// impl is only for debug dumps.
-#[derive(Serialize)]
-pub(crate) struct Scheduler {
-    nodes: HashMap<NodeId, SchedulerNode>,
-}
-
-/// Score for soft constraint scheduling: lower scores are preferred to higher scores.
-///
-/// For example, we may set an affinity score based on the number of shards from the same
-/// tenant already on a node, to implicitly prefer to balance out shards.
-#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord)]
-pub(crate) struct AffinityScore(pub(crate) usize);
-
-impl AffinityScore {
-    /// If we have no anti-affinity at all toward a node, this is its score.  It means
-    /// the scheduler has a free choice amongst nodes with this score, and may pick a node
-    /// based on other information such as total utilization.
-    pub(crate) const FREE: Self = Self(0);
-
-    pub(crate) fn inc(&mut self) {
-        self.0 += 1;
-    }
-}
-
-impl std::ops::Add for AffinityScore {
-    type Output = Self;
-
-    fn add(self, rhs: Self) -> Self::Output {
-        Self(self.0 + rhs.0)
-    }
-}
-
-// For carrying state between multiple calls to [`TenantState::schedule`], e.g. when calling
-// it for many shards in the same tenant.
-#[derive(Debug, Default)]
-pub(crate) struct ScheduleContext {
-    /// Sparse map of nodes: omitting a node implicitly makes its affinity [`AffinityScore::FREE`]
-    pub(crate) nodes: HashMap<NodeId, AffinityScore>,
-
-    /// Specifically how many _attached_ locations are on each node
-    pub(crate) attached_nodes: HashMap<NodeId, usize>,
-}
-
-impl ScheduleContext {
-    /// Input is a list of nodes we would like to avoid using again within this context.  The more
-    /// times a node is passed into this call, the less inclined we are to use it.
-    pub(crate) fn avoid(&mut self, nodes: &[NodeId]) {
-        for node_id in nodes {
-            let entry = self.nodes.entry(*node_id).or_insert(AffinityScore::FREE);
-            entry.inc()
-        }
-    }
-
-    pub(crate) fn push_attached(&mut self, node_id: NodeId) {
-        let entry = self.attached_nodes.entry(node_id).or_default();
-        *entry += 1;
-    }
-
-    pub(crate) fn get_node_affinity(&self, node_id: NodeId) -> AffinityScore {
-        self.nodes
-            .get(&node_id)
-            .copied()
-            .unwrap_or(AffinityScore::FREE)
-    }
-
-    pub(crate) fn get_node_attachments(&self, node_id: NodeId) -> usize {
-        self.attached_nodes.get(&node_id).copied().unwrap_or(0)
-    }
-}
-
-impl Scheduler {
-    pub(crate) fn new<'a>(nodes: impl Iterator<Item = &'a Node>) -> Self {
-        let mut scheduler_nodes = HashMap::new();
-        for node in nodes {
-            scheduler_nodes.insert(
-                node.get_id(),
-                SchedulerNode {
-                    shard_count: 0,
-                    may_schedule: node.may_schedule(),
-                },
-            );
-        }
-
-        Self {
-            nodes: scheduler_nodes,
-        }
-    }
-
-    /// For debug/support: check that our internal statistics are in sync with the state of
-    /// the nodes & tenant shards.
-    ///
-    /// If anything is inconsistent, log details and return an error.
-    pub(crate) fn consistency_check<'a>(
-        &self,
-        nodes: impl Iterator<Item = &'a Node>,
-        shards: impl Iterator<Item = &'a TenantState>,
-    ) -> anyhow::Result<()> {
-        let mut expect_nodes: HashMap<NodeId, SchedulerNode> = HashMap::new();
-        for node in nodes {
-            expect_nodes.insert(
-                node.get_id(),
-                SchedulerNode {
-                    shard_count: 0,
-                    may_schedule: node.may_schedule(),
-                },
-            );
-        }
-
-        for shard in shards {
-            if let Some(node_id) = shard.intent.get_attached() {
-                match expect_nodes.get_mut(node_id) {
-                    Some(node) => node.shard_count += 1,
-                    None => anyhow::bail!(
-                        "Tenant {} references nonexistent node {}",
-                        shard.tenant_shard_id,
-                        node_id
-                    ),
-                }
-            }
-
-            for node_id in shard.intent.get_secondary() {
-                match expect_nodes.get_mut(node_id) {
-                    Some(node) => node.shard_count += 1,
-                    None => anyhow::bail!(
-                        "Tenant {} references nonexistent node {}",
-                        shard.tenant_shard_id,
-                        node_id
-                    ),
-                }
-            }
-        }
-
-        for (node_id, expect_node) in &expect_nodes {
-            let Some(self_node) = self.nodes.get(node_id) else {
-                anyhow::bail!("Node {node_id} not found in Self")
-            };
-
-            if self_node != expect_node {
-                tracing::error!("Inconsistency detected in scheduling state for node {node_id}");
-                tracing::error!("Expected state: {}", serde_json::to_string(expect_node)?);
-                tracing::error!("Self state: {}", serde_json::to_string(self_node)?);
-
-                anyhow::bail!("Inconsistent state on {node_id}");
-            }
-        }
-
-        if expect_nodes.len() != self.nodes.len() {
-            // We just checked that all the expected nodes are present.  If the lengths don't match,
-            // it means that we have nodes in Self that are unexpected.
-            for node_id in self.nodes.keys() {
-                if !expect_nodes.contains_key(node_id) {
-                    anyhow::bail!("Node {node_id} found in Self but not in expected nodes");
-                }
-            }
-        }
-
-        Ok(())
-    }
-
-    /// Increment the reference count of a node.  This reference count is used to guide scheduling
-    /// decisions, not for memory management: it represents one tenant shard whose IntentState targets
-    /// this node.
-    ///
-    /// It is an error to call this for a node that is not known to the scheduler (i.e. passed into
-    /// [`Self::new`] or [`Self::node_upsert`])
-    pub(crate) fn node_inc_ref(&mut self, node_id: NodeId) {
-        let Some(node) = self.nodes.get_mut(&node_id) else {
-            tracing::error!("Scheduler missing node {node_id}");
-            debug_assert!(false);
-            return;
-        };
-
-        node.shard_count += 1;
-    }
-
-    /// Decrement a node's reference count.  Inverse of [`Self::node_inc_ref`].
-    pub(crate) fn node_dec_ref(&mut self, node_id: NodeId) {
-        let Some(node) = self.nodes.get_mut(&node_id) else {
-            debug_assert!(false);
-            tracing::error!("Scheduler missing node {node_id}");
-            return;
-        };
-
-        node.shard_count -= 1;
-    }
-
-    pub(crate) fn node_upsert(&mut self, node: &Node) {
-        use std::collections::hash_map::Entry::*;
-        match self.nodes.entry(node.get_id()) {
-            Occupied(mut entry) => {
-                entry.get_mut().may_schedule = node.may_schedule();
-            }
-            Vacant(entry) => {
-                entry.insert(SchedulerNode {
-                    shard_count: 0,
-                    may_schedule: node.may_schedule(),
-                });
-            }
-        }
-    }
-
-    pub(crate) fn node_remove(&mut self, node_id: NodeId) {
-        if self.nodes.remove(&node_id).is_none() {
-            tracing::warn!(node_id=%node_id, "Removed non-existent node from scheduler");
-        }
-    }
-
-    /// Where we have several nodes to choose from, for example when picking a secondary location
-    /// to promote to an attached location, this method may be used to pick the best choice based
-    /// on the scheduler's knowledge of utilization and availability.
-    ///
-    /// If the input is empty, or all the nodes are not elegible for scheduling, return None: the
-    /// caller can pick a node some other way.
-    pub(crate) fn node_preferred(&self, nodes: &[NodeId]) -> Option<NodeId> {
-        if nodes.is_empty() {
-            return None;
-        }
-
-        // TODO: When the utilization score returned by the pageserver becomes meaningful,
-        // schedule based on that instead of the shard count.
-        let node = nodes
-            .iter()
-            .map(|node_id| {
-                let may_schedule = self
-                    .nodes
-                    .get(node_id)
-                    .map(|n| n.may_schedule != MaySchedule::No)
-                    .unwrap_or(false);
-                (*node_id, may_schedule)
-            })
-            .max_by_key(|(_n, may_schedule)| *may_schedule);
-
-        // If even the preferred node has may_schedule==false, return None
-        node.and_then(|(node_id, may_schedule)| if may_schedule { Some(node_id) } else { None })
-    }
-
-    /// hard_exclude: it is forbidden to use nodes in this list, typically becacuse they
-    /// are already in use by this shard -- we use this to avoid picking the same node
-    /// as both attached and secondary location.  This is a hard constraint: if we cannot
-    /// find any nodes that aren't in this list, then we will return a [`ScheduleError::ImpossibleConstraint`].
-    ///
-    /// context: we prefer to avoid using nodes identified in the context, according
-    /// to their anti-affinity score.  We use this to prefeer to avoid placing shards in
-    /// the same tenant on the same node.  This is a soft constraint: the context will never
-    /// cause us to fail to schedule a shard.
-    pub(crate) fn schedule_shard(
-        &self,
-        hard_exclude: &[NodeId],
-        context: &ScheduleContext,
-    ) -> Result<NodeId, ScheduleError> {
-        if self.nodes.is_empty() {
-            return Err(ScheduleError::NoPageservers);
-        }
-
-        let mut scores: Vec<(NodeId, AffinityScore, usize)> = self
-            .nodes
-            .iter()
-            .filter_map(|(k, v)| {
-                if hard_exclude.contains(k) || v.may_schedule == MaySchedule::No {
-                    None
-                } else {
-                    Some((
-                        *k,
-                        context.nodes.get(k).copied().unwrap_or(AffinityScore::FREE),
-                        v.shard_count,
-                    ))
-                }
-            })
-            .collect();
-
-        // Sort by, in order of precedence:
-        //  1st: Affinity score.  We should never pick a higher-score node if a lower-score node is available
-        //  2nd: Utilization.  Within nodes with the same affinity, use the least loaded nodes.
-        //  3rd: Node ID.  This is a convenience to make selection deterministic in tests and empty systems.
-        scores.sort_by_key(|i| (i.1, i.2, i.0));
-
-        if scores.is_empty() {
-            // After applying constraints, no pageservers were left.  We log some detail about
-            // the state of nodes to help understand why this happened.  This is not logged as an error because
-            // it is legitimately possible for enough nodes to be Offline to prevent scheduling a shard.
-            tracing::info!("Scheduling failure, while excluding {hard_exclude:?}, node states:");
-            for (node_id, node) in &self.nodes {
-                tracing::info!(
-                    "Node {node_id}: may_schedule={} shards={}",
-                    node.may_schedule != MaySchedule::No,
-                    node.shard_count
-                );
-            }
-
-            return Err(ScheduleError::ImpossibleConstraint);
-        }
-
-        // Lowest score wins
-        let node_id = scores.first().unwrap().0;
-        tracing::info!(
-            "scheduler selected node {node_id} (elegible nodes {:?}, hard exclude: {hard_exclude:?}, soft exclude: {context:?})",
-            scores.iter().map(|i| i.0 .0).collect::<Vec<_>>()
-        );
-
-        // Note that we do not update shard count here to reflect the scheduling: that
-        // is IntentState's job when the scheduled location is used.
-
-        Ok(node_id)
-    }
-
-    /// Unit test access to internal state
-    #[cfg(test)]
-    pub(crate) fn get_node_shard_count(&self, node_id: NodeId) -> usize {
-        self.nodes.get(&node_id).unwrap().shard_count
-    }
-}
-
-#[cfg(test)]
-pub(crate) mod test_utils {
-
-    use crate::node::Node;
-    use pageserver_api::controller_api::{NodeAvailability, UtilizationScore};
-    use std::collections::HashMap;
-    use utils::id::NodeId;
-    /// Test helper: synthesize the requested number of nodes, all in active state.
-    ///
-    /// Node IDs start at one.
-    pub(crate) fn make_test_nodes(n: u64) -> HashMap<NodeId, Node> {
-        (1..n + 1)
-            .map(|i| {
-                (NodeId(i), {
-                    let mut node = Node::new(
-                        NodeId(i),
-                        format!("httphost-{i}"),
-                        80 + i as u16,
-                        format!("pghost-{i}"),
-                        5432 + i as u16,
-                    );
-                    node.set_availability(NodeAvailability::Active(UtilizationScore::worst()));
-                    assert!(node.is_available());
-                    node
-                })
-            })
-            .collect()
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    use crate::tenant_state::IntentState;
-    #[test]
-    fn scheduler_basic() -> anyhow::Result<()> {
-        let nodes = test_utils::make_test_nodes(2);
-
-        let mut scheduler = Scheduler::new(nodes.values());
-        let mut t1_intent = IntentState::new();
-        let mut t2_intent = IntentState::new();
-
-        let context = ScheduleContext::default();
-
-        let scheduled = scheduler.schedule_shard(&[], &context)?;
-        t1_intent.set_attached(&mut scheduler, Some(scheduled));
-        let scheduled = scheduler.schedule_shard(&[], &context)?;
-        t2_intent.set_attached(&mut scheduler, Some(scheduled));
-
-        assert_eq!(scheduler.nodes.get(&NodeId(1)).unwrap().shard_count, 1);
-        assert_eq!(scheduler.nodes.get(&NodeId(2)).unwrap().shard_count, 1);
-
-        let scheduled = scheduler.schedule_shard(&t1_intent.all_pageservers(), &context)?;
-        t1_intent.push_secondary(&mut scheduler, scheduled);
-
-        assert_eq!(scheduler.nodes.get(&NodeId(1)).unwrap().shard_count, 1);
-        assert_eq!(scheduler.nodes.get(&NodeId(2)).unwrap().shard_count, 2);
-
-        t1_intent.clear(&mut scheduler);
-        assert_eq!(scheduler.nodes.get(&NodeId(1)).unwrap().shard_count, 0);
-        assert_eq!(scheduler.nodes.get(&NodeId(2)).unwrap().shard_count, 1);
-
-        if cfg!(debug_assertions) {
-            // Dropping an IntentState without clearing it causes a panic in debug mode,
-            // because we have failed to properly update scheduler shard counts.
-            let result = std::panic::catch_unwind(move || {
-                drop(t2_intent);
-            });
-            assert!(result.is_err());
-        } else {
-            t2_intent.clear(&mut scheduler);
-            assert_eq!(scheduler.nodes.get(&NodeId(1)).unwrap().shard_count, 0);
-            assert_eq!(scheduler.nodes.get(&NodeId(2)).unwrap().shard_count, 0);
-        }
-
-        Ok(())
-    }
-}

control_plane/attachment_service/src/schema.rs

@@ -1,29 +0,0 @@
-// @generated automatically by Diesel CLI.
-
-diesel::table! {
-    nodes (node_id) {
-        node_id -> Int8,
-        scheduling_policy -> Varchar,
-        listen_http_addr -> Varchar,
-        listen_http_port -> Int4,
-        listen_pg_addr -> Varchar,
-        listen_pg_port -> Int4,
-    }
-}
-
-diesel::table! {
-    tenant_shards (tenant_id, shard_number, shard_count) {
-        tenant_id -> Varchar,
-        shard_number -> Int4,
-        shard_count -> Int4,
-        shard_stripe_size -> Int4,
-        generation -> Nullable<Int4>,
-        generation_pageserver -> Nullable<Int8>,
-        placement_policy -> Varchar,
-        splitting -> Int2,
-        config -> Text,
-        scheduling_policy -> Varchar,
-    }
-}
-
-diesel::allow_tables_to_appear_in_same_query!(nodes, tenant_shards,);

control_plane/storcon_cli/src/main.rs

@@ -223,7 +223,7 @@ impl Client {
         }
     }
 
-    /// Simple HTTP request wrapper for calling into attachment service
+    /// Simple HTTP request wrapper for calling into storage controller
     async fn dispatch<RQ, RS>(
         &self,
         method: hyper::Method,