remove "next" info_span!, it showed up in flamegraphs (not too meaningful though)

Not worth the trouble.

.github/workflows/benchmarking.yml

@@ -222,20 +222,13 @@ jobs:
       id: create-allure-report
       if: ${{ !cancelled() }}
       uses: ./.github/actions/allure-report-generate
-      with:
-        store-test-results-into-db: true
-      env:
-        REGRESS_TEST_RESULT_CONNSTR_NEW: ${{ secrets.REGRESS_TEST_RESULT_CONNSTR_NEW }}
 
     - name: Post to a Slack channel
       if: ${{ github.event.schedule && failure() }}
       uses: slackapi/slack-github-action@v1
       with:
-        channel-id: "C06T9AMNDQQ" # on-call-compute-staging-stream
-        slack-message: |
-          Periodic replication testing: ${{ job.status }}
-          <${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}|GitHub Run>
-          <${{ steps.create-allure-report.outputs.report-url }}|Allure report>
+        channel-id: "C033QLM5P7D" # dev-staging-stream
+        slack-message: "Periodic replication testing: ${{ job.status }}\n${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}"
       env:
         SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}
 
@@ -337,7 +330,7 @@ jobs:
   prepare_AWS_RDS_databases:
     uses: ./.github/workflows/_benchmarking_preparation.yml
     secrets: inherit
-
+  
   pgbench-compare:
     if: ${{ github.event.inputs.run_only_pgvector_tests == 'false' || github.event.inputs.run_only_pgvector_tests == null }}
     needs: [ generate-matrices, prepare_AWS_RDS_databases ]

.github/workflows/label-for-external-users.yml

@@ -4,7 +4,7 @@ on:
   issues:
     types:
       - opened
-  pull_request_target:
+  pull_request:
     types:
       - opened
 
@@ -25,7 +25,7 @@ jobs:
     - name: Check whether `${{ github.actor }}` is a member of `${{ github.repository_owner }}`
       id: check-user
       env:
-        GH_TOKEN: ${{ secrets.CI_ACCESS_TOKEN }}
+        GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       run: |
         if gh api -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" "/orgs/${GITHUB_REPOSITORY_OWNER}/members/${GITHUB_ACTOR}"; then
           is_member=true
@@ -45,10 +45,10 @@ jobs:
       issues: write        # for `gh issue edit`
 
     steps:
-    - name: Add `${{ env.LABEL }}` label
+    - name: Label new ${{ github.event_name }}
       env:
         GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        ITEM_NUMBER: ${{ github.event[github.event_name == 'pull_request_target' && 'pull_request' || 'issue'].number }}
-        GH_CLI_COMMAND: ${{ github.event_name == 'pull_request_target' && 'pr' || 'issue' }}
+        ITEM_NUMBER: ${{ github.event[github.event_name == 'pull_request' && 'pull_request' || 'issue'].number }}
+        GH_CLI_COMMAND: ${{ github.event_name == 'pull_request' && 'pr' || 'issue' }}
       run: |
         gh ${GH_CLI_COMMAND} --repo ${GITHUB_REPOSITORY} edit --add-label=${LABEL} ${ITEM_NUMBER}

libs/pageserver_api/src/controller_api.rs

