wip

.github/workflows/_build-and-test-locally.yml

@@ -391,6 +391,10 @@ jobs:
           PAGESERVER_VIRTUAL_FILE_IO_ENGINE: tokio-epoll-uring
           USE_LFC: ${{ matrix.lfc_state == 'with-lfc' && 'true' || 'false' }}
 
+      # Temporary disable this step until we figure out why it's so flaky
+      # Ref https://github.com/neondatabase/neon/issues/4540
       - name: Merge and upload coverage data
-        if: inputs.build-type == 'debug'
+        if: |
+          false &&
+          inputs.build-type == 'debug' && matrix.pg_version == 'v16'
         uses: ./.github/actions/save-coverage-data

.github/workflows/build_and_test.yml

@@ -484,6 +484,21 @@ jobs:
       - name: Merge coverage data
         run: scripts/coverage "--profraw-prefix=$GITHUB_JOB" --dir=/tmp/coverage merge
 
+      - name: Build coverage report
+        env:
+          COMMIT_URL: ${{ github.server_url }}/${{ github.repository }}/commit/${{ github.event.pull_request.head.sha || github.sha }}
+        run: |
+          scripts/coverage --dir=/tmp/coverage \
+            report \
+            --input-objects=/tmp/coverage/binaries.list \
+            --commit-url=${COMMIT_URL} \
+            --format=github
+
+          scripts/coverage --dir=/tmp/coverage \
+            report \
+            --input-objects=/tmp/coverage/binaries.list \
+            --format=lcov
+
       - name: Build coverage report NEW
         id: upload-coverage-report-new
         env:
@@ -496,13 +511,6 @@ jobs:
           CURRENT="${COMMIT_SHA}"
           BASELINE="$(git merge-base $BASE_SHA $CURRENT)"
 
-          mkdir /tmp/coverage/report
-
-          scripts/coverage --dir=/tmp/coverage \
-            report \
-            --input-objects=/tmp/coverage/binaries.list \
-            --format=lcov
-
           cp /tmp/coverage/report/lcov.info ./${CURRENT}.info
 
           GENHTML_ARGS="--ignore-errors path,unmapped,empty --synthesize-missing --demangle-cpp rustfilt --output-directory lcov-html ${CURRENT}.info"

compute_tools/src/bin/compute_ctl.rs

@@ -36,8 +36,6 @@
 use std::ffi::OsString;
 use std::fs::File;
 use std::process::exit;
-use std::sync::Arc;
-use std::sync::atomic::AtomicU64;
 use std::sync::mpsc;
 use std::thread;
 use std::time::Duration;
@@ -192,9 +190,7 @@ fn main() -> Result<()> {
             cgroup: cli.cgroup,
             #[cfg(target_os = "linux")]
             vm_monitor_addr: cli.vm_monitor_addr,
-            installed_extensions_collection_interval: Arc::new(AtomicU64::new(
-                cli.installed_extensions_collection_interval,
-            )),
+            installed_extensions_collection_interval: cli.installed_extensions_collection_interval,
         },
         config,
     )?;

compute_tools/src/compute.rs

@@ -25,7 +25,7 @@ use std::os::unix::fs::{PermissionsExt, symlink};
 use std::path::Path;
 use std::process::{Command, Stdio};
 use std::str::FromStr;
-use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::{Arc, Condvar, Mutex, RwLock};
 use std::time::{Duration, Instant};
 use std::{env, fs};
@@ -70,7 +70,6 @@ pub static BUILD_TAG: Lazy<String> = Lazy::new(|| {
         .unwrap_or(BUILD_TAG_DEFAULT)
         .to_string()
 });
-const DEFAULT_INSTALLED_EXTENSIONS_COLLECTION_INTERVAL: u64 = 3600;
 
 /// Static configuration params that don't change after startup. These mostly
 /// come from the CLI args, or are derived from them.
@@ -104,7 +103,7 @@ pub struct ComputeNodeParams {
     pub remote_ext_base_url: Option<Url>,
 
     /// Interval for installed extensions collection
-    pub installed_extensions_collection_interval: Arc<AtomicU64>,
+    pub installed_extensions_collection_interval: u64,
 }
 
 /// Compute node info shared across several `compute_ctl` threads.
@@ -127,9 +126,6 @@ pub struct ComputeNode {
     // key: ext_archive_name, value: started download time, download_completed?
     pub ext_download_progress: RwLock<HashMap<String, (DateTime<Utc>, bool)>>,
     pub compute_ctl_config: ComputeCtlConfig,
-
-    /// Handle to the extension stats collection task
-    extension_stats_task: Mutex<Option<tokio::task::JoinHandle<()>>>,
 }
 
 // store some metrics about download size that might impact startup time
@@ -432,7 +428,6 @@ impl ComputeNode {
             state_changed: Condvar::new(),
             ext_download_progress: RwLock::new(HashMap::new()),
             compute_ctl_config: config.compute_ctl_config,
-            extension_stats_task: Mutex::new(None),
         })
     }
 
@@ -520,9 +515,6 @@ impl ComputeNode {
             None
         };
 
-        // Terminate the extension stats collection task
-        this.terminate_extension_stats_task();
-
         // Terminate the vm_monitor so it releases the file watcher on
         // /sys/fs/cgroup/neon-postgres.
         // Note: the vm-monitor only runs on linux because it requires cgroups.
@@ -1679,8 +1671,6 @@ impl ComputeNode {
             tls_config = self.compute_ctl_config.tls.clone();
         }
 
-        self.update_installed_extensions_collection_interval(&spec);
-
         let max_concurrent_connections = self.max_service_connections(compute_state, &spec);
 
         // Merge-apply spec & changes to PostgreSQL state.