@@ -313,17 +313,20 @@ pub struct MetadataHealthUpdateRequest {
 pub struct MetadataHealthUpdateResponse {}
 
 #[derive(Serialize, Deserialize, Debug)]
+
 pub struct MetadataHealthListUnhealthyResponse {
     pub unhealthy_tenant_shards: Vec<TenantShardId>,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
+
 pub struct MetadataHealthListOutdatedRequest {
     #[serde(with = "humantime_serde")]
     pub not_scrubbed_for: Duration,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
+
 pub struct MetadataHealthListOutdatedResponse {
     pub health_records: Vec<MetadataHealthRecord>,
 }

pageserver/src/tenant/timeline.rs

@@ -1645,20 +1645,6 @@ impl Timeline {
         self.last_record_lsn.shutdown();
 
         if try_freeze_and_flush {
-            if let Some((open, frozen)) = self
-                .layers
-                .read()
-                .await
-                .layer_map()
-                .map(|lm| (lm.open_layer.is_some(), lm.frozen_layers.len()))
-                .ok()
-                .filter(|(open, frozen)| *open || *frozen > 0)
-            {
-                tracing::info!(?open, frozen, "flushing and freezing on shutdown");
-            } else {
-                // this is double-shutdown, ignore it
-            }
-
             // we shut down walreceiver above, so, we won't add anything more
             // to the InMemoryLayer; freeze it and wait for all frozen layers
             // to reach the disk & upload queue, then shut the upload queue and

pageserver/src/tenant/timeline/compaction.rs

@@ -855,11 +855,21 @@ impl Timeline {
                 merge_iter: MergeIterator<'a>,
             },
             ValidatingStreamingKmergeBypassingPageCache {
-                mode: CompactL0BypassPageCacheValidation,
+                what: ValidationWhat,
+                concurrency: ValidationIoConcurrency,
                 merge_iter: MergeIterator<'a>,
                 all_keys_iter: VecIter<'a>,
             },
         }
+        enum ValidationIoConcurrency {
+            Sequential,
+            Concurrent,
+        }
+        enum ValidationWhat {
+            Nothing,
+            KeyLsn,
+            KeyLsnValue,
+        }
         type VecIter<'a> = std::slice::Iter<'a, DeltaEntry<'a>>; // TODO: distinguished lifetimes
         impl AllValuesIter<'_> {
             async fn next_all_keys_iter(
@@ -887,10 +897,18 @@ impl Timeline {
                       Self::next_all_keys_iter(iter, ctx).await
                     }
                     AllValuesIter::StreamingKmergeBypassingPageCache { merge_iter } => merge_iter.next().await,
-                    AllValuesIter::ValidatingStreamingKmergeBypassingPageCache { mode, merge_iter, all_keys_iter } => async {
-                        // advance both iterators
-                        let all_keys_iter_item = Self::next_all_keys_iter(all_keys_iter, ctx).await;
-                        let merge_iter_item = merge_iter.next().await;
+                    AllValuesIter::ValidatingStreamingKmergeBypassingPageCache { what, concurrency, merge_iter, all_keys_iter } => async {
+                        // advance both iterators. Use concurrency but no parallelism.
+                        let all_keys_iter_item_fut = Self::next_all_keys_iter(all_keys_iter, ctx);
+                        let merge_iter_item_fut = merge_iter.next();
+                        let (all_keys_iter_item, merge_iter_item) = match concurrency {
+                            ValidationIoConcurrency::Sequential => {
+                                (all_keys_iter_item_fut.await, merge_iter_item_fut.await)
+                            },
+                            ValidationIoConcurrency::Concurrent => {
+                                futures::future::join(all_keys_iter_item_fut, merge_iter_item_fut).await
+                            },
+                        };
                         // compare results & log warnings as needed
                         macro_rules! rate_limited_warn {
                             ($($arg:tt)*) => {{
@@ -928,16 +946,17 @@ impl Timeline {
                                 rate_limited_warn!(?merge, "all_keys_iter returned None where merge returned Some");
                             }
                             (Ok(Some((all_keys_key, all_keys_lsn, all_keys_value))), Ok(Some((merge_key, merge_lsn, merge_value)))) => {
-                                match mode {
+                                match what {
+                                    ValidationWhat::Nothing => { }
                                     // TODO: in this mode, we still load the value from disk for both iterators, even though we only need the all_keys_iter one
-                                    CompactL0BypassPageCacheValidation::KeyLsn => {
+                                    ValidationWhat::KeyLsn => {
                                         let all_keys = (all_keys_key, all_keys_lsn);
                                         let merge = (merge_key, merge_lsn);
                                         if all_keys != merge {
                                             rate_limited_warn!(?all_keys, ?merge, "merge returned a different (Key,LSN) than all_keys_iter");
                                         }
                                     }
-                                    CompactL0BypassPageCacheValidation::KeyLsnValue => {
+                                    ValidationWhat::KeyLsnValue => {
                                         let all_keys = (all_keys_key, all_keys_lsn, all_keys_value);
                                         let merge = (merge_key, merge_lsn, merge_value);
                                         if all_keys != merge {
@@ -949,7 +968,7 @@ impl Timeline {
                         }
                         // in case of mismatch, trust the legacy all_keys_iter_item
                         all_keys_iter_item
-                    }.instrument(info_span!("next")).await
+                    }.await
                 }
             }
         }
@@ -969,7 +988,32 @@ impl Timeline {
                 match validate {
                     None => AllValuesIter::StreamingKmergeBypassingPageCache { merge_iter },
                     Some(validate) => AllValuesIter::ValidatingStreamingKmergeBypassingPageCache {
-                        mode: validate.clone(),
+                        what: match &validate {
+                            CompactL0BypassPageCacheValidation::JustReadBoth
+                            | CompactL0BypassPageCacheValidation::JustReadBothConcurrentIo => {
+                                ValidationWhat::Nothing
+                            }
+                            CompactL0BypassPageCacheValidation::KeyLsn
+                            | CompactL0BypassPageCacheValidation::KeyLsnConcurrentIo => {
+                                ValidationWhat::KeyLsn
+                            }
+                            CompactL0BypassPageCacheValidation::KeyLsnValue
+                            | CompactL0BypassPageCacheValidation::KeyLsnValueConcurrentIo => {
+                                ValidationWhat::KeyLsnValue
+                            }
+                        },
+                        concurrency: match validate {
+                            CompactL0BypassPageCacheValidation::JustReadBothConcurrentIo
+                            | CompactL0BypassPageCacheValidation::KeyLsnConcurrentIo
+                            | CompactL0BypassPageCacheValidation::KeyLsnValueConcurrentIo => {
+                                ValidationIoConcurrency::Concurrent
+                            }
+                            CompactL0BypassPageCacheValidation::JustReadBoth
+                            | CompactL0BypassPageCacheValidation::KeyLsn
+                            | CompactL0BypassPageCacheValidation::KeyLsnValue => {
+                                ValidationIoConcurrency::Sequential
+                            }
+                        },
                         merge_iter,
                         all_keys_iter: all_keys.iter(),
                     },
@@ -1389,11 +1433,18 @@ pub enum CompactL0Phase1ValueAccess {
 /// See [`CompactL0Phase1ValueAccess::StreamingKmerge`].
 #[derive(Debug, PartialEq, Eq, Clone, serde::Deserialize, serde::Serialize)]
 #[serde(rename_all = "kebab-case")]
+#[allow(clippy::enum_variant_names)]
 pub enum CompactL0BypassPageCacheValidation {
+    JustReadBoth,
+    JustReadBothConcurrentIo,
     /// Validate that the series of (key, lsn) pairs are the same.
     KeyLsn,
+    // Like [`KeyLsn`], but perform the IO concurrently.
+    KeyLsnConcurrentIo,
     /// Validate that the entire output of old and new way is identical.
     KeyLsnValue,
+    // Like [`KeyLsnValue`], but perform the IO concurrently.
+    KeyLsnValueConcurrentIo,
 }
 
 impl Default for CompactL0Phase1ValueAccess {

pgxn/neon/neon.c

@@ -192,13 +192,6 @@ LogicalSlotsMonitorMain(Datum main_arg)
 	{
 		XLogRecPtr	cutoff_lsn;
 
-		/* In case of a SIGHUP, just reload the configuration. */
-		if (ConfigReloadPending)
-		{
-			ConfigReloadPending = false;
-			ProcessConfigFile(PGC_SIGHUP);
-		}
-
 		/*
 		 * If there are too many .snap files, just drop all logical slots to
 		 * prevent aux files bloat.

proxy/src/cancellation.rs

@@ -151,34 +151,21 @@ impl<P: CancellationPublisherMut> CancellationHandler<Option<Arc<Mutex<P>>>> {
 #[derive(Clone)]
 pub struct CancelClosure {
     socket_addr: SocketAddr,
-    cancel_token: Option<CancelToken>,
+    cancel_token: CancelToken,
 }
 
 impl CancelClosure {
     pub fn new(socket_addr: SocketAddr, cancel_token: CancelToken) -> Self {
         Self {
             socket_addr,
-            cancel_token: Some(cancel_token),
+            cancel_token,
         }
     }
-
-    #[cfg(test)]
-    pub fn test() -> Self {
-        use std::net::{Ipv4Addr, SocketAddrV4};
-
-        Self {
-            socket_addr: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::from_bits(0), 0)),
-            cancel_token: None,
-        }
-    }
-
     /// Cancels the query running on user's compute node.
     pub async fn try_cancel_query(self) -> Result<(), CancelError> {
-        if let Some(cancel_token) = self.cancel_token {
-            let socket = TcpStream::connect(self.socket_addr).await?;
-            cancel_token.cancel_query_raw(socket, NoTls).await?;
-            info!("query was cancelled");
-        }
+        let socket = TcpStream::connect(self.socket_addr).await?;
+        self.cancel_token.cancel_query_raw(socket, NoTls).await?;
+        info!("query was cancelled");
         Ok(())
     }
 }

proxy/src/compute.rs

@@ -16,10 +16,8 @@ use rustls::{client::danger::ServerCertVerifier, pki_types::InvalidDnsNameError}
 use std::{io, net::SocketAddr, sync::Arc, time::Duration};
 use thiserror::Error;
 use tokio::net::TcpStream;
-use tokio_postgres::{
-    tls::{MakeTlsConnect, NoTlsError},
-    Client, Connection,
-};
+use tokio_postgres::tls::MakeTlsConnect;
+use tokio_postgres_rustls::MakeRustlsConnect;
 use tracing::{error, info, warn};
 
 const COULD_NOT_CONNECT: &str = "Couldn't connect to compute node";
@@ -44,12 +42,6 @@ pub enum ConnectionError {
     TooManyConnectionAttempts(#[from] ApiLockError),
 }
 
-impl From<NoTlsError> for ConnectionError {
-    fn from(value: NoTlsError) -> Self {
-        Self::CouldNotConnect(io::Error::new(io::ErrorKind::Other, value.to_string()))
-    }
-}
-
 impl UserFacingError for ConnectionError {
     fn to_string_client(&self) -> String {
         use ConnectionError::*;
@@ -281,30 +273,6 @@ pub struct PostgresConnection {
 }
 
 impl ConnCfg {
-    /// Connect to a corresponding compute node.
-    pub async fn managed_connect<M: MakeTlsConnect<tokio::net::TcpStream>>(
-        &self,
-        ctx: &RequestMonitoring,
-        timeout: Duration,
-        mktls: &mut M,
-    ) -> Result<(SocketAddr, Client, Connection<TcpStream, M::Stream>), ConnectionError>
-    where
-        ConnectionError: From<M::Error>,
-    {
-        let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
-        let (socket_addr, stream, host) = self.connect_raw(timeout).await?;
-        drop(pause);
-
-        let tls = mktls.make_tls_connect(host)?;
-
-        // connect_raw() will not use TLS if sslmode is "disable"
-        let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
-        let (client, connection) = self.0.connect_raw(stream, tls).await?;
-        drop(pause);
-
-        Ok((socket_addr, client, connection))
-    }
-
     /// Connect to a corresponding compute node.
     pub async fn connect(
         &self,
@@ -313,6 +281,10 @@ impl ConnCfg {
         aux: MetricsAuxInfo,
         timeout: Duration,
     ) -> Result<PostgresConnection, ConnectionError> {
+        let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
+        let (socket_addr, stream, host) = self.connect_raw(timeout).await?;
+        drop(pause);
+
         let client_config = if allow_self_signed_compute {
             // Allow all certificates for creating the connection
             let verifier = Arc::new(AcceptEverythingVerifier) as Arc<dyn ServerCertVerifier>;
@@ -326,15 +298,21 @@ impl ConnCfg {
         let client_config = client_config.with_no_client_auth();
 
         let mut mk_tls = tokio_postgres_rustls::MakeRustlsConnect::new(client_config);
+        let tls = <MakeRustlsConnect as MakeTlsConnect<tokio::net::TcpStream>>::make_tls_connect(
+            &mut mk_tls,
+            host,
+        )?;
 
-        let (socket_addr, client, connection) =
-            self.managed_connect(ctx, timeout, &mut mk_tls).await?;
+        // connect_raw() will not use TLS if sslmode is "disable"
+        let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
+        let (client, connection) = self.0.connect_raw(stream, tls).await?;
+        drop(pause);
         tracing::Span::current().record("pid", tracing::field::display(client.get_process_id()));
         let stream = connection.stream.into_inner();
 
         info!(
             cold_start_info = ctx.cold_start_info().as_str(),
-            "connected to compute node ({socket_addr}) sslmode={:?}",
+            "connected to compute node at {host} ({socket_addr}) sslmode={:?}",
             self.0.get_ssl_mode()
         );
 

proxy/src/serverless/backend.rs

@@ -5,8 +5,7 @@ use tracing::{field::display, info};
 
 use crate::{
     auth::{backend::ComputeCredentials, check_peer_addr_is_in_list, AuthError},
-    cancellation::CancelClosure,
-    compute::{self, ConnectionError},
+    compute,
     config::{AuthenticationConfig, ProxyConfig},
     console::{
         errors::{GetAuthInfoError, WakeComputeError},
@@ -143,7 +142,7 @@ pub enum HttpConnError {
     #[error("pooled connection closed at inconsistent state")]
     ConnectionClosedAbruptly(#[from] tokio::sync::watch::error::SendError<uuid::Uuid>),
     #[error("could not connection to compute")]
-    ConnectionError(#[from] ConnectionError),
+    ConnectionError(#[from] tokio_postgres::Error),
 
     #[error("could not get auth info")]
     GetAuthInfo(#[from] GetAuthInfoError),
@@ -230,16 +229,17 @@ impl ConnectMechanism for TokioMechanism {
         let host = node_info.config.get_host()?;
         let permit = self.locks.get_permit(&host).await?;
 
-        let (socket_addr, client, connection) = permit.release_result(
-            node_info
-                .config
-                .managed_connect(ctx, timeout, &mut tokio_postgres::NoTls)
-                .await,
-        )?;
+        let mut config = (*node_info.config).clone();
+        let config = config
+            .user(&self.conn_info.user_info.user)
+            .password(&*self.conn_info.password)
+            .dbname(&self.conn_info.dbname)
+            .connect_timeout(timeout);
 
-        // NB: CancelToken is supposed to hold socket_addr, but we use connect_raw.
-        // Yet another reason to rework the connection establishing code.
-        let cancel_closure = CancelClosure::new(socket_addr, client.cancel_token());
+        let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
+        let res = config.connect(tokio_postgres::NoTls).await;
+        drop(pause);
+        let (client, connection) = permit.release_result(res)?;
 
         tracing::Span::current().record("pid", tracing::field::display(client.get_process_id()));
         Ok(poll_client(
@@ -250,14 +250,8 @@ impl ConnectMechanism for TokioMechanism {
             connection,
             self.conn_id,
             node_info.aux.clone(),
-            cancel_closure,
         ))
     }
 
-    fn update_connect_config(&self, config: &mut compute::ConnCfg) {
-        config
-            .user(&self.conn_info.user_info.user)
-            .dbname(&self.conn_info.dbname)
-            .password(&self.conn_info.password);
-    }
+    fn update_connect_config(&self, _config: &mut compute::ConnCfg) {}
 }

proxy/src/serverless/conn_pool.rs

@@ -12,13 +12,11 @@ use std::{
     ops::Deref,
     sync::atomic::{self, AtomicUsize},
 };
-use tokio::net::TcpStream;
 use tokio::time::Instant;
 use tokio_postgres::tls::NoTlsStream;
-use tokio_postgres::{AsyncMessage, ReadyForQueryStatus};
+use tokio_postgres::{AsyncMessage, ReadyForQueryStatus, Socket};
 use tokio_util::sync::CancellationToken;
 
-use crate::cancellation::CancelClosure;
 use crate::console::messages::{ColdStartInfo, MetricsAuxInfo};
 use crate::metrics::{HttpEndpointPoolsGuard, Metrics};
 use crate::usage_metrics::{Ids, MetricCounter, USAGE_METRICS};
@@ -465,16 +463,14 @@ impl<C: ClientInnerExt> GlobalConnPool<C> {
     }
 }
 
-#[allow(clippy::too_many_arguments)]
 pub fn poll_client<C: ClientInnerExt>(
     global_pool: Arc<GlobalConnPool<C>>,
     ctx: &RequestMonitoring,
     conn_info: ConnInfo,
     client: C,
-    mut connection: tokio_postgres::Connection<TcpStream, NoTlsStream>,
+    mut connection: tokio_postgres::Connection<Socket, NoTlsStream>,
     conn_id: uuid::Uuid,
     aux: MetricsAuxInfo,
-    cancel_closure: CancelClosure,
 ) -> Client<C> {
     let conn_gauge = Metrics::get().proxy.db_connections.guard(ctx.protocol());
     let mut session_id = ctx.session_id();
@@ -576,7 +572,6 @@ pub fn poll_client<C: ClientInnerExt>(
         cancel,
         aux,
         conn_id,
-        cancel_closure,
     };
     Client::new(inner, conn_info, pool_clone)
 }
@@ -587,7 +582,6 @@ struct ClientInner<C: ClientInnerExt> {
     cancel: CancellationToken,
     aux: MetricsAuxInfo,
     conn_id: uuid::Uuid,
-    cancel_closure: CancelClosure,
 }
 
 impl<C: ClientInnerExt> Drop for ClientInner<C> {
@@ -652,7 +646,7 @@ impl<C: ClientInnerExt> Client<C> {
             pool,
         }
     }
-    pub fn inner(&mut self) -> (&mut C, &CancelClosure, Discard<'_, C>) {
+    pub fn inner(&mut self) -> (&mut C, Discard<'_, C>) {
         let Self {
             inner,
             pool,
@@ -660,11 +654,7 @@ impl<C: ClientInnerExt> Client<C> {
             span: _,
         } = self;
         let inner = inner.as_mut().expect("client inner should not be removed");
-        (
-            &mut inner.inner,
-            &inner.cancel_closure,
-            Discard { pool, conn_info },
-        )
+        (&mut inner.inner, Discard { pool, conn_info })
     }
 }
 
@@ -761,7 +751,6 @@ mod tests {
                 cold_start_info: crate::console::messages::ColdStartInfo::Warm,
             },
             conn_id: uuid::Uuid::new_v4(),
-            cancel_closure: CancelClosure::test(),
         }
     }
 
@@ -796,7 +785,7 @@ mod tests {
         {
             let mut client = Client::new(create_inner(), conn_info.clone(), ep_pool.clone());
             assert_eq!(0, pool.get_global_connections_count());
-            client.inner().2.discard();
+            client.inner().1.discard();
             // Discard should not add the connection from the pool.
             assert_eq!(0, pool.get_global_connections_count());
         }

proxy/src/serverless/sql_over_http.rs

@@ -26,6 +26,7 @@ use tokio_postgres::error::ErrorPosition;
 use tokio_postgres::error::SqlState;
 use tokio_postgres::GenericClient;
 use tokio_postgres::IsolationLevel;
+use tokio_postgres::NoTls;
 use tokio_postgres::ReadyForQueryStatus;
 use tokio_postgres::Transaction;
 use tokio_util::sync::CancellationToken;
@@ -260,9 +261,7 @@ pub async fn handle(
 
             let mut message = e.to_string_client();
             let db_error = match &e {
-                SqlOverHttpError::ConnectCompute(HttpConnError::ConnectionError(
-                    crate::compute::ConnectionError::Postgres(e),
-                ))
+                SqlOverHttpError::ConnectCompute(HttpConnError::ConnectionError(e))
                 | SqlOverHttpError::Postgres(e) => e.as_db_error(),
                 _ => None,
             };
@@ -623,7 +622,8 @@ impl QueryData {
         client: &mut Client<tokio_postgres::Client>,
         parsed_headers: HttpHeaders,
     ) -> Result<String, SqlOverHttpError> {
-        let (inner, cancel_token, mut discard) = client.inner();
+        let (inner, mut discard) = client.inner();
+        let cancel_token = inner.cancel_token();
 
         let res = match select(
             pin!(query_to_json(&*inner, self, &mut 0, parsed_headers)),
@@ -647,7 +647,7 @@ impl QueryData {
             // The query was cancelled.
             Either::Right((_cancelled, query)) => {
                 tracing::info!("cancelling query");
-                if let Err(err) = cancel_token.clone().try_cancel_query().await {
+                if let Err(err) = cancel_token.cancel_query(NoTls).await {
                     tracing::error!(?err, "could not cancel query");
                 }
                 // wait for the query cancellation
@@ -663,9 +663,7 @@ impl QueryData {
                     // query failed or was cancelled.
                     Ok(Err(error)) => {
                         let db_error = match &error {
-                            SqlOverHttpError::ConnectCompute(HttpConnError::ConnectionError(
-                                crate::compute::ConnectionError::Postgres(e),
-                            ))
+                            SqlOverHttpError::ConnectCompute(HttpConnError::ConnectionError(e))
                             | SqlOverHttpError::Postgres(e) => e.as_db_error(),
                             _ => None,
                         };
@@ -696,7 +694,8 @@ impl BatchQueryData {
         parsed_headers: HttpHeaders,
     ) -> Result<String, SqlOverHttpError> {
         info!("starting transaction");
-        let (inner, cancel_token, mut discard) = client.inner();
+        let (inner, mut discard) = client.inner();
+        let cancel_token = inner.cancel_token();
         let mut builder = inner.build_transaction();
         if let Some(isolation_level) = parsed_headers.txn_isolation_level {
             builder = builder.isolation_level(isolation_level);
@@ -729,7 +728,7 @@ impl BatchQueryData {
                     json_output
                 }
                 Err(SqlOverHttpError::Cancelled(_)) => {
-                    if let Err(err) = cancel_token.clone().try_cancel_query().await {
+                    if let Err(err) = cancel_token.cancel_query(NoTls).await {
                         tracing::error!(?err, "could not cancel query");
                     }
                     // TODO: after cancelling, wait to see if we can get a status. maybe the connection is still safe.

storage_controller/src/http.rs

@@ -500,7 +500,7 @@ async fn handle_node_configure(mut req: Request<Body>) -> Result<Response<Body>,
         StatusCode::OK,
         state
             .service
-            .external_node_configure(
+            .node_configure(
                 config_req.node_id,
                 config_req.availability.map(NodeAvailability::from),
                 config_req.scheduling,

storage_controller/src/main.rs

@@ -196,26 +196,14 @@ async fn migration_run(database_url: &str) -> anyhow::Result<()> {
 }
 
 fn main() -> anyhow::Result<()> {
-    logging::init(
-        LogFormat::Plain,
-        logging::TracingErrorLayerEnablement::Disabled,
-        logging::Output::Stdout,
-    )?;
-
-    // log using tracing so we don't get confused output by default hook writing to stderr
-    utils::logging::replace_panic_hook_with_tracing_panic_hook().forget();
-
-    let _sentry_guard = init_sentry(Some(GIT_VERSION.into()), &[]);
-
-    let hook = std::panic::take_hook();
+    let default_panic = std::panic::take_hook();
     std::panic::set_hook(Box::new(move |info| {
-        // let sentry send a message (and flush)
-        // and trace the error
-        hook(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.
@@ -229,6 +217,12 @@ fn main() -> anyhow::Result<()> {
 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();

storage_controller/src/service.rs

@@ -4912,26 +4912,6 @@ impl Service {
         Ok(())
     }
 
-    /// Wrapper around [`Self::node_configure`] which only allows changes while there is no ongoing
-    /// operation for HTTP api.
-    pub(crate) async fn external_node_configure(
-        &self,
-        node_id: NodeId,
-        availability: Option<NodeAvailability>,
-        scheduling: Option<NodeSchedulingPolicy>,
-    ) -> Result<(), ApiError> {
-        {
-            let locked = self.inner.read().unwrap();
-            if let Some(op) = locked.ongoing_operation.as_ref().map(|op| op.operation) {
-                return Err(ApiError::PreconditionFailed(
-                    format!("Ongoing background operation forbids configuring: {op}").into(),
-                ));
-            }
-        }
-
-        self.node_configure(node_id, availability, scheduling).await
-    }
-
     pub(crate) async fn start_node_drain(
         self: &Arc<Self>,
         node_id: NodeId,
@@ -4989,8 +4969,6 @@ impl Service {
                     cancel: cancel.clone(),
                 });
 
-                let span = tracing::info_span!(parent: None, "drain_node", %node_id);
-
                 tokio::task::spawn({
                     let service = self.clone();
                     let cancel = cancel.clone();
@@ -5007,21 +4985,21 @@ impl Service {
                             }
                         }
 
-                        tracing::info!("Drain background operation starting");
+                        tracing::info!(%node_id, "Drain background operation starting");
                         let res = service.drain_node(node_id, cancel).await;
                         match res {
                             Ok(()) => {
-                                tracing::info!("Drain background operation completed successfully");
+                                tracing::info!(%node_id, "Drain background operation completed successfully");
                             }
                             Err(OperationError::Cancelled) => {
-                                tracing::info!("Drain background operation was cancelled");
+                                tracing::info!(%node_id, "Drain background operation was cancelled");
                             }
                             Err(err) => {
-                                tracing::error!("Drain background operation encountered: {err}")
+                                tracing::error!(%node_id, "Drain background operation encountered: {err}")
                             }
                         }
                     }
-                }.instrument(span));
+                });
             }
             NodeSchedulingPolicy::Draining => {
                 return Err(ApiError::Conflict(format!(
@@ -5039,14 +5017,14 @@ impl Service {
     }
 
     pub(crate) async fn cancel_node_drain(&self, node_id: NodeId) -> Result<(), ApiError> {
-        let node_available = {
+        let (node_available, node_policy) = {
             let locked = self.inner.read().unwrap();
             let nodes = &locked.nodes;
             let node = nodes.get(&node_id).ok_or(ApiError::NotFound(
                 anyhow::anyhow!("Node {} not registered", node_id).into(),
             ))?;
 
-            node.is_available()
+            (node.is_available(), node.get_scheduling())
         };
 
         if !node_available {
@@ -5055,6 +5033,12 @@ impl Service {
             ));
         }
 
+        if !matches!(node_policy, NodeSchedulingPolicy::Draining) {
+            return Err(ApiError::PreconditionFailed(
+                format!("Node {node_id} has no drain in progress").into(),
+            ));
+        }
+
         if let Some(op_handler) = self.inner.read().unwrap().ongoing_operation.as_ref() {
             if let Operation::Drain(drain) = op_handler.operation {
                 if drain.node_id == node_id {
@@ -5120,8 +5104,6 @@ impl Service {
                     cancel: cancel.clone(),
                 });
 
-                let span = tracing::info_span!(parent: None, "fill_node", %node_id);
-
                 tokio::task::spawn({
                     let service = self.clone();
                     let cancel = cancel.clone();
@@ -5138,21 +5120,21 @@ impl Service {
                             }
                         }
 
-                        tracing::info!("Fill background operation starting");
+                        tracing::info!(%node_id, "Fill background operation starting");
                         let res = service.fill_node(node_id, cancel).await;
                         match res {
                             Ok(()) => {
-                                tracing::info!("Fill background operation completed successfully");
+                                tracing::info!(%node_id, "Fill background operation completed successfully");
                             }
                             Err(OperationError::Cancelled) => {
-                                tracing::info!("Fill background operation was cancelled");
+                                tracing::info!(%node_id, "Fill background operation was cancelled");
                             }
                             Err(err) => {
-                                tracing::error!("Fill background operation encountered: {err}")
+                                tracing::error!(%node_id, "Fill background operation encountered: {err}")
                             }
                         }
                     }
-                }.instrument(span));
+                });
             }
             NodeSchedulingPolicy::Filling => {
                 return Err(ApiError::Conflict(format!(
@@ -5170,14 +5152,14 @@ impl Service {
     }
 
     pub(crate) async fn cancel_node_fill(&self, node_id: NodeId) -> Result<(), ApiError> {
-        let node_available = {
+        let (node_available, node_policy) = {
             let locked = self.inner.read().unwrap();
             let nodes = &locked.nodes;
             let node = nodes.get(&node_id).ok_or(ApiError::NotFound(
                 anyhow::anyhow!("Node {} not registered", node_id).into(),
             ))?;
 
-            node.is_available()
+            (node.is_available(), node.get_scheduling())
         };
 
         if !node_available {
@@ -5186,6 +5168,12 @@ impl Service {
             ));
         }
 
+        if !matches!(node_policy, NodeSchedulingPolicy::Filling) {
+            return Err(ApiError::PreconditionFailed(
+                format!("Node {node_id} has no fill in progress").into(),
+            ));
+        }
+
         if let Some(op_handler) = self.inner.read().unwrap().ongoing_operation.as_ref() {
             if let Operation::Fill(fill) = op_handler.operation {
                 if fill.node_id == node_id {
@@ -5994,7 +5982,7 @@ impl Service {
                 .await_waiters_remainder(waiters, SHORT_RECONCILE_TIMEOUT)
                 .await;
 
-            failpoint_support::sleep_millis_async!("sleepy-drain-loop", &cancel);
+            failpoint_support::sleep_millis_async!("sleepy-drain-loop");
         }
 
         while !waiters.is_empty() {

test_runner/performance/test_logical_replication.py

@@ -262,85 +262,3 @@ def test_publisher_restart(
             sub_workload.terminate()
     finally:
         pub_workload.terminate()
-
-
-@pytest.mark.remote_cluster
-@pytest.mark.timeout(2 * 60 * 60)
-def test_snap_files(
-    pg_bin: PgBin,
-    benchmark_project_pub: NeonApiEndpoint,
-    zenbenchmark: NeonBenchmarker,
-):
-    """
-    Creates a node with a replication slot. Generates pgbench into the replication slot,
-    then runs pgbench inserts while generating large numbers of snapfiles. Then restarts
-    the node and tries to peek the replication changes.
-    """
-    test_duration_min = 60
-    test_interval_min = 5
-    pgbench_duration = f"-T{test_duration_min * 60 * 2}"
-
-    env = benchmark_project_pub.pgbench_env
-    connstr = benchmark_project_pub.connstr
-    pg_bin.run_capture(["pgbench", "-i", "-s100"], env=env)
-
-    with psycopg2.connect(connstr) as conn:
-        conn.autocommit = True
-        with conn.cursor() as cur:
-            cur.execute("SELECT rolsuper FROM pg_roles WHERE rolname = 'neondb_owner'")
-            is_super = cur.fetchall()[0]
-            assert is_super, "This benchmark won't work if we don't have superuser"
-
-    conn = psycopg2.connect(connstr)
-    conn.autocommit = True
-    cur = conn.cursor()
-    cur.execute("ALTER SYSTEM SET neon.logical_replication_max_snap_files = -1")
-
-    with psycopg2.connect(connstr) as conn:
-        conn.autocommit = True
-        with conn.cursor() as cur:
-            cur.execute("SELECT pg_reload_conf()")
-
-    with psycopg2.connect(connstr) as conn:
-        conn.autocommit = True
-        with conn.cursor() as cur:
-            cur.execute(
-                """
-                DO $$
-                    BEGIN
-                    IF EXISTS (
-                        SELECT 1
-                        FROM pg_replication_slots
-                        WHERE slot_name = 'slotter'
-                    ) THEN
-                        PERFORM pg_drop_replication_slot('slotter');
-                    END IF;
-                END $$;
-            """
-            )
-            cur.execute("SELECT pg_create_logical_replication_slot('slotter', 'test_decoding')")
-
-    workload = pg_bin.run_nonblocking(["pgbench", "-c10", pgbench_duration, "-Mprepared"], env=env)
-    try:
-        start = time.time()
-        prev_measurement = time.time()
-        while time.time() - start < test_duration_min * 60:
-            with psycopg2.connect(connstr) as conn:
-                with conn.cursor() as cur:
-                    cur.execute(
-                        "SELECT count(*) FROM (SELECT pg_log_standby_snapshot() FROM generate_series(1, 10000) g) s"
-                    )
-                    check_pgbench_still_running(workload)
-                    cur.execute(
-                        "SELECT pg_replication_slot_advance('slotter', pg_current_wal_lsn())"
-                    )
-
-            # Measure storage
-            if time.time() - prev_measurement > test_interval_min * 60:
-                storage = benchmark_project_pub.get_synthetic_storage_size()
-                zenbenchmark.record("storage", storage, "B", MetricReport.LOWER_IS_BETTER)
-                prev_measurement = time.time()
-            time.sleep(test_interval_min * 60 / 3)
-
-    finally:
-        workload.terminate()

test_runner/regress/test_storage_controller.py

@@ -2091,47 +2091,3 @@ def test_storage_controller_step_down(neon_env_builder: NeonEnvBuilder):
         )
         == 0
     )
-
-
-def test_storage_controller_ps_restarted_during_drain(neon_env_builder: NeonEnvBuilder):
-    # single unsharded tenant, two locations
-    neon_env_builder.num_pageservers = 2
-
-    env = neon_env_builder.init_start()
-
-    env.storage_controller.tenant_policy_update(env.initial_tenant, {"placement": {"Attached": 1}})
-    env.storage_controller.reconcile_until_idle()
-
-    attached_id = int(env.storage_controller.locate(env.initial_tenant)[0]["node_id"])
-    attached = next((ps for ps in env.pageservers if ps.id == attached_id))
-
-    def attached_is_draining():
-        details = env.storage_controller.node_status(attached.id)
-        assert details["scheduling"] == "Draining"
-
-    env.storage_controller.configure_failpoints(("sleepy-drain-loop", "return(10000)"))
-    env.storage_controller.node_drain(attached.id)
-
-    wait_until(10, 0.5, attached_is_draining)
-
-    attached.restart()
-
-    # we are unable to reconfigure node while the operation is still ongoing
-    with pytest.raises(
-        StorageControllerApiException,
-        match="Precondition failed: Ongoing background operation forbids configuring: drain.*",
-    ):
-        env.storage_controller.node_configure(attached.id, {"scheduling": "Pause"})
-    with pytest.raises(
-        StorageControllerApiException,
-        match="Precondition failed: Ongoing background operation forbids configuring: drain.*",
-    ):
-        env.storage_controller.node_configure(attached.id, {"availability": "Offline"})
-
-    env.storage_controller.cancel_node_drain(attached.id)
-
-    def reconfigure_node_again():
-        env.storage_controller.node_configure(attached.id, {"scheduling": "Pause"})
-
-    # allow for small delay between actually having cancelled and being able reconfigure again
-    wait_until(4, 0.5, reconfigure_node_again)