@@ -1745,8 +1735,6 @@ impl ComputeNode {
 
         let tls_config = self.tls_config(&spec);
 
-        self.update_installed_extensions_collection_interval(&spec);
-
         if let Some(ref pgbouncer_settings) = spec.pgbouncer_settings {
             info!("tuning pgbouncer");
 
@@ -2351,20 +2339,10 @@ LIMIT 100",
     }
 
     pub fn spawn_extension_stats_task(&self) {
-        // Cancel any existing task
-        if let Some(handle) = self.extension_stats_task.lock().unwrap().take() {
-            handle.abort();
-        }
-
         let conf = self.tokio_conn_conf.clone();
-        let atomic_interval = self.params.installed_extensions_collection_interval.clone();
-        let mut installed_extensions_collection_interval =
-            2 * atomic_interval.load(std::sync::atomic::Ordering::SeqCst);
-        info!(
-            "[NEON_EXT_SPAWN] Spawning background installed extensions worker with Timeout: {}",
-            installed_extensions_collection_interval
-        );
-        let handle = tokio::spawn(async move {
+        let installed_extensions_collection_interval =
+            self.params.installed_extensions_collection_interval;
+        tokio::spawn(async move {
             // An initial sleep is added to ensure that two collections don't happen at the same time.
             // The first collection happens during compute startup.
             tokio::time::sleep(tokio::time::Duration::from_secs(
@@ -2377,48 +2355,8 @@ LIMIT 100",
             loop {
                 interval.tick().await;
                 let _ = installed_extensions(conf.clone()).await;
-                // Acquire a read lock on the compute spec and then update the interval if necessary
-                interval = tokio::time::interval(tokio::time::Duration::from_secs(std::cmp::max(
-                    installed_extensions_collection_interval,
-                    2 * atomic_interval.load(std::sync::atomic::Ordering::SeqCst),
-                )));
-                installed_extensions_collection_interval = interval.period().as_secs();
             }
         });
-
-        // Store the new task handle
-        *self.extension_stats_task.lock().unwrap() = Some(handle);
-    }
-
-    fn terminate_extension_stats_task(&self) {
-        if let Some(handle) = self.extension_stats_task.lock().unwrap().take() {
-            handle.abort();
-        }
-    }
-
-    fn update_installed_extensions_collection_interval(&self, spec: &ComputeSpec) {
-        // Update the interval for collecting installed extensions statistics
-        // If the value is -1, we never suspend so set the value to default collection.
-        // If the value is 0, it means default, we will just continue to use the default.
-        if spec.suspend_timeout_seconds == -1 || spec.suspend_timeout_seconds == 0 {
-            info!(
-                "[NEON_EXT_INT_UPD] Spec Timeout: {}, New Timeout: {}",
-                spec.suspend_timeout_seconds, DEFAULT_INSTALLED_EXTENSIONS_COLLECTION_INTERVAL
-            );
-            self.params.installed_extensions_collection_interval.store(
-                DEFAULT_INSTALLED_EXTENSIONS_COLLECTION_INTERVAL,
-                std::sync::atomic::Ordering::SeqCst,
-            );
-        } else {
-            info!(
-                "[NEON_EXT_INT_UPD] Spec Timeout: {}",
-                spec.suspend_timeout_seconds
-            );
-            self.params.installed_extensions_collection_interval.store(
-                spec.suspend_timeout_seconds as u64,
-                std::sync::atomic::Ordering::SeqCst,
-            );
-        }
     }
 }
 

compute_tools/tests/cluster_spec.json

@@ -3,8 +3,7 @@
 
   "timestamp": "2021-05-23T18:25:43.511Z",
   "operation_uuid": "0f657b36-4b0f-4a2d-9c2e-1dcd615e7d8b",
-  "suspend_timeout_seconds": 3600,
-  
+
   "cluster": {
     "cluster_id": "test-cluster-42",
     "name": "Zenith Test",

control_plane/src/bin/neon_local.rs

@@ -1649,9 +1649,7 @@ async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Res
             // If --safekeepers argument is given, use only the listed
             // safekeeper nodes; otherwise all from the env.
             let safekeepers = parse_safekeepers(&args.safekeepers)?;
-            endpoint
-                .reconfigure(Some(pageservers), None, safekeepers, None)
-                .await?;
+            endpoint.reconfigure(pageservers, None, safekeepers).await?;
         }
         EndpointCmd::Stop(args) => {
             let endpoint_id = &args.endpoint_id;

control_plane/src/endpoint.rs

@@ -780,7 +780,6 @@ impl Endpoint {
                 endpoint_storage_addr: Some(endpoint_storage_addr),
                 endpoint_storage_token: Some(endpoint_storage_token),
                 autoprewarm: false,
-                suspend_timeout_seconds: -1, // Only used in neon_local.
             };
 
             // this strange code is needed to support respec() in tests
@@ -975,11 +974,12 @@ impl Endpoint {
 
     pub async fn reconfigure(
         &self,
-        pageservers: Option<Vec<(PageserverProtocol, Host, u16)>>,
+        pageservers: Vec<(PageserverProtocol, Host, u16)>,
         stripe_size: Option<ShardStripeSize>,
         safekeepers: Option<Vec<NodeId>>,
-        safekeeper_generation: Option<SafekeeperGeneration>,
     ) -> Result<()> {
+        anyhow::ensure!(!pageservers.is_empty(), "no pageservers provided");
+
         let (mut spec, compute_ctl_config) = {
             let config_path = self.endpoint_path().join("config.json");
             let file = std::fs::File::open(config_path)?;
@@ -991,24 +991,16 @@ impl Endpoint {
         let postgresql_conf = self.read_postgresql_conf()?;
         spec.cluster.postgresql_conf = Some(postgresql_conf);
 
-        // If pageservers are not specified, don't change them.
-        if let Some(pageservers) = pageservers {
-            anyhow::ensure!(!pageservers.is_empty(), "no pageservers provided");
-
-            let pageserver_connstr = Self::build_pageserver_connstr(&pageservers);
-            spec.pageserver_connstring = Some(pageserver_connstr);
-            if stripe_size.is_some() {
-                spec.shard_stripe_size = stripe_size.map(|s| s.0 as usize);
-            }
+        let pageserver_connstr = Self::build_pageserver_connstr(&pageservers);
+        spec.pageserver_connstring = Some(pageserver_connstr);
+        if stripe_size.is_some() {
+            spec.shard_stripe_size = stripe_size.map(|s| s.0 as usize);
         }
 
         // If safekeepers are not specified, don't change them.
         if let Some(safekeepers) = safekeepers {
             let safekeeper_connstrings = self.build_safekeepers_connstrs(safekeepers)?;
             spec.safekeeper_connstrings = safekeeper_connstrings;
-            if let Some(g) = safekeeper_generation {
-                spec.safekeepers_generation = Some(g.into_inner());
-            }
         }
 
         let client = reqwest::Client::builder()
@@ -1046,24 +1038,6 @@ impl Endpoint {
         }
     }
 
-    pub async fn reconfigure_pageservers(
-        &self,
-        pageservers: Vec<(PageserverProtocol, Host, u16)>,
-        stripe_size: Option<ShardStripeSize>,
-    ) -> Result<()> {
-        self.reconfigure(Some(pageservers), stripe_size, None, None)
-            .await
-    }
-
-    pub async fn reconfigure_safekeepers(
-        &self,
-        safekeepers: Vec<NodeId>,
-        generation: SafekeeperGeneration,
-    ) -> Result<()> {
-        self.reconfigure(None, None, Some(safekeepers), Some(generation))
-            .await
-    }
-
     pub async fn stop(
         &self,
         mode: EndpointTerminateMode,

docker-compose/compute_wrapper/var/db/postgres/configs/config.json

@@ -4,7 +4,6 @@
 
         "timestamp": "2022-10-12T18:00:00.000Z",
         "operation_uuid": "0f657b36-4b0f-4a2d-9c2e-1dcd615e7d8c",
-        "suspend_timeout_seconds": -1,
 
         "cluster": {
             "cluster_id": "docker_compose",

libs/compute_api/src/spec.rs

@@ -184,11 +184,6 @@ pub struct ComputeSpec {
     /// Download LFC state from endpoint_storage and pass it to Postgres on startup
     #[serde(default)]
     pub autoprewarm: bool,
-
-    /// Suspend timeout in seconds.
-    ///
-    /// We use this value to derive other values, such as the installed extensions metric.
-    pub suspend_timeout_seconds: i64,
 }
 
 /// Feature flag to signal `compute_ctl` to enable certain experimental functionality.

libs/compute_api/tests/cluster_spec.json

@@ -3,7 +3,6 @@
 
     "timestamp": "2021-05-23T18:25:43.511Z",
     "operation_uuid": "0f657b36-4b0f-4a2d-9c2e-1dcd615e7d8b",
-    "suspend_timeout_seconds": 3600,
 
     "cluster": {
         "cluster_id": "test-cluster-42",

pageserver/src/context.rs

@@ -92,6 +92,7 @@
 use std::{sync::Arc, time::Duration};
 
 use once_cell::sync::Lazy;
+use tokio_util::sync::CancellationToken;
 use tracing::warn;
 use utils::{id::TimelineId, shard::TenantShardId};
 
@@ -117,6 +118,7 @@ pub struct RequestContext {
     scope: Scope,
     perf_span: Option<PerfSpan>,
     perf_span_dispatch: Option<Dispatch>,
+    cancel: CancellationToken,
 }
 
 #[derive(Clone)]
@@ -263,6 +265,10 @@ pub struct RequestContextBuilder {
 impl RequestContextBuilder {
     /// A new builder with default settings
     pub fn new(task_kind: TaskKind) -> Self {
+        Self::new_with_cancel(task_kind, CancellationToken::new())
+    }
+    /// A new builder with default settings, with ability to specify the cancellation token
+    pub(crate) fn new_with_cancel(task_kind: TaskKind, cancel: CancellationToken) -> Self {
         Self {
             inner: RequestContext {
                 task_kind,
@@ -273,6 +279,7 @@ impl RequestContextBuilder {
                 scope: Scope::new_global(),
                 perf_span: None,
                 perf_span_dispatch: None,
+                cancel,
             },
         }
     }
@@ -358,11 +365,18 @@ impl RequestContextBuilder {
         self.inner
     }
 
-    pub fn attached_child(self) -> RequestContext {
+    pub fn attached_child(mut self) -> RequestContext {
+        self.inner.cancel = self.inner.cancel.child_token();
         self.inner
     }
 
-    pub fn detached_child(self) -> RequestContext {
+    pub fn detached_child(mut self) -> RequestContext {
+        self.inner.cancel = CancellationToken::new();
+        self.inner
+    }
+
+    pub fn detached_child_with_cancel(mut self, cancel: CancellationToken) -> RequestContext {
+        self.inner.cancel = cancel;
         self.inner
     }
 }
@@ -382,6 +396,7 @@ impl RequestContext {
             scope: self.scope.clone(),
             perf_span: self.perf_span.clone(),
             perf_span_dispatch: self.perf_span_dispatch.clone(),
+            cancel: self.cancel.clone(),
         }
     }
 
@@ -427,6 +442,19 @@ impl RequestContext {
             .detached_child()
     }
 
+    /// Like [`Self::detached_child`], but with the ability to specify the cancellation token
+    pub fn detached_child_with_cancel(
+        &self,
+        task_kind: TaskKind,
+        download_behavior: DownloadBehavior,
+        cancel: CancellationToken,
+    ) -> Self {
+        RequestContextBuilder::from(self)
+            .task_kind(task_kind)
+            .download_behavior(download_behavior)
+            .detached_child_with_cancel(cancel)
+    }
+
     /// Create a child of context `self` for a task that shall not outlive `self`.
     ///
     /// Use this when fanning-out work to other async tasks.
@@ -603,6 +631,10 @@ impl RequestContext {
     pub(crate) fn has_perf_span(&self) -> bool {
         self.perf_span.is_some()
     }
+
+    pub(crate) fn cancellation_token(&self) -> &CancellationToken {
+        &self.cancel
+    }
 }
 
 /// [`Future`] extension trait that allow for creating performance

pageserver/src/tenant/remote_timeline_client.rs

@@ -727,7 +727,7 @@ impl RemoteTimelineClient {
                     reason: "no need for a downloads gauge",
                 },
             );
-            download::download_layer_file(
+            let fut = download::download_layer_file(
                 self.conf,
                 &self.storage_impl,
                 self.tenant_shard_id,
@@ -744,8 +744,11 @@ impl RemoteTimelineClient {
                 RemoteOpFileKind::Layer,
                 RemoteOpKind::Download,
                 Arc::clone(&self.metrics),
-            )
-            .await?
+            );
+            /*tokio::select! {
+                res = fut => res,
+                _ = ctx.cancellation_token().cancelled() => return Err(DownloadError::DownloadCancelled),
+            }*/fut.await?
         };
 
         REMOTE_ONDEMAND_DOWNLOADED_LAYERS.inc();

pageserver/src/tenant/remote_timeline_client/download.rs

@@ -190,7 +190,13 @@ async fn download_object(
         .download(src_path, &DownloadOpts::default(), cancel)
         .await?;
 
-    pausable_failpoint!("before-downloading-layer-stream-pausable");
+    tracing::info!("Starting layer download");
+
+    pausable_failpoint!(
+        "before-downloading-layer-stream-pausable",
+        ctx.cancellation_token()
+    )
+    .map_err(|_| DownloadError::Cancelled)?;
 
     let dst_path = destination_file.path().to_owned();
     let mut buffered = owned_buffers_io::write::BufferedWriter::<IoBufferMut, _>::new(

pageserver/src/tenant/storage_layer/layer.rs

@@ -337,10 +337,14 @@ impl Layer {
                 })
                 .attached_child();
 
-            self.0
+            // good
+            let fut = self.0
                 .get_or_maybe_download(true, &ctx)
-                .maybe_perf_instrument(&ctx, |crnt_perf_context| crnt_perf_context.clone())
-                .await
+                .maybe_perf_instrument(&ctx, |crnt_perf_context| crnt_perf_context.clone());
+            /*tokio::select! {
+                res = fut => res,
+                _ = ctx.cancellation_token().cancelled() => return Err(GetVectoredError::Cancelled),
+            }*/fut.await
                 .map_err(|err| match err {
                     DownloadError::TimelineShutdown | DownloadError::DownloadCancelled => {
                         GetVectoredError::Cancelled
@@ -1041,10 +1045,11 @@ impl LayerInner {
         //
         // if we are cancelled while doing this `stat` the `self.inner` will be uninitialized. a
         // pending eviction will try to evict even upon finding an uninitialized `self.inner`.
-        let needs_download = self
-            .needs_download()
-            .await
-            .map_err(DownloadError::PreStatFailed);
+        let needs_download = tokio::select! {
+            dl = self.needs_download() => dl,
+            _ = ctx.cancellation_token().cancelled() => return Err(DownloadError::DownloadCancelled),
+        }
+        .map_err(DownloadError::PreStatFailed);
 
         scopeguard::ScopeGuard::into_inner(init_cancelled);
 
@@ -1054,8 +1059,12 @@ impl LayerInner {
             // the file is present locally because eviction has not had a chance to run yet
 
             #[cfg(test)]
-            self.failpoint(failpoints::FailpointKind::AfterDeterminingLayerNeedsNoDownload)
-                .await?;
+            //self.failpoint(failpoints::FailpointKind::AfterDeterminingLayerNeedsNoDownload)
+                //.await?;
+            tokio::select! {
+                dl = self.failpoint(failpoints::FailpointKind::AfterDeterminingLayerNeedsNoDownload) => dl?,
+                _ = ctx.cancellation_token().cancelled() => return Err(DownloadError::DownloadCancelled),
+            }
 
             LAYER_IMPL_METRICS.inc_init_needed_no_download();
 
@@ -1092,10 +1101,16 @@ impl LayerInner {
             tracing::info!(%reason, "downloading on-demand");
 
             let init_cancelled = scopeguard::guard((), |_| LAYER_IMPL_METRICS.inc_init_cancelled());
-            let res = self
+            let fut = self
                 .download_init_and_wait(timeline, permit, ctx.attached_child())
                 .maybe_perf_instrument(&ctx, |current_perf_span| current_perf_span.clone())
-                .await?;
+                ;
+            let res = tokio::select! {
+                res = fut => res,
+                _ = ctx.cancellation_token().cancelled() => return Err(DownloadError::DownloadCancelled),
+            }?;
+            // bad
+            //fut.await?;
 
             scopeguard::ScopeGuard::into_inner(init_cancelled);
             Ok(res)

pageserver/src/tenant/timeline.rs

@@ -1345,9 +1345,12 @@ impl Timeline {
                 })
                 .attached_child();
 
-            self.get_vectored_reconstruct_data(query.clone(), reconstruct_state, &ctx)
-                .maybe_perf_instrument(&ctx, |crnt_perf_span| crnt_perf_span.clone())
-                .await
+            let fut = self.get_vectored_reconstruct_data(query.clone(), reconstruct_state, &ctx)
+                .maybe_perf_instrument(&ctx, |crnt_perf_span| crnt_perf_span.clone());
+            /*tokio::select! {
+                res = fut => res,
+                _ = ctx.cancellation_token().cancelled() => return Err(GetVectoredError::Cancelled),
+            }*/fut.await
         };
 
         if let Err(err) = traversal_res {
@@ -4441,7 +4444,7 @@ impl Timeline {
         let mut image_covered_keyspace = KeySpaceRandomAccum::new();
 
         while let Some((layer_to_read, keyspace_to_read, lsn_range)) = fringe.next_layer() {
-            if cancel.is_cancelled() {
+            if cancel.is_cancelled() || ctx.cancellation_token().is_cancelled() {
                 return Err(GetVectoredError::Cancelled);
             }
 

pageserver/src/tenant/timeline/detach_ancestor.rs

@@ -885,7 +885,7 @@ async fn remote_copy(
                 }
                 tracing::info!("Deleting orphan layer file to make way for hard linking");
                 // Delete orphan layer file and try again, to ensure this layer has a well understood source
-                std::fs::remove_file(&adoptee_path)
+                std::fs::remove_file(adopted_path)
                     .map_err(|e| Error::launder(e.into(), Error::Prepare))?;
                 std::fs::hard_link(adopted_path, &adoptee_path)
                     .map_err(|e| Error::launder(e.into(), Error::Prepare))?;

pageserver/src/tenant/timeline/walreceiver/connection_manager.rs

@@ -36,7 +36,7 @@ use utils::postgres_client::{ConnectionConfigArgs, wal_stream_connection_config}
 
 use super::walreceiver_connection::{WalConnectionStatus, WalReceiverError};
 use super::{TaskEvent, TaskHandle, TaskStateUpdate, WalReceiverConf};
-use crate::context::{DownloadBehavior, RequestContext};
+use crate::context::{DownloadBehavior, RequestContext, RequestContextBuilder};
 use crate::metrics::{
     WALRECEIVER_ACTIVE_MANAGERS, WALRECEIVER_BROKER_UPDATES, WALRECEIVER_CANDIDATES_ADDED,
     WALRECEIVER_CANDIDATES_REMOVED, WALRECEIVER_SWITCHES,
@@ -536,16 +536,18 @@ impl ConnectionManagerState {
         let protocol = self.conf.protocol;
         let validate_wal_contiguity = self.conf.validate_wal_contiguity;
         let timeline = Arc::clone(&self.timeline);
-        let ctx = ctx.detached_child(
-            TaskKind::WalReceiverConnectionHandler,
-            DownloadBehavior::Download,
-        );
+
+        let ctx_builder = RequestContextBuilder::from(ctx)
+            .task_kind(TaskKind::WalReceiverConnectionHandler)
+            .download_behavior(DownloadBehavior::Download);
 
         let span = info_span!("connection", %node_id);
         let connection_handle = self.spawn(move |events_sender, cancellation| {
             async move {
                 debug_assert_current_span_has_tenant_and_timeline_id();
 
+                let ctx = ctx_builder.detached_child_with_cancel(cancellation.clone());
+
                 let res = super::walreceiver_connection::handle_walreceiver_connection(
                     timeline,
                     protocol,

pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs

@@ -275,20 +275,12 @@ pub(super) async fn handle_walreceiver_connection(
     let copy_stream = replication_client.copy_both_simple(&query).await?;
     let mut physical_stream = pin!(ReplicationStream::new(copy_stream));
 
-    let walingest_future = WalIngest::new(timeline.as_ref(), startpoint, &ctx);
-    let walingest_res = select! {
-        walingest_res = walingest_future => walingest_res,
-        _ = cancellation.cancelled() => {
-            // We are doing reads in WalIngest::new, and those can hang as they come from the network.
-            // Timeline cancellation hits the walreceiver cancellation token before it hits the timeline global one.
-            debug!("Connection cancelled");
-            return Err(WalReceiverError::Cancelled);
-        },
-    };
-    let mut walingest = walingest_res.map_err(|e| match e.kind {
-        crate::walingest::WalIngestErrorKind::Cancelled => WalReceiverError::Cancelled,
-        _ => WalReceiverError::Other(e.into()),
-    })?;
+    let mut walingest = WalIngest::new(timeline.as_ref(), startpoint, &ctx)
+        .await
+        .map_err(|e| match e.kind {
+            crate::walingest::WalIngestErrorKind::Cancelled => WalReceiverError::Cancelled,
+            _ => WalReceiverError::Other(e.into()),
+        })?;
 
     let (format, compression) = match protocol {
         PostgresClientProtocol::Interpreted {

pageserver/src/virtual_file/owned_buffers_io/write/flush.rs

@@ -299,7 +299,7 @@ where
             //
             let mut request_storage = Some(request);
             for attempt in 1.. {
-                if self.cancel.is_cancelled() {
+                if self.cancel.is_cancelled() || self.ctx.cancellation_token().is_cancelled() {
                     return Err(FlushTaskError::Cancelled);
                 }
                 let result = async {

storage_controller/src/compute_hook.rs

@@ -14,12 +14,11 @@ use pageserver_api::config::DEFAULT_GRPC_LISTEN_PORT;
 use pageserver_api::controller_api::AvailabilityZone;
 use pageserver_api::shard::{ShardCount, ShardNumber, ShardStripeSize, TenantShardId};
 use postgres_connection::parse_host_port;
-use safekeeper_api::membership::SafekeeperGeneration;
 use serde::{Deserialize, Serialize};
 use tokio_util::sync::CancellationToken;
 use tracing::{Instrument, info_span};
 use utils::backoff::{self};
-use utils::id::{NodeId, TenantId, TenantTimelineId, TimelineId};
+use utils::id::{NodeId, TenantId};
 
 use crate::service::Config;
 
@@ -37,7 +36,7 @@ struct UnshardedComputeHookTenant {
     preferred_az: Option<AvailabilityZone>,
 
     // Must hold this lock to send a notification.
-    send_lock: Arc<tokio::sync::Mutex<Option<ComputeRemoteTenantState>>>,
+    send_lock: Arc<tokio::sync::Mutex<Option<ComputeRemoteState>>>,
 }
 struct ShardedComputeHookTenant {
     stripe_size: ShardStripeSize,
@@ -50,7 +49,7 @@ struct ShardedComputeHookTenant {
     // Must hold this lock to send a notification.  The contents represent
     // the last successfully sent notification, and are used to coalesce multiple
     // updates by only sending when there is a chance since our last successful send.
-    send_lock: Arc<tokio::sync::Mutex<Option<ComputeRemoteTenantState>>>,
+    send_lock: Arc<tokio::sync::Mutex<Option<ComputeRemoteState>>>,
 }
 
 /// Represents our knowledge of the compute's state: we can update this when we get a
@@ -58,9 +57,9 @@ struct ShardedComputeHookTenant {
 ///
 /// Should be wrapped in an Option<>, as we cannot always know the remote state.
 #[derive(PartialEq, Eq, Debug)]
-struct ComputeRemoteState<R> {
+struct ComputeRemoteState {
     // The request body which was acked by the compute
-    request: R,
+    request: ComputeHookNotifyRequest,
 
     // Whether the cplane indicated that the state was applied to running computes, or just
     // persisted.  In the Neon control plane, this is the difference between a 423 response (meaning
@@ -68,36 +67,6 @@ struct ComputeRemoteState<R> {
     applied: bool,
 }
 
-type ComputeRemoteTenantState = ComputeRemoteState<NotifyAttachRequest>;
-type ComputeRemoteTimelineState = ComputeRemoteState<NotifySafekeepersRequest>;
-
-/// The trait which define the handler-specific types and methods.
-/// We have two implementations of this trait so far:
-/// - [`ComputeHookTenant`] for tenant attach notifications ("/notify-attach")
-/// - [`ComputeHookTimeline`] for safekeeper change notifications ("/notify-safekeepers")
-trait ApiMethod {
-    /// Type of the key which identifies the resource.
-    /// It's either TenantId for tenant attach notifications,
-    /// or TenantTimelineId for safekeeper change notifications.
-    type Key: std::cmp::Eq + std::hash::Hash + Clone;
-
-    type Request: serde::Serialize + std::fmt::Debug;
-
-    const API_PATH: &'static str;
-
-    fn maybe_send(
-        &self,
-        key: Self::Key,
-        lock: Option<tokio::sync::OwnedMutexGuard<Option<ComputeRemoteState<Self::Request>>>>,
-    ) -> MaybeSendResult<Self::Request, Self::Key>;
-
-    async fn notify_local(
-        env: &LocalEnv,
-        cplane: &ComputeControlPlane,
-        req: &Self::Request,
-    ) -> Result<(), NotifyError>;
-}
-
 enum ComputeHookTenant {
     Unsharded(UnshardedComputeHookTenant),
     Sharded(ShardedComputeHookTenant),
@@ -128,7 +97,7 @@ impl ComputeHookTenant {
         }
     }
 
-    fn get_send_lock(&self) -> &Arc<tokio::sync::Mutex<Option<ComputeRemoteTenantState>>> {
+    fn get_send_lock(&self) -> &Arc<tokio::sync::Mutex<Option<ComputeRemoteState>>> {
         match self {
             Self::Unsharded(unsharded_tenant) => &unsharded_tenant.send_lock,
             Self::Sharded(sharded_tenant) => &sharded_tenant.send_lock,
@@ -222,136 +191,19 @@ impl ComputeHookTenant {
     }
 }
 
-/// The state of a timeline we need to notify the compute about.
-struct ComputeHookTimeline {
-    generation: SafekeeperGeneration,
-    safekeepers: Vec<SafekeeperInfo>,
-
-    send_lock: Arc<tokio::sync::Mutex<Option<ComputeRemoteTimelineState>>>,
-}
-
-impl ComputeHookTimeline {
-    /// Construct a new ComputeHookTimeline with the given safekeepers and generation.
-    fn new(generation: SafekeeperGeneration, safekeepers: Vec<SafekeeperInfo>) -> Self {
-        Self {
-            generation,
-            safekeepers,
-            send_lock: Arc::default(),
-        }
-    }
-
-    /// Update the state with a new SafekeepersUpdate.
-    /// Noop if the update generation is not greater than the current generation.
-    fn update(&mut self, sk_update: SafekeepersUpdate) {
-        if sk_update.generation > self.generation {
-            self.generation = sk_update.generation;
-            self.safekeepers = sk_update.safekeepers;
-        }
-    }
-}
-
-impl ApiMethod for ComputeHookTimeline {
-    type Key = TenantTimelineId;
-    type Request = NotifySafekeepersRequest;
-
-    const API_PATH: &'static str = "notify-safekeepers";
-
-    fn maybe_send(
-        &self,
-        ttid: TenantTimelineId,
-        lock: Option<tokio::sync::OwnedMutexGuard<Option<ComputeRemoteTimelineState>>>,
-    ) -> MaybeSendNotifySafekeepersResult {
-        let locked = match lock {
-            Some(already_locked) => already_locked,
-            None => {
-                // Lock order: this _must_ be only a try_lock, because we are called inside of the [`ComputeHook::timelines`] lock.
-                let Ok(locked) = self.send_lock.clone().try_lock_owned() else {
-                    return MaybeSendResult::AwaitLock((ttid, self.send_lock.clone()));
-                };
-                locked
-            }
-        };
-
-        if locked
-            .as_ref()
-            .is_some_and(|s| s.request.generation >= self.generation)
-        {
-            return MaybeSendResult::Noop;
-        }
-
-        MaybeSendResult::Transmit((
-            NotifySafekeepersRequest {
-                tenant_id: ttid.tenant_id,
-                timeline_id: ttid.timeline_id,
-                generation: self.generation,
-                safekeepers: self.safekeepers.clone(),
-            },
-            locked,
-        ))
-    }
-
-    async fn notify_local(
-        _env: &LocalEnv,
-        cplane: &ComputeControlPlane,
-        req: &NotifySafekeepersRequest,
-    ) -> Result<(), NotifyError> {
-        let NotifySafekeepersRequest {
-            tenant_id,
-            timeline_id,
-            generation,
-            safekeepers,
-        } = req;
-
-        for (endpoint_name, endpoint) in &cplane.endpoints {
-            if endpoint.tenant_id == *tenant_id
-                && endpoint.timeline_id == *timeline_id
-                && endpoint.status() == EndpointStatus::Running
-            {
-                tracing::info!("Reconfiguring safekeepers for endpoint {endpoint_name}");
-
-                let safekeepers = safekeepers.iter().map(|sk| sk.id).collect::<Vec<_>>();
-
-                endpoint
-                    .reconfigure_safekeepers(safekeepers, *generation)
-                    .await
-                    .map_err(NotifyError::NeonLocal)?;
-            }
-        }
-
-        Ok(())
-    }
-}
-
 #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
-struct NotifyAttachRequestShard {
+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, Eq, PartialEq)]
-struct NotifyAttachRequest {
+struct ComputeHookNotifyRequest {
     tenant_id: TenantId,
     preferred_az: Option<String>,
     stripe_size: Option<ShardStripeSize>,
-    shards: Vec<NotifyAttachRequestShard>,
-}
-
-#[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Clone)]
-pub(crate) struct SafekeeperInfo {
-    pub id: NodeId,
-    /// Hostname of the safekeeper.
-    /// It exists for better debuggability. Might be missing.
-    /// Should not be used for anything else.
-    pub hostname: Option<String>,
-}
-
-#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
-struct NotifySafekeepersRequest {
-    tenant_id: TenantId,
-    timeline_id: TimelineId,
-    generation: SafekeeperGeneration,
-    safekeepers: Vec<SafekeeperInfo>,
+    shards: Vec<ComputeHookNotifyRequestShard>,
 }
 
 /// Error type for attempts to call into the control plane compute notification hook
@@ -383,50 +235,42 @@ pub(crate) enum NotifyError {
     NeonLocal(anyhow::Error),
 }
 
-enum MaybeSendResult<R, K> {
+enum MaybeSendResult {
     // Please send this request while holding the lock, and if you succeed then write
     // the request into the lock.
     Transmit(
         (
-            R,
-            tokio::sync::OwnedMutexGuard<Option<ComputeRemoteState<R>>>,
+            ComputeHookNotifyRequest,
+            tokio::sync::OwnedMutexGuard<Option<ComputeRemoteState>>,
         ),
     ),
     // Something requires sending, but you must wait for a current sender then call again
-    AwaitLock((K, Arc<tokio::sync::Mutex<Option<ComputeRemoteState<R>>>>)),
+    AwaitLock(Arc<tokio::sync::Mutex<Option<ComputeRemoteState>>>),
     // Nothing requires sending
     Noop,
 }
 
-type MaybeSendNotifyAttachResult = MaybeSendResult<NotifyAttachRequest, TenantId>;
-type MaybeSendNotifySafekeepersResult = MaybeSendResult<NotifySafekeepersRequest, TenantTimelineId>;
-
-impl ApiMethod for ComputeHookTenant {
-    type Key = TenantId;
-    type Request = NotifyAttachRequest;
-
-    const API_PATH: &'static str = "notify-attach";
-
+impl ComputeHookTenant {
     fn maybe_send(
         &self,
         tenant_id: TenantId,
-        lock: Option<tokio::sync::OwnedMutexGuard<Option<ComputeRemoteTenantState>>>,
-    ) -> MaybeSendNotifyAttachResult {
+        lock: Option<tokio::sync::OwnedMutexGuard<Option<ComputeRemoteState>>>,
+    ) -> MaybeSendResult {
         let locked = match lock {
             Some(already_locked) => already_locked,
             None => {
-                // Lock order: this _must_ be only a try_lock, because we are called inside of the [`ComputeHook::tenants`] lock.
+                // Lock order: this _must_ be only a try_lock, because we are called inside of the [`ComputeHook::state`] lock.
                 let Ok(locked) = self.get_send_lock().clone().try_lock_owned() else {
-                    return MaybeSendResult::AwaitLock((tenant_id, self.get_send_lock().clone()));
+                    return MaybeSendResult::AwaitLock(self.get_send_lock().clone());
                 };
                 locked
             }
         };
 
         let request = match self {
-            Self::Unsharded(unsharded_tenant) => Some(NotifyAttachRequest {
+            Self::Unsharded(unsharded_tenant) => Some(ComputeHookNotifyRequest {
                 tenant_id,
-                shards: vec![NotifyAttachRequestShard {
+                shards: vec![ComputeHookNotifyRequestShard {
                     shard_number: ShardNumber(0),
                     node_id: unsharded_tenant.node_id,
                 }],
@@ -439,12 +283,12 @@ impl ApiMethod for ComputeHookTenant {
             Self::Sharded(sharded_tenant)
                 if sharded_tenant.shards.len() == sharded_tenant.shard_count.count() as usize =>
             {
-                Some(NotifyAttachRequest {
+                Some(ComputeHookNotifyRequest {
                     tenant_id,
                     shards: sharded_tenant
                         .shards
                         .iter()
-                        .map(|(shard_number, node_id)| NotifyAttachRequestShard {
+                        .map(|(shard_number, node_id)| ComputeHookNotifyRequestShard {
                             shard_number: *shard_number,
                             node_id: *node_id,
                         })
@@ -489,22 +333,98 @@ impl ApiMethod for ComputeHookTenant {
             }
         }
     }
+}
 
-    async fn notify_local(
-        env: &LocalEnv,
-        cplane: &ComputeControlPlane,
-        req: &NotifyAttachRequest,
+/// 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: std::sync::Mutex<HashMap<TenantId, ComputeHookTenant>>,
+    authorization_header: Option<String>,
+
+    // Concurrency limiter, so that we do not overload the cloud control plane when updating
+    // large numbers of tenants (e.g. when failing over after a node failure)
+    api_concurrency: tokio::sync::Semaphore,
+
+    // This lock is only used in testing enviroments, to serialize calls into neon_lock
+    neon_local_lock: tokio::sync::Mutex<()>,
+
+    // We share a client across all notifications to enable connection re-use etc when
+    // sending large numbers of notifications
+    client: reqwest::Client,
+}
+
+/// Callers may give us a list of these when asking us to send a bulk batch
+/// of notifications in the background.  This is a 'notification' in the sense of
+/// other code notifying us of a shard's status, rather than being the final notification
+/// that we send upwards to the control plane for the whole tenant.
+pub(crate) struct ShardUpdate<'a> {
+    pub(crate) tenant_shard_id: TenantShardId,
+    pub(crate) node_id: NodeId,
+    pub(crate) stripe_size: ShardStripeSize,
+    pub(crate) preferred_az: Option<Cow<'a, AvailabilityZone>>,
+}
+
+impl ComputeHook {
+    pub(super) fn new(config: Config) -> anyhow::Result<Self> {
+        let authorization_header = config
+            .control_plane_jwt_token
+            .clone()
+            .map(|jwt| format!("Bearer {jwt}"));
+
+        let mut client = reqwest::ClientBuilder::new().timeout(NOTIFY_REQUEST_TIMEOUT);
+        for cert in &config.ssl_ca_certs {
+            client = client.add_root_certificate(cert.clone());
+        }
+        let client = client
+            .build()
+            .context("Failed to build http client for compute hook")?;
+
+        Ok(Self {
+            state: Default::default(),
+            config,
+            authorization_header,
+            neon_local_lock: Default::default(),
+            api_concurrency: tokio::sync::Semaphore::new(API_CONCURRENCY),
+            client,
+        })
+    }
+
+    /// For test environments: use neon_local's LocalEnv to update compute
+    async fn do_notify_local(
+        &self,
+        reconfigure_request: &ComputeHookNotifyRequest,
     ) -> Result<(), NotifyError> {
-        let NotifyAttachRequest {
+        // neon_local updates are not safe to call concurrently, use a lock to serialize
+        // all calls to this function
+        let _locked = self.neon_local_lock.lock().await;
+
+        let Some(repo_dir) = self.config.neon_local_repo_dir.as_deref() else {
+            tracing::warn!(
+                "neon_local_repo_dir not set, likely a bug in neon_local; skipping compute update"
+            );
+            return Ok(());
+        };
+        let env = match LocalEnv::load_config(repo_dir) {
+            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,
             preferred_az: _preferred_az,
-        } = req;
+        } = reconfigure_request;
 
         for (endpoint_name, endpoint) in &cplane.endpoints {
             if endpoint.tenant_id == *tenant_id && endpoint.status() == EndpointStatus::Running {
-                tracing::info!("Reconfiguring pageservers for endpoint {endpoint_name}");
+                tracing::info!("Reconfiguring endpoint {endpoint_name}");
 
                 let pageservers = shards
                     .iter()
@@ -526,7 +446,7 @@ impl ApiMethod for ComputeHookTenant {
                     .collect::<Vec<_>>();
 
                 endpoint
-                    .reconfigure_pageservers(pageservers, *stripe_size)
+                    .reconfigure(pageservers, *stripe_size, None)
                     .await
                     .map_err(NotifyError::NeonLocal)?;
             }
@@ -534,102 +454,11 @@ impl ApiMethod for ComputeHookTenant {
 
         Ok(())
     }
-}
 
-/// The compute hook is a destination for notifications about changes to tenant:pageserver
-/// mapping.  It aggregates updates for the shards in a tenant, and when appropriate reconfigures
-/// the compute connection string.
-pub(super) struct ComputeHook {
-    config: Config,
-    tenants: std::sync::Mutex<HashMap<TenantId, ComputeHookTenant>>,
-    timelines: std::sync::Mutex<HashMap<TenantTimelineId, ComputeHookTimeline>>,
-    authorization_header: Option<String>,
-
-    // Concurrency limiter, so that we do not overload the cloud control plane when updating
-    // large numbers of tenants (e.g. when failing over after a node failure)
-    api_concurrency: tokio::sync::Semaphore,
-
-    // This lock is only used in testing enviroments, to serialize calls into neon_local
-    neon_local_lock: tokio::sync::Mutex<()>,
-
-    // We share a client across all notifications to enable connection re-use etc when
-    // sending large numbers of notifications
-    client: reqwest::Client,
-}
-
-/// Callers may give us a list of these when asking us to send a bulk batch
-/// of notifications in the background.  This is a 'notification' in the sense of
-/// other code notifying us of a shard's status, rather than being the final notification
-/// that we send upwards to the control plane for the whole tenant.
-pub(crate) struct ShardUpdate<'a> {
-    pub(crate) tenant_shard_id: TenantShardId,
-    pub(crate) node_id: NodeId,
-    pub(crate) stripe_size: ShardStripeSize,
-    pub(crate) preferred_az: Option<Cow<'a, AvailabilityZone>>,
-}
-
-pub(crate) struct SafekeepersUpdate {
-    pub(crate) tenant_id: TenantId,
-    pub(crate) timeline_id: TimelineId,
-    pub(crate) generation: SafekeeperGeneration,
-    pub(crate) safekeepers: Vec<SafekeeperInfo>,
-}
-
-impl ComputeHook {
-    pub(super) fn new(config: Config) -> anyhow::Result<Self> {
-        let authorization_header = config
-            .control_plane_jwt_token
-            .clone()
-            .map(|jwt| format!("Bearer {jwt}"));
-
-        let mut client = reqwest::ClientBuilder::new().timeout(NOTIFY_REQUEST_TIMEOUT);
-        for cert in &config.ssl_ca_certs {
-            client = client.add_root_certificate(cert.clone());
-        }
-        let client = client
-            .build()
-            .context("Failed to build http client for compute hook")?;
-
-        Ok(Self {
-            tenants: Default::default(),
-            timelines: Default::default(),
-            config,
-            authorization_header,
-            neon_local_lock: Default::default(),
-            api_concurrency: tokio::sync::Semaphore::new(API_CONCURRENCY),
-            client,
-        })
-    }
-
-    /// For test environments: use neon_local's LocalEnv to update compute
-    async fn do_notify_local<M: ApiMethod>(&self, req: &M::Request) -> Result<(), NotifyError> {
-        // neon_local updates are not safe to call concurrently, use a lock to serialize
-        // all calls to this function
-        let _locked = self.neon_local_lock.lock().await;
-
-        let Some(repo_dir) = self.config.neon_local_repo_dir.as_deref() else {
-            tracing::warn!(
-                "neon_local_repo_dir not set, likely a bug in neon_local; skipping compute update"
-            );
-            return Ok(());
-        };
-        let env = match LocalEnv::load_config(repo_dir) {
-            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");
-
-        M::notify_local(&env, &cplane, req).await
-    }
-
-    async fn do_notify_iteration<Req: serde::Serialize + std::fmt::Debug>(
+    async fn do_notify_iteration(
         &self,
         url: &String,
-        reconfigure_request: &Req,
+        reconfigure_request: &ComputeHookNotifyRequest,
         cancel: &CancellationToken,
     ) -> Result<(), NotifyError> {
         let req = self.client.request(reqwest::Method::PUT, url);
@@ -651,7 +480,9 @@ impl ComputeHook {
         };
 
         // Treat all 2xx responses as success
-        if response.status().is_success() {
+        if response.status() >= reqwest::StatusCode::OK
+            && response.status() < reqwest::StatusCode::MULTIPLE_CHOICES
+        {
             if response.status() != reqwest::StatusCode::OK {
                 // Non-200 2xx response: it doesn't make sense to retry, but this is unexpected, so
                 // log a warning.
@@ -702,10 +533,10 @@ impl ComputeHook {
         }
     }
 
-    async fn do_notify<R: serde::Serialize + std::fmt::Debug>(
+    async fn do_notify(
         &self,
         url: &String,
-        reconfigure_request: &R,
+        reconfigure_request: &ComputeHookNotifyRequest,
         cancel: &CancellationToken,
     ) -> Result<(), NotifyError> {
         // We hold these semaphore units across all retries, rather than only across each
@@ -737,13 +568,13 @@ impl ComputeHook {
     }
 
     /// Synchronous phase: update the per-tenant state for the next intended notification
-    fn notify_attach_prepare(&self, shard_update: ShardUpdate) -> MaybeSendNotifyAttachResult {
-        let mut tenants_locked = self.tenants.lock().unwrap();
+    fn notify_prepare(&self, shard_update: ShardUpdate) -> MaybeSendResult {
+        let mut state_locked = self.state.lock().unwrap();
 
         use std::collections::hash_map::Entry;
         let tenant_shard_id = shard_update.tenant_shard_id;
 
-        let tenant = match tenants_locked.entry(tenant_shard_id.tenant_id) {
+        let tenant = match state_locked.entry(tenant_shard_id.tenant_id) {
             Entry::Vacant(e) => {
                 let ShardUpdate {
                     tenant_shard_id,
@@ -767,37 +598,10 @@ impl ComputeHook {
         tenant.maybe_send(tenant_shard_id.tenant_id, None)
     }
 
-    fn notify_safekeepers_prepare(
+    async fn notify_execute(
         &self,
-        safekeepers_update: SafekeepersUpdate,
-    ) -> MaybeSendNotifySafekeepersResult {
-        let mut timelines_locked = self.timelines.lock().unwrap();
-
-        let ttid = TenantTimelineId {
-            tenant_id: safekeepers_update.tenant_id,
-            timeline_id: safekeepers_update.timeline_id,
-        };
-
-        use std::collections::hash_map::Entry;
-        let timeline = match timelines_locked.entry(ttid) {
-            Entry::Vacant(e) => e.insert(ComputeHookTimeline::new(
-                safekeepers_update.generation,
-                safekeepers_update.safekeepers,
-            )),
-            Entry::Occupied(e) => {
-                let timeline = e.into_mut();
-                timeline.update(safekeepers_update);
-                timeline
-            }
-        };
-
-        timeline.maybe_send(ttid, None)
-    }
-
-    async fn notify_execute<M: ApiMethod>(
-        &self,
-        state: &std::sync::Mutex<HashMap<M::Key, M>>,
-        maybe_send_result: MaybeSendResult<M::Request, M::Key>,
+        maybe_send_result: MaybeSendResult,
+        tenant_shard_id: TenantShardId,
         cancel: &CancellationToken,
     ) -> Result<(), NotifyError> {
         // Process result: we may get an update to send, or we may have to wait for a lock
@@ -806,7 +610,7 @@ impl ComputeHook {
             MaybeSendResult::Noop => {
                 return Ok(());
             }
-            MaybeSendResult::AwaitLock((key, send_lock)) => {
+            MaybeSendResult::AwaitLock(send_lock) => {
                 let send_locked = tokio::select! {
                     guard = send_lock.lock_owned() => {guard},
                     _ = cancel.cancelled() => {
@@ -817,11 +621,11 @@ impl ComputeHook {
                 // Lock order: maybe_send is called within the `[Self::state]` lock, and takes the send lock, but here
                 // we have acquired the send lock and take `[Self::state]` lock.  This is safe because maybe_send only uses
                 // try_lock.
-                let state_locked = state.lock().unwrap();
-                let Some(resource_state) = state_locked.get(&key) else {
+                let state_locked = self.state.lock().unwrap();
+                let Some(tenant) = state_locked.get(&tenant_shard_id.tenant_id) else {
                     return Ok(());
                 };
-                match resource_state.maybe_send(key, Some(send_locked)) {
+                match tenant.maybe_send(tenant_shard_id.tenant_id, Some(send_locked)) {
                     MaybeSendResult::AwaitLock(_) => {
                         unreachable!("We supplied lock guard")
                     }
@@ -840,18 +644,14 @@ impl ComputeHook {
                     .control_plane_url
                     .as_ref()
                     .map(|control_plane_url| {
-                        format!(
-                            "{}/{}",
-                            control_plane_url.trim_end_matches('/'),
-                            M::API_PATH
-                        )
+                        format!("{}/notify-attach", control_plane_url.trim_end_matches('/'))
                     });
 
             // We validate this at startup
             let notify_url = compute_hook_url.as_ref().unwrap();
             self.do_notify(notify_url, &request, cancel).await
         } else {
-            self.do_notify_local::<M>(&request).await.map_err(|e| {
+            self.do_notify_local(&request).await.map_err(|e| {
                 // This path is for testing only, so munge the error into our prod-style error type.
                 tracing::error!("neon_local notification hook failed: {e}");
                 NotifyError::Fatal(StatusCode::INTERNAL_SERVER_ERROR)
@@ -887,7 +687,7 @@ impl ComputeHook {
     /// Infallible synchronous fire-and-forget version of notify(), that sends its results to
     /// a channel.  Something should consume the channel and arrange to try notifying again
     /// if something failed.
-    pub(super) fn notify_attach_background(
+    pub(super) fn notify_background(
         self: &Arc<Self>,
         notifications: Vec<ShardUpdate>,
         result_tx: tokio::sync::mpsc::Sender<Result<(), (TenantShardId, NotifyError)>>,
@@ -896,7 +696,7 @@ impl ComputeHook {
         let mut maybe_sends = Vec::new();
         for shard_update in notifications {
             let tenant_shard_id = shard_update.tenant_shard_id;
-            let maybe_send_result = self.notify_attach_prepare(shard_update);
+            let maybe_send_result = self.notify_prepare(shard_update);
             maybe_sends.push((tenant_shard_id, maybe_send_result))
         }
 
@@ -915,10 +715,10 @@ impl ComputeHook {
 
                     async move {
                         this
-                            .notify_execute(&this.tenants, maybe_send_result, &cancel)
+                            .notify_execute(maybe_send_result, tenant_shard_id, &cancel)
                             .await.map_err(|e| (tenant_shard_id, e))
                     }.instrument(info_span!(
-                        "notify_attach_background", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug()
+                        "notify_background", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug()
                     ))
                 })
                 .buffered(API_CONCURRENCY);
@@ -961,23 +761,14 @@ impl ComputeHook {
     /// 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=%shard_update.tenant_shard_id.tenant_id, shard_id=%shard_update.tenant_shard_id.shard_slug(), node_id))]
-    pub(super) async fn notify_attach<'a>(
+    pub(super) async fn notify<'a>(
         &self,
         shard_update: ShardUpdate<'a>,
         cancel: &CancellationToken,
     ) -> Result<(), NotifyError> {
-        let maybe_send_result = self.notify_attach_prepare(shard_update);
-        self.notify_execute(&self.tenants, maybe_send_result, cancel)
-            .await
-    }
-
-    pub(super) async fn notify_safekeepers(
-        &self,
-        safekeepers_update: SafekeepersUpdate,
-        cancel: &CancellationToken,
-    ) -> Result<(), NotifyError> {
-        let maybe_send_result = self.notify_safekeepers_prepare(safekeepers_update);
-        self.notify_execute(&self.timelines, maybe_send_result, cancel)
+        let tenant_shard_id = shard_update.tenant_shard_id;
+        let maybe_send_result = self.notify_prepare(shard_update);
+        self.notify_execute(maybe_send_result, tenant_shard_id, cancel)
             .await
     }
 
@@ -993,8 +784,8 @@ impl ComputeHook {
     ) {
         use std::collections::hash_map::Entry;
 
-        let mut tenants_locked = self.tenants.lock().unwrap();
-        match tenants_locked.entry(tenant_shard_id.tenant_id) {
+        let mut state_locked = self.state.lock().unwrap();
+        match state_locked.entry(tenant_shard_id.tenant_id) {
             Entry::Vacant(_) => {
                 // This is a valid but niche case, where the tenant was previously attached
                 // as a Secondary location and then detached, so has no previously notified

storage_controller/src/reconciler.rs

@@ -65,7 +65,7 @@ pub(super) struct Reconciler {
     pub(crate) compute_hook: Arc<ComputeHook>,
 
     /// To avoid stalling if the cloud control plane is unavailable, we may proceed
-    /// past failures in [`ComputeHook::notify_attach`], but we _must_ remember that we failed
+    /// past failures in [`ComputeHook::notify`], but we _must_ remember that we failed
     /// so that we can set [`crate::tenant_shard::TenantShard::pending_compute_notification`] to ensure a later retry.
     pub(crate) compute_notify_failure: bool,
 
@@ -1023,7 +1023,7 @@ impl Reconciler {
         if let Some(node) = &self.intent.attached {
             let result = self
                 .compute_hook
-                .notify_attach(
+                .notify(
                     compute_hook::ShardUpdate {
                         tenant_shard_id: self.tenant_shard_id,
                         node_id: node.get_id(),

storage_controller/src/service.rs

@@ -878,18 +878,18 @@ impl Service {
         // Emit compute hook notifications for all tenants which are already stably attached.  Other tenants
         // will emit compute hook notifications when they reconcile.
         //
-        // Ordering: our calls to notify_attach_background synchronously establish a relative order for these notifications vs. any later
+        // Ordering: our calls to notify_background synchronously establish a relative order for these notifications vs. any later
         // calls into the ComputeHook for the same tenant: we can leave these to run to completion in the background and any later
         // calls will be correctly ordered wrt these.
         //
-        // Concurrency: we call notify_attach_background for all tenants, which will create O(N) tokio tasks, but almost all of them
+        // Concurrency: we call notify_background for all tenants, which will create O(N) tokio tasks, but almost all of them
         // will just wait on the ComputeHook::API_CONCURRENCY semaphore immediately, so very cheap until they get that semaphore
         // unit and start doing I/O.
         tracing::info!(
             "Sending {} compute notifications",
             compute_notifications.len()
         );
-        self.compute_hook.notify_attach_background(
+        self.compute_hook.notify_background(
             compute_notifications,
             bg_compute_notify_result_tx.clone(),
             &self.cancel,
@@ -6281,7 +6281,7 @@ impl Service {
         for (child_id, child_ps, stripe_size) in child_locations {
             if let Err(e) = self
                 .compute_hook
-                .notify_attach(
+                .notify(
                     compute_hook::ShardUpdate {
                         tenant_shard_id: child_id,
                         node_id: child_ps,

storage_controller/src/service/safekeeper_service.rs

@@ -4,7 +4,6 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use super::safekeeper_reconciler::ScheduleRequest;
-use crate::compute_hook;
 use crate::heartbeater::SafekeeperState;
 use crate::id_lock_map::trace_shared_lock;
 use crate::metrics;
@@ -1199,11 +1198,7 @@ impl Service {
         // 4. Call PUT configuration on safekeepers from the current set,
         // delivering them joint_conf.
 
-        // Notify cplane/compute about the membership change BEFORE changing the membership on safekeepers.
-        // This way the compute will know about new safekeepers from joint_config before we require to
-        // collect a quorum from them.
-        self.cplane_notify_safekeepers(tenant_id, timeline_id, &joint_config)
-            .await?;
+        // TODO(diko): need to notify cplane with an updated set of safekeepers.
 
         let results = self
             .tenant_timeline_set_membership_quorum(
@@ -1310,55 +1305,8 @@ impl Service {
         )
         .await?;
 
-        // Notify cplane/compute about the membership change AFTER changing the membership on safekeepers.
-        // This way the compute will stop talking to excluded safekeepers only after we stop requiring to
-        // collect a quorum from them.
-        self.cplane_notify_safekeepers(tenant_id, timeline_id, &new_conf)
-            .await?;
+        // TODO(diko): need to notify cplane with an updated set of safekeepers.
 
         Ok(())
     }
-
-    /// Notify cplane about safekeeper membership change.
-    /// The cplane will receive a joint set of safekeepers as a safekeeper list.
-    async fn cplane_notify_safekeepers(
-        &self,
-        tenant_id: TenantId,
-        timeline_id: TimelineId,
-        mconf: &membership::Configuration,
-    ) -> Result<(), ApiError> {
-        let mut safekeepers = Vec::new();
-        let mut ids: HashSet<_> = HashSet::new();
-
-        for member in mconf
-            .members
-            .m
-            .iter()
-            .chain(mconf.new_members.iter().flat_map(|m| m.m.iter()))
-        {
-            if ids.insert(member.id) {
-                safekeepers.push(compute_hook::SafekeeperInfo {
-                    id: member.id,
-                    hostname: Some(member.host.clone()),
-                });
-            }
-        }
-
-        self.compute_hook
-            .notify_safekeepers(
-                compute_hook::SafekeepersUpdate {
-                    tenant_id,
-                    timeline_id,
-                    generation: mconf.generation,
-                    safekeepers,
-                },
-                &self.cancel,
-            )
-            .await
-            .map_err(|err| {
-                ApiError::InternalServerError(anyhow::anyhow!(
-                    "failed to notify cplane about safekeeper membership change: {err}"
-                ))
-            })
-    }
 }

test_runner/regress/test_safekeeper_migration.py

@@ -32,13 +32,10 @@ def test_safekeeper_migration_simple(neon_env_builder: NeonEnvBuilder):
     )
 
     ep = env.endpoints.create("main", tenant_id=env.initial_tenant)
-
-    mconf = env.storage_controller.timeline_locate(env.initial_tenant, env.initial_timeline)
-    assert mconf["new_sk_set"] is None
-    assert len(mconf["sk_set"]) == 1
-    assert mconf["generation"] == 1
-
-    ep.start(safekeeper_generation=1, safekeepers=mconf["sk_set"])
+    # We specify all safekeepers, so compute will connect to all of them.
+    # Only those from the current membership configuration will be used.
+    # TODO(diko): set only current safekeepers when cplane notify is implemented.
+    ep.start(safekeeper_generation=1, safekeepers=[1, 2, 3])
     ep.safe_psql("CREATE EXTENSION neon_test_utils;")
     ep.safe_psql("CREATE TABLE t(a int)")
 
@@ -61,16 +58,7 @@ def test_safekeeper_migration_simple(neon_env_builder: NeonEnvBuilder):
 
     assert ep.safe_psql("SELECT * FROM t") == [(i,) for i in range(1, 4)]
 
-    # 1 initial generation + 2 migrations on each loop iteration.
-    expected_gen = 1 + 2 * 3
-
-    mconf = env.storage_controller.timeline_locate(env.initial_tenant, env.initial_timeline)
-    assert mconf["generation"] == expected_gen
-
-    assert ep.safe_psql("SHOW neon.safekeepers")[0][0].startswith(f"g#{expected_gen}:")
-
-    # Restart and check again to make sure data is persistent.
     ep.stop()
-    ep.start(safekeeper_generation=1, safekeepers=[3])
+    ep.start(safekeeper_generation=1, safekeepers=[1, 2, 3])
 
     assert ep.safe_psql("SELECT * FROM t") == [(i,) for i in range(1, 4)]