standardise logging

## Problem

https://github.com/neondatabase/cloud/issues/30539

If the current leader cancels the `call` function, then it has removed
the jobs from the queue, but will never finish sending the responses.
Because of this, it is not cancellation safe.

## Summary of changes

Document these functions as not cancellation safe. Move cancellation of
the queued jobs into the queue itself.

## Alternatives considered

1. We could spawn the task that runs the batch, since that won't get
cancelled.
* This requires `fn call(self: Arc<Self>)` or `fn call(&'static self)`.
2. We could add another scopeguard and return the requests back to the
queue.
* This requires that requests are always retry safe, and also requires
requests to be `Clone`.

libs/postgres_ffi/benches/waldecoder.rs

@@ -6,7 +6,7 @@ use postgres_ffi::v17::waldecoder_handler::WalStreamDecoderHandler;
 use postgres_ffi::waldecoder::WalStreamDecoder;
 use postgres_versioninfo::PgMajorVersion;
 use pprof::criterion::{Output, PProfProfiler};
-use utils::lsn::{Lsn, WalSegmentSize};
+use utils::lsn::Lsn;
 
 const KB: usize = 1024;
 
@@ -22,26 +22,23 @@ criterion_main!(benches);
 fn bench_complete_record(c: &mut Criterion) {
     let mut g = c.benchmark_group("complete_record");
     for size in [64, KB, 8 * KB, 128 * KB] {
-        let value_size = size as WalSegmentSize;
         // Kind of weird to change the group throughput per benchmark, but it's the only way
         // to vary it per benchmark. It works.
-        g.throughput(criterion::Throughput::Bytes(value_size as u64));
-        g.bench_function(format!("size={size}"), |b| {
-            run_bench(b, value_size).unwrap()
-        });
+        g.throughput(criterion::Throughput::Bytes(size as u64));
+        g.bench_function(format!("size={size}"), |b| run_bench(b, size).unwrap());
     }
 
-    fn run_bench(b: &mut Bencher, size: WalSegmentSize) -> anyhow::Result<()> {
+    fn run_bench(b: &mut Bencher, size: usize) -> anyhow::Result<()> {
         const PREFIX: &CStr = c"";
         let value_size = LogicalMessageGenerator::make_value_size(size, PREFIX);
-        let value = vec![1; value_size as usize];
+        let value = vec![1; value_size];
 
         let mut decoder = WalStreamDecoder::new(Lsn(0), PgMajorVersion::PG17);
         let msg = LogicalMessageGenerator::new(PREFIX, &value)
             .next()
             .unwrap()
             .encode(Lsn(0));
-        assert_eq!(msg.len(), size as usize);
+        assert_eq!(msg.len(), size);
 
         b.iter(|| {
             let msg = msg.clone(); // Bytes::clone() is cheap

libs/postgres_ffi/src/lib.rs

@@ -12,7 +12,7 @@
 
 use bytes::Bytes;
 use utils::bin_ser::SerializeError;
-use utils::lsn::{Lsn, WalSegmentSize};
+use utils::lsn::Lsn;
 
 pub use postgres_versioninfo::PgMajorVersion;
 
@@ -241,7 +241,7 @@ pub use v14::xlog_utils::{
 pub const BLCKSZ: u16 = 8192;
 pub const RELSEG_SIZE: u32 = 1024 * 1024 * 1024 / (BLCKSZ as u32);
 pub const XLOG_BLCKSZ: usize = 8192;
-pub const WAL_SEGMENT_SIZE: WalSegmentSize = 16 * 1024 * 1024;
+pub const WAL_SEGMENT_SIZE: usize = 16 * 1024 * 1024;
 
 pub const MAX_SEND_SIZE: usize = XLOG_BLCKSZ * 16;
 

libs/postgres_ffi/src/wal_generator.rs

@@ -2,7 +2,7 @@ use std::ffi::{CStr, CString};
 
 use bytes::{Bytes, BytesMut};
 use crc32c::crc32c_append;
-use utils::lsn::{Lsn, WalSegmentSize};
+use utils::lsn::Lsn;
 
 use super::bindings::{RmgrId, XLogLongPageHeaderData, XLogPageHeaderData, XLOG_PAGE_MAGIC};
 use super::xlog_utils::{
@@ -39,7 +39,7 @@ impl Record {
 
         // Construct the WAL record header.
         let mut header = XLogRecord {
-            xl_tot_len: XLOG_SIZE_OF_XLOG_RECORD + data_header.len() as WalSegmentSize + self.data.len() as WalSegmentSize,
+            xl_tot_len: (XLOG_SIZE_OF_XLOG_RECORD + data_header.len() + self.data.len()) as u32,
             xl_xid: 0,
             xl_prev: prev_lsn.into(),
             xl_info: self.info,
@@ -158,7 +158,7 @@ impl<R: RecordGenerator> WalGenerator<R> {
                     XLogLongPageHeaderData {
                         std: page_header,
                         xlp_sysid: Self::SYS_ID,
-                        xlp_seg_size: WAL_SEGMENT_SIZE,
+                        xlp_seg_size: WAL_SEGMENT_SIZE as u32,
                         xlp_xlog_blcksz: XLOG_BLCKSZ as u32,
                     }
                     .encode()
@@ -234,10 +234,10 @@ impl LogicalMessageGenerator {
 
     /// Computes how large a value must be to get a record of the given size. Convenience method to
     /// construct records of pre-determined size. Panics if the record size is too small.
-    pub fn make_value_size(record_size: WalSegmentSize, prefix: &CStr) -> WalSegmentSize {
+    pub fn make_value_size(record_size: usize, prefix: &CStr) -> usize {
         let xlog_header_size = XLOG_SIZE_OF_XLOG_RECORD;
-        let lm_header_size = size_of::<XlLogicalMessage>() as WalSegmentSize;
-        let prefix_size = prefix.to_bytes_with_nul().len() as WalSegmentSize;
+        let lm_header_size = size_of::<XlLogicalMessage>();
+        let prefix_size = prefix.to_bytes_with_nul().len();
         let data_header_size = match record_size - xlog_header_size - 2 {
             0..=255 => 2,
             256..=258 => panic!("impossible record_size {record_size}"),

libs/postgres_ffi/src/waldecoder_handler.rs

@@ -108,7 +108,7 @@ impl WalStreamDecoderHandler for WalStreamDecoder {
                     if self.lsn.segment_offset(WAL_SEGMENT_SIZE) == 0 {
                         // parse long header
 
-                        if self.inputbuf.remaining() < XLOG_SIZE_OF_XLOG_LONG_PHD as usize{
+                        if self.inputbuf.remaining() < XLOG_SIZE_OF_XLOG_LONG_PHD {
                             return Ok(None);
                         }
 
@@ -123,7 +123,7 @@ impl WalStreamDecoderHandler for WalStreamDecoder {
 
                         self.lsn += XLOG_SIZE_OF_XLOG_LONG_PHD as u64;
                     } else if self.lsn.block_offset() == 0 {
-                        if self.inputbuf.remaining() < XLOG_SIZE_OF_XLOG_SHORT_PHD as usize{
+                        if self.inputbuf.remaining() < XLOG_SIZE_OF_XLOG_SHORT_PHD {
                             return Ok(None);
                         }
 
@@ -153,7 +153,7 @@ impl WalStreamDecoderHandler for WalStreamDecoder {
                     // peek xl_tot_len at the beginning of the record.
                     // FIXME: assumes little-endian
                     let xl_tot_len = (&self.inputbuf[0..4]).get_u32_le();
-                    if xl_tot_len < XLOG_SIZE_OF_XLOG_RECORD {
+                    if (xl_tot_len as usize) < XLOG_SIZE_OF_XLOG_RECORD {
                         return Err(WalDecodeError {
                             msg: format!("invalid xl_tot_len {xl_tot_len}"),
                             lsn: self.lsn,
@@ -216,7 +216,7 @@ impl WalStreamDecoderHandler for WalStreamDecoder {
     fn complete_record(&mut self, recordbuf: Bytes) -> Result<(Lsn, Bytes), WalDecodeError> {
         // We now have a record in the 'recordbuf' local variable.
         let xlogrec =
-            XLogRecord::from_slice(&recordbuf[0..XLOG_SIZE_OF_XLOG_RECORD as usize]).map_err(|e| {
+            XLogRecord::from_slice(&recordbuf[0..XLOG_SIZE_OF_XLOG_RECORD]).map_err(|e| {
                 WalDecodeError {
                     msg: format!("xlog record deserialization failed {e}"),
                     lsn: self.lsn,

libs/postgres_ffi/src/walrecord.rs

@@ -266,7 +266,7 @@ pub fn decode_wal_record(
         xlogrec.xl_info
     );
 
-    let remaining: usize = (xlogrec.xl_tot_len - XLOG_SIZE_OF_XLOG_RECORD) as usize;
+    let remaining: usize = xlogrec.xl_tot_len as usize - XLOG_SIZE_OF_XLOG_RECORD;
 
     if buf.remaining() != remaining {
         //TODO error

libs/postgres_ffi/src/xlog_utils.rs

@@ -35,7 +35,7 @@ use std::time::SystemTime;
 use utils::bin_ser::DeserializeError;
 use utils::bin_ser::SerializeError;
 
-use utils::lsn::{Lsn, WalSegmentSize};
+use utils::lsn::Lsn;
 
 pub const XLOG_FNAME_LEN: usize = 24;
 pub const XLP_BKP_REMOVABLE: u16 = 0x0004;
@@ -43,9 +43,9 @@ pub const XLP_FIRST_IS_CONTRECORD: u16 = 0x0001;
 pub const XLP_REM_LEN_OFFS: usize = 2 + 2 + 4 + 8;
 pub const XLOG_RECORD_CRC_OFFS: usize = 4 + 4 + 8 + 1 + 1 + 2;
 
-pub const XLOG_SIZE_OF_XLOG_SHORT_PHD: WalSegmentSize = size_of::<XLogPageHeaderData>() as WalSegmentSize;
-pub const XLOG_SIZE_OF_XLOG_LONG_PHD: WalSegmentSize = size_of::<XLogLongPageHeaderData>() as WalSegmentSize;
-pub const XLOG_SIZE_OF_XLOG_RECORD: WalSegmentSize = size_of::<XLogRecord>() as WalSegmentSize;
+pub const XLOG_SIZE_OF_XLOG_SHORT_PHD: usize = size_of::<XLogPageHeaderData>();
+pub const XLOG_SIZE_OF_XLOG_LONG_PHD: usize = size_of::<XLogLongPageHeaderData>();
+pub const XLOG_SIZE_OF_XLOG_RECORD: usize = size_of::<XLogRecord>();
 #[allow(clippy::identity_op)]
 pub const SIZE_OF_XLOG_RECORD_DATA_HEADER_SHORT: usize = 1 * 2;
 
@@ -58,19 +58,19 @@ pub const SIZE_OF_XLOG_RECORD_DATA_HEADER_SHORT: usize = 1 * 2;
 /// in order to let CLOG_TRUNCATE mechanism correctly extend CLOG.
 const XID_CHECKPOINT_INTERVAL: u32 = 1024;
 
-pub fn XLogSegmentsPerXLogId(wal_segsz_bytes: WalSegmentSize) -> XLogSegNo {
+pub fn XLogSegmentsPerXLogId(wal_segsz_bytes: usize) -> XLogSegNo {
     (0x100000000u64 / wal_segsz_bytes as u64) as XLogSegNo
 }
 
 pub fn XLogSegNoOffsetToRecPtr(
     segno: XLogSegNo,
     offset: u32,
-    wal_segsz_bytes: WalSegmentSize,
+    wal_segsz_bytes: usize,
 ) -> XLogRecPtr {
     segno * (wal_segsz_bytes as u64) + (offset as u64)
 }
 
-pub fn XLogFileName(tli: TimeLineID, logSegNo: XLogSegNo, wal_segsz_bytes: WalSegmentSize) -> String {
+pub fn XLogFileName(tli: TimeLineID, logSegNo: XLogSegNo, wal_segsz_bytes: usize) -> String {
     format!(
         "{:>08X}{:>08X}{:>08X}",
         tli,
@@ -81,7 +81,7 @@ pub fn XLogFileName(tli: TimeLineID, logSegNo: XLogSegNo, wal_segsz_bytes: WalSe
 
 pub fn XLogFromFileName(
     fname: &OsStr,
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
 ) -> anyhow::Result<(XLogSegNo, TimeLineID)> {
     if let Some(fname_str) = fname.to_str() {
         let tli = u32::from_str_radix(&fname_str[0..8], 16)?;
@@ -111,7 +111,7 @@ pub fn IsPartialXLogFileName(fname: &OsStr) -> bool {
 
 /// If LSN points to the beginning of the page, then shift it to first record,
 /// otherwise align on 8-bytes boundary (required for WAL records)
-pub fn normalize_lsn(lsn: Lsn, seg_sz: WalSegmentSize) -> Lsn {
+pub fn normalize_lsn(lsn: Lsn, seg_sz: usize) -> Lsn {
     if lsn.0 % XLOG_BLCKSZ as u64 == 0 {
         let hdr_size = if lsn.0 % seg_sz as u64 == 0 {
             XLOG_SIZE_OF_XLOG_LONG_PHD
@@ -227,7 +227,7 @@ pub use timestamp_conversions::{to_pg_timestamp, try_from_pg_timestamp};
 // back.
 pub fn find_end_of_wal(
     data_dir: &Path,
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
     start_lsn: Lsn, // start reading WAL at this point; must point at record start_lsn.
 ) -> anyhow::Result<Lsn> {
     let mut result = start_lsn;
@@ -431,14 +431,14 @@ impl CheckPoint {
 /// page of the segment and the page that contains the given LSN.
 /// We need this segment to start compute node.
 pub fn generate_wal_segment(segno: u64, system_id: u64, lsn: Lsn) -> Result<Bytes, SerializeError> {
-    let mut seg_buf = BytesMut::with_capacity(WAL_SEGMENT_SIZE as usize);
+    let mut seg_buf = BytesMut::with_capacity(WAL_SEGMENT_SIZE);
 
     let pageaddr = XLogSegNoOffsetToRecPtr(segno, 0, WAL_SEGMENT_SIZE);
 
     let page_off = lsn.block_offset();
     let seg_off = lsn.segment_offset(WAL_SEGMENT_SIZE);
 
-    let first_page_only = seg_off < XLOG_BLCKSZ as WalSegmentSize;
+    let first_page_only = seg_off < XLOG_BLCKSZ;
     // If first records starts in the middle of the page, pretend in page header
     // there is a fake record which ends where first real record starts. This
     // makes pg_waldump etc happy.
@@ -460,12 +460,12 @@ pub fn generate_wal_segment(segno: u64, system_id: u64, lsn: Lsn) -> Result<Byte
                 xlp_info: pg_constants::XLP_LONG_HEADER | infoflags,
                 xlp_tli: PG_TLI,
                 xlp_pageaddr: pageaddr,
-                xlp_rem_len: shdr_rem_len,
+                xlp_rem_len: shdr_rem_len as u32,
                 ..Default::default() // Put 0 in padding fields.
             }
         },
         xlp_sysid: system_id,
-        xlp_seg_size: WAL_SEGMENT_SIZE,
+        xlp_seg_size: WAL_SEGMENT_SIZE as u32,
         xlp_xlog_blcksz: XLOG_BLCKSZ as u32,
     };
 
@@ -473,7 +473,7 @@ pub fn generate_wal_segment(segno: u64, system_id: u64, lsn: Lsn) -> Result<Byte
     seg_buf.extend_from_slice(&hdr_bytes);
 
     //zero out the rest of the file
-    seg_buf.resize(WAL_SEGMENT_SIZE as usize, 0);
+    seg_buf.resize(WAL_SEGMENT_SIZE, 0);
 
     if !first_page_only {
         let block_offset = lsn.page_offset_in_segment(WAL_SEGMENT_SIZE) as usize;

libs/postgres_ffi/wal_craft/src/lib.rs

@@ -389,7 +389,7 @@ impl Crafter for LastWalRecordXlogSwitchEndsOnPageBoundary {
             let xlog_switch_record_end: PgLsn =
                 client.query_one("SELECT pg_switch_wal()", &[])?.get(0);
 
-            if (u64::from(xlog_switch_record_end) % XLOG_BLCKSZ as u64) as u32
+            if u64::from(xlog_switch_record_end) as usize % XLOG_BLCKSZ
                 != XLOG_SIZE_OF_XLOG_SHORT_PHD
             {
                 warn!(

libs/postgres_ffi/wal_craft/src/xlog_utils_test.rs

@@ -81,10 +81,10 @@ fn test_end_of_wal<C: crate::Crafter>(test_name: &str) {
                 continue;
             }
             let mut f = File::options().write(true).open(file.path()).unwrap();
-            static ZEROS: [u8; WAL_SEGMENT_SIZE as usize] = [0u8; WAL_SEGMENT_SIZE as usize];
+            static ZEROS: [u8; WAL_SEGMENT_SIZE] = [0u8; WAL_SEGMENT_SIZE];
             f.write_all(
                 &ZEROS[0..min(
-                    WAL_SEGMENT_SIZE as usize,
+                    WAL_SEGMENT_SIZE,
                     (u64::from(*start_lsn) - seg_start_lsn) as usize,
                 )],
             )

libs/proxy/tokio-postgres2/src/config.rs

@@ -12,7 +12,9 @@ use tokio::net::TcpStream;
 
 use crate::connect::connect;
 use crate::connect_raw::{RawConnection, connect_raw};
-use crate::tls::{MakeTlsConnect, TlsConnect};
+use crate::connect_tls::connect_tls;
+use crate::maybe_tls_stream::MaybeTlsStream;
+use crate::tls::{MakeTlsConnect, TlsConnect, TlsStream};
 use crate::{Client, Connection, Error};
 
 /// TLS configuration.
@@ -238,7 +240,7 @@ impl Config {
         connect(tls, self).await
     }
 
-    pub async fn connect_raw<S, T>(
+    pub async fn tls_and_authenticate<S, T>(
         &self,
         stream: S,
         tls: T,
@@ -247,7 +249,19 @@ impl Config {
         S: AsyncRead + AsyncWrite + Unpin,
         T: TlsConnect<S>,
     {
-        connect_raw(stream, tls, self).await
+        let stream = connect_tls(stream, self.ssl_mode, tls).await?;
+        connect_raw(stream, self).await
+    }
+
+    pub async fn authenticate<S, T>(
+        &self,
+        stream: MaybeTlsStream<S, T>,
+    ) -> Result<RawConnection<S, T>, Error>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+        T: TlsStream + Unpin,
+    {
+        connect_raw(stream, self).await
     }
 }
 

libs/proxy/tokio-postgres2/src/connect.rs

@@ -1,14 +1,16 @@
 use std::net::IpAddr;
 
 use postgres_protocol2::message::backend::Message;
+use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::TcpStream;
 use tokio::sync::mpsc;
 
 use crate::client::SocketConfig;
 use crate::codec::BackendMessage;
-use crate::config::Host;
+use crate::config::{Host, SslMode};
 use crate::connect_raw::connect_raw;
 use crate::connect_socket::connect_socket;
+use crate::connect_tls::connect_tls;
 use crate::tls::{MakeTlsConnect, TlsConnect};
 use crate::{Client, Config, Connection, Error, RawConnection};
 
@@ -44,13 +46,8 @@ where
     T: TlsConnect<TcpStream>,
 {
     let socket = connect_socket(host_addr, host, port, config.connect_timeout).await?;
-    let RawConnection {
-        stream,
-        parameters,
-        delayed_notice,
-        process_id,
-        secret_key,
-    } = connect_raw(socket, tls, config).await?;
+    let stream = connect_tls(socket, config.ssl_mode, tls).await?;
+    let raw = connect_raw(stream, config).await?;
 
     let socket_config = SocketConfig {
         host_addr,
@@ -59,24 +56,46 @@ where
         connect_timeout: config.connect_timeout,
     };
 
-    let (client_tx, conn_rx) = mpsc::unbounded_channel();
-    let (conn_tx, client_rx) = mpsc::channel(4);
-    let client = Client::new(
-        client_tx,
-        client_rx,
-        socket_config,
-        config.ssl_mode,
-        process_id,
-        secret_key,
-    );
-
-    // delayed notices are always sent as "Async" messages.
-    let delayed = delayed_notice
-        .into_iter()
-        .map(|m| BackendMessage::Async(Message::NoticeResponse(m)))
-        .collect();
-
-    let connection = Connection::new(stream, delayed, parameters, conn_tx, conn_rx);
-
-    Ok((client, connection))
+    Ok(raw.into_managed_conn(socket_config, config.ssl_mode))
+}
+
+impl<S, T> RawConnection<S, T>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    pub fn into_managed_conn(
+        self,
+        socket_config: SocketConfig,
+        ssl_mode: SslMode,
+    ) -> (Client, Connection<S, T>) {
+        let RawConnection {
+            stream,
+            parameters,
+            delayed_notice,
+            process_id,
+            secret_key,
+        } = self;
+
+        let (client_tx, conn_rx) = mpsc::unbounded_channel();
+        let (conn_tx, client_rx) = mpsc::channel(4);
+        let client = Client::new(
+            client_tx,
+            client_rx,
+            socket_config,
+            ssl_mode,
+            process_id,
+            secret_key,
+        );
+
+        // delayed notices are always sent as "Async" messages.
+        let delayed = delayed_notice
+            .into_iter()
+            .map(|m| BackendMessage::Async(Message::NoticeResponse(m)))
+            .collect();
+
+        let connection = Connection::new(stream, delayed, parameters, conn_tx, conn_rx);
+
+        (client, connection)
+    }
 }

libs/proxy/tokio-postgres2/src/connect_raw.rs

@@ -16,9 +16,8 @@ use tokio_util::codec::Framed;
 use crate::Error;
 use crate::codec::{BackendMessage, BackendMessages, FrontendMessage, PostgresCodec};
 use crate::config::{self, AuthKeys, Config};
-use crate::connect_tls::connect_tls;
 use crate::maybe_tls_stream::MaybeTlsStream;
-use crate::tls::{TlsConnect, TlsStream};
+use crate::tls::TlsStream;
 
 pub struct StartupStream<S, T> {
     inner: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
@@ -87,16 +86,13 @@ pub struct RawConnection<S, T> {
 }
 
 pub async fn connect_raw<S, T>(
-    stream: S,
-    tls: T,
+    stream: MaybeTlsStream<S, T>,
     config: &Config,
-) -> Result<RawConnection<S, T::Stream>, Error>
+) -> Result<RawConnection<S, T>, Error>
 where
     S: AsyncRead + AsyncWrite + Unpin,
-    T: TlsConnect<S>,
+    T: TlsStream + Unpin,
 {
-    let stream = connect_tls(stream, config.ssl_mode, tls).await?;
-
     let mut stream = StartupStream {
         inner: Framed::new(stream, PostgresCodec),
         buf: BackendMessages::empty(),

libs/utils/src/lsn.rs

@@ -17,9 +17,6 @@ pub const XLOG_BLCKSZ: u32 = 8192;
 #[derive(Clone, Copy, Default, Eq, Ord, PartialEq, PartialOrd, Hash)]
 pub struct Lsn(pub u64);
 
-/// Size of a Postgres WAL segment.  These are always small enough to fit in a u32.
-pub type WalSegmentSize = u32;
-
 impl Serialize for Lsn {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where
@@ -166,19 +163,19 @@ impl Lsn {
 
     /// Compute the offset into a segment
     #[inline]
-    pub fn segment_offset(self, seg_sz: WalSegmentSize) -> WalSegmentSize {
-        (self.0 % seg_sz as u64) as WalSegmentSize
+    pub fn segment_offset(self, seg_sz: usize) -> usize {
+        (self.0 % seg_sz as u64) as usize
     }
 
     /// Compute LSN of the segment start.
     #[inline]
-    pub fn segment_lsn(self, seg_sz: WalSegmentSize) -> Lsn {
+    pub fn segment_lsn(self, seg_sz: usize) -> Lsn {
         Lsn(self.0 - (self.0 % seg_sz as u64))
     }
 
     /// Compute the segment number
     #[inline]
-    pub fn segment_number(self, seg_sz: WalSegmentSize) -> u64 {
+    pub fn segment_number(self, seg_sz: usize) -> u64 {
         self.0 / seg_sz as u64
     }
 
@@ -199,7 +196,7 @@ impl Lsn {
     /// Compute the block offset of the first byte of this Lsn within this
     /// segment
     #[inline]
-    pub fn page_offset_in_segment(self, seg_sz: WalSegmentSize) -> u64 {
+    pub fn page_offset_in_segment(self, seg_sz: usize) -> u64 {
         (self.0 - self.block_offset()) - self.segment_lsn(seg_sz).0
     }
 
@@ -466,7 +463,7 @@ mod tests {
         assert_eq!(Lsn(u64::MAX).widening_sub(0u64), i128::from(u64::MAX));
         assert_eq!(Lsn(0).widening_sub(u64::MAX), -i128::from(u64::MAX));
 
-        let seg_sz: WalSegmentSize = 16 * 1024 * 1024;
+        let seg_sz: usize = 16 * 1024 * 1024;
         assert_eq!(Lsn(0x1000007).segment_offset(seg_sz), 7);
         assert_eq!(Lsn(0x1000007).segment_number(seg_sz), 1u64);
 

libs/wal_decoder/benches/bench_interpret_wal.rs

@@ -226,9 +226,9 @@ fn decode_interpret_main(bench: &BenchmarkData, shards: &[ShardIdentity]) {
 
 fn decode_interpret(bench: &BenchmarkData, shard: &[ShardIdentity]) -> anyhow::Result<()> {
     let mut decoder = WalStreamDecoder::new(bench.meta.start_lsn, bench.meta.pg_version);
-    let xlogoff = bench.meta.start_lsn.segment_offset(WAL_SEGMENT_SIZE);
+    let xlogoff: usize = bench.meta.start_lsn.segment_offset(WAL_SEGMENT_SIZE);
 
-    for chunk in bench.wal[xlogoff as usize..].chunks(MAX_SEND_SIZE) {
+    for chunk in bench.wal[xlogoff..].chunks(MAX_SEND_SIZE) {
         decoder.feed_bytes(chunk);
         while let Some((lsn, recdata)) = decoder.poll_decode().unwrap() {
             assert!(lsn.is_aligned());

pageserver/src/basebackup.rs

@@ -29,7 +29,7 @@ use tokio::io;
 use tokio::io::AsyncWrite;
 use tokio_tar::{Builder, EntryType, Header};
 use tracing::*;
-use utils::lsn::{Lsn, WalSegmentSize};
+use utils::lsn::Lsn;
 
 use crate::context::RequestContext;
 use crate::pgdatadir_mapping::Version;
@@ -773,7 +773,7 @@ where
             self.lsn,
         )
         .map_err(|e| anyhow!(e).context("Failed generating wal segment"))?;
-        if WalSegmentSize::try_from(wal_seg.len()) != Ok(WAL_SEGMENT_SIZE) {
+        if wal_seg.len() != WAL_SEGMENT_SIZE {
             return Err(BasebackupError::Server(anyhow!(
                 "wal_seg.len() != WAL_SEGMENT_SIZE, wal_seg.len()={}",
                 wal_seg.len()

pageserver/src/import_datadir.rs

@@ -301,7 +301,7 @@ async fn import_wal(
 
         use std::io::Read;
         let nread = file.read_to_end(&mut buf)?;
-        if nread != WAL_SEGMENT_SIZE as usize - offset as usize {
+        if nread != WAL_SEGMENT_SIZE - offset {
             // Maybe allow this for .partial files?
             error!("read only {} bytes from WAL file", nread);
         }
@@ -455,7 +455,7 @@ pub async fn import_wal_from_tar(
             }
         };
 
-        waldecoder.feed_bytes(&bytes[offset as usize..]);
+        waldecoder.feed_bytes(&bytes[offset..]);
 
         let mut modification = tline.begin_modification(last_lsn);
         while last_lsn <= end_lsn {

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

@@ -275,12 +275,20 @@ 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 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 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 (format, compression) = match protocol {
         PostgresClientProtocol::Interpreted {

pageserver/src/walingest.rs

@@ -2383,17 +2383,17 @@ mod tests {
         let started_at = std::time::Instant::now();
 
         // Initialize walingest
-        let xlogoff = startpoint.segment_offset(WAL_SEGMENT_SIZE);
+        let xlogoff: usize = startpoint.segment_offset(WAL_SEGMENT_SIZE);
         let mut decoder = WalStreamDecoder::new(startpoint, pg_version);
         let mut walingest = WalIngest::new(tline.as_ref(), startpoint, &ctx)
             .await
             .unwrap();
         let mut modification = tline.begin_modification(startpoint);
-        println!("decoding {} bytes", bytes.len() - xlogoff as usize);
+        println!("decoding {} bytes", bytes.len() - xlogoff);
 
         // Decode and ingest wal. We process the wal in chunks because
         // that's what happens when we get bytes from safekeepers.
-        for chunk in bytes[xlogoff as usize..].chunks(50) {
+        for chunk in bytes[xlogoff..].chunks(50) {
             decoder.feed_bytes(chunk);
             while let Some((lsn, recdata)) = decoder.poll_decode().unwrap() {
                 let interpreted = InterpretedWalRecord::from_bytes_filtered(

proxy/src/batch.rs

@@ -6,7 +6,6 @@ use std::collections::BTreeMap;
 use std::pin::pin;
 use std::sync::Mutex;
 
-use futures::future::Either;
 use scopeguard::ScopeGuard;
 use tokio::sync::oneshot::error::TryRecvError;
 
@@ -49,37 +48,67 @@ impl<P: QueueProcessing> BatchQueue<P> {
         }
     }
 
-    pub async fn call(&self, req: P::Req) -> P::Res {
+    /// Perform a single request-response process, this may be batched internally.
+    ///
+    /// This function is not cancel safe.
+    pub async fn call<R>(
+        &self,
+        req: P::Req,
+        cancelled: impl Future<Output = R>,
+    ) -> Result<P::Res, R> {
         let (id, mut rx) = self.inner.lock_propagate_poison().register_job(req);
-        let guard = scopeguard::guard(id, move |id| {
-            let mut inner = self.inner.lock_propagate_poison();
-            if inner.queue.remove(&id).is_some() {
-                tracing::debug!("batched task cancelled before completion");
-            }
-        });
 
+        let mut cancelled = pin!(cancelled);
         let resp = loop {
             // try become the leader, or try wait for success.
-            let mut processor = match futures::future::select(rx, pin!(self.processor.lock())).await
-            {
-                // we got the resp.
-                Either::Left((resp, _)) => break resp.ok(),
-                // we are the leader.
-                Either::Right((p, rx_)) => {
-                    rx = rx_;
-                    p
-                }
+            let mut processor = tokio::select! {
+                // try become leader.
+                p = self.processor.lock() => p,
+                // wait for success.
+                resp = &mut rx => break resp.ok(),
+                // wait for cancellation.
+                cancel = cancelled.as_mut() => {
+                    let mut inner = self.inner.lock_propagate_poison();
+                    if inner.queue.remove(&id).is_some() {
+                        tracing::warn!("batched task cancelled before completion");
+                    }
+                    return Err(cancel);
+                },
             };
 
+            tracing::debug!(id, "batch: became leader");
             let (reqs, resps) = self.inner.lock_propagate_poison().get_batch(&processor);
 
+            // snitch incase the task gets cancelled.
+            let cancel_safety = scopeguard::guard((), |()| {
+                if !std::thread::panicking() {
+                    tracing::error!(
+                        id,
+                        "batch: leader cancelled, despite not being cancellation safe"
+                    );
+                }
+            });
+
             // apply a batch.
+            // if this is cancelled, jobs will not be completed and will panic.
             let values = processor.apply(reqs).await;
 
+            // good: we didn't get cancelled.
+            ScopeGuard::into_inner(cancel_safety);
+
+            if values.len() != resps.len() {
+                tracing::error!(
+                    "batch: invalid response size, expected={}, got={}",
+                    resps.len(),
+                    values.len()
+                );
+            }
+
             // send response values.
             for (tx, value) in std::iter::zip(resps, values) {
-                // sender hung up but that's fine.
-                drop(tx.send(value));
+                if tx.send(value).is_err() {
+                    // receiver hung up but that's fine.
+                }
             }
 
             match rx.try_recv() {
@@ -98,10 +127,9 @@ impl<P: QueueProcessing> BatchQueue<P> {
             }
         };
 
-        // already removed.
-        ScopeGuard::into_inner(guard);
+        tracing::debug!(id, "batch: job completed");
 
-        resp.expect("no response found. batch processer should not panic")
+        Ok(resp.expect("no response found. batch processer should not panic"))
     }
 }
 
@@ -125,6 +153,8 @@ impl<P: QueueProcessing> BatchQueueInner<P> {
 
         self.queue.insert(id, BatchJob { req, res: tx });
 
+        tracing::debug!(id, "batch: registered job in the queue");
+
         (id, rx)
     }
 
@@ -132,15 +162,19 @@ impl<P: QueueProcessing> BatchQueueInner<P> {
         let batch_size = p.batch_size(self.queue.len());
         let mut reqs = Vec::with_capacity(batch_size);
         let mut resps = Vec::with_capacity(batch_size);
+        let mut ids = Vec::with_capacity(batch_size);
 
         while reqs.len() < batch_size {
-            let Some((_, job)) = self.queue.pop_first() else {
+            let Some((id, job)) = self.queue.pop_first() else {
                 break;
             };
             reqs.push(job.req);
             resps.push(job.res);
+            ids.push(id);
         }
 
+        tracing::debug!(ids=?ids, "batch: acquired jobs");
+
         (reqs, resps)
     }
 }

proxy/src/cancellation.rs

@@ -1,5 +1,6 @@
 use std::convert::Infallible;
 use std::net::{IpAddr, SocketAddr};
+use std::pin::pin;
 use std::sync::{Arc, OnceLock};
 use std::time::Duration;
 
@@ -98,7 +99,6 @@ impl Pipeline {
 
 impl CancelKeyOp {
     fn register(&self, pipe: &mut Pipeline) {
-        #[allow(clippy::used_underscore_binding)]
         match self {
             CancelKeyOp::StoreCancelKey { key, value, expire } => {
                 let key = KeyPrefix::Cancel(*key).build_redis_key();
@@ -224,6 +224,7 @@ impl CancellationHandler {
         }
     }
 
+    /// This is not cancel safe
     async fn get_cancel_key(
         &self,
         key: CancelKeyData,
@@ -240,16 +241,21 @@ impl CancellationHandler {
         };
 
         const TIMEOUT: Duration = Duration::from_secs(5);
-        let result = timeout(TIMEOUT, tx.call((guard, op)))
-            .await
-            .map_err(|_| {
-                tracing::warn!("timed out waiting to receive GetCancelData response");
-                CancelError::RateLimit
-            })?
-            .map_err(|e| {
-                tracing::warn!("failed to receive GetCancelData response: {e}");
-                CancelError::InternalError
-            })?;
+        let result = timeout(
+            TIMEOUT,
+            tx.call((guard, op), std::future::pending::<Infallible>()),
+        )
+        .await
+        .map_err(|_| {
+            tracing::warn!("timed out waiting to receive GetCancelData response");
+            CancelError::RateLimit
+        })?
+        // cannot be cancelled
+        .unwrap_or_else(|x| match x {})
+        .map_err(|e| {
+            tracing::warn!("failed to receive GetCancelData response: {e}");
+            CancelError::InternalError
+        })?;
 
         let cancel_state_str = String::from_owned_redis_value(result).map_err(|e| {
             tracing::warn!("failed to receive GetCancelData response: {e}");
@@ -271,6 +277,8 @@ impl CancellationHandler {
     /// Will fetch IP allowlist internally.
     ///
     /// return Result primarily for tests
+    ///
+    /// This is not cancel safe
     pub(crate) async fn cancel_session<T: ControlPlaneApi>(
         &self,
         key: CancelKeyData,
@@ -342,7 +350,7 @@ impl CancellationHandler {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CancelClosure {
     socket_addr: SocketAddr,
-    cancel_token: RawCancelToken,
+    pub cancel_token: RawCancelToken,
     hostname: String, // for pg_sni router
     user_info: ComputeUserInfo,
 }
@@ -394,6 +402,8 @@ impl Session {
 
     /// Ensure the cancel key is continously refreshed,
     /// but stop when the channel is dropped.
+    ///
+    /// This is not cancel safe
     pub(crate) async fn maintain_cancel_key(
         &self,
         session_id: uuid::Uuid,
@@ -401,27 +411,6 @@ impl Session {
         cancel_closure: &CancelClosure,
         compute_config: &ComputeConfig,
     ) {
-        futures::future::select(
-            std::pin::pin!(self.maintain_redis_cancel_key(cancel_closure)),
-            cancel,
-        )
-        .await;
-
-        if let Err(err) = cancel_closure
-            .try_cancel_query(compute_config)
-            .boxed()
-            .await
-        {
-            tracing::warn!(
-                ?session_id,
-                ?err,
-                "could not cancel the query in the database"
-            );
-        }
-    }
-
-    // Ensure the cancel key is continously refreshed.
-    async fn maintain_redis_cancel_key(&self, cancel_closure: &CancelClosure) -> ! {
         let Some(tx) = self.cancellation_handler.tx.get() else {
             tracing::warn!("cancellation handler is not available");
             // don't exit, as we only want to exit if cancelled externally.
@@ -432,6 +421,8 @@ impl Session {
             .expect("serialising to json string should not fail")
             .into_boxed_str();
 
+        let mut cancel = pin!(cancel);
+
         loop {
             let guard = Metrics::get()
                 .proxy
@@ -449,9 +440,35 @@ impl Session {
                 "registering cancellation key"
             );
 
-            if tx.call((guard, op)).await.is_ok() {
-                tokio::time::sleep(CANCEL_KEY_REFRESH).await;
+            match tx.call((guard, op), cancel.as_mut()).await {
+                Ok(Ok(_)) => {
+                    tracing::debug!(
+                        src=%self.key,
+                        dest=?cancel_closure.cancel_token,
+                        "registered cancellation key"
+                    );
+
+                    // wait before continuing.
+                    tokio::time::sleep(CANCEL_KEY_REFRESH).await;
+                }
+                // retry immediately.
+                Ok(Err(error)) => {
+                    tracing::warn!(?error, "error registering cancellation key");
+                }
+                Err(Err(_cancelled)) => break,
             }
         }
+
+        if let Err(err) = cancel_closure
+            .try_cancel_query(compute_config)
+            .boxed()
+            .await
+        {
+            tracing::warn!(
+                ?session_id,
+                ?err,
+                "could not cancel the query in the database"
+            );
+        }
     }
 }

proxy/src/compute/mod.rs

@@ -6,7 +6,7 @@ use std::net::{IpAddr, SocketAddr};
 
 use futures::{FutureExt, TryFutureExt};
 use itertools::Itertools;
-use postgres_client::config::{AuthKeys, SslMode};
+use postgres_client::config::{AuthKeys, ChannelBinding, SslMode};
 use postgres_client::maybe_tls_stream::MaybeTlsStream;
 use postgres_client::tls::MakeTlsConnect;
 use postgres_client::{NoTls, RawCancelToken, RawConnection};
@@ -86,6 +86,14 @@ pub(crate) enum ConnectionError {
 
     #[error("error acquiring resource permit: {0}")]
     TooManyConnectionAttempts(#[from] ApiLockError),
+
+    #[cfg(test)]
+    #[error("retryable: {retryable}, wakeable: {wakeable}, kind: {kind:?}")]
+    TestError {
+        retryable: bool,
+        wakeable: bool,
+        kind: crate::error::ErrorKind,
+    },
 }
 
 impl UserFacingError for ConnectionError {
@@ -96,6 +104,8 @@ impl UserFacingError for ConnectionError {
                 "Failed to acquire permit to connect to the database. Too many database connection attempts are currently ongoing.".to_owned()
             }
             ConnectionError::TlsError(_) => COULD_NOT_CONNECT.to_owned(),
+            #[cfg(test)]
+            ConnectionError::TestError { .. } => self.to_string(),
         }
     }
 }
@@ -106,6 +116,8 @@ impl ReportableError for ConnectionError {
             ConnectionError::TlsError(_) => crate::error::ErrorKind::Compute,
             ConnectionError::WakeComputeError(e) => e.get_error_kind(),
             ConnectionError::TooManyConnectionAttempts(e) => e.get_error_kind(),
+            #[cfg(test)]
+            ConnectionError::TestError { kind, .. } => *kind,
         }
     }
 }
@@ -129,6 +141,8 @@ pub(crate) struct AuthInfo {
     auth: Option<Auth>,
     server_params: StartupMessageParams,
 
+    channel_binding: ChannelBinding,
+
     /// Console redirect sets user and database, we shouldn't re-use those from the params.
     skip_db_user: bool,
 }
@@ -152,6 +166,8 @@ impl AuthInfo {
             auth: pw.map(|pw| Auth::Password(pw.as_bytes().to_owned())),
             server_params,
             skip_db_user: true,
+            // pg-sni-router is a mitm so this would fail.
+            channel_binding: ChannelBinding::Disable,
         }
     }
 
@@ -165,6 +181,7 @@ impl AuthInfo {
             },
             server_params: StartupMessageParams::default(),
             skip_db_user: false,
+            channel_binding: ChannelBinding::Prefer,
         }
     }
 }
@@ -187,6 +204,7 @@ impl AuthInfo {
             Some(Auth::Password(pw)) => config.password(pw),
             None => &mut config,
         };
+        config.channel_binding(self.channel_binding);
         for (k, v) in self.server_params.iter() {
             config.set_param(k, v);
         }
@@ -241,9 +259,24 @@ impl AuthInfo {
         let tmp_config = self.enrich(tmp_config);
 
         let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
-        let connection = tmp_config.connect_raw(&mut compute.stream, NoTls).await?;
+        let connection = tmp_config
+            .tls_and_authenticate(&mut compute.stream, NoTls)
+            .await?;
         drop(pause);
 
+        // TODO: lots of useful info but maybe we can move it elsewhere (eg traces?)
+        info!(
+            compute_id = %compute.aux.compute_id,
+            pid = connection.process_id,
+            cold_start_info = ctx.cold_start_info().as_str(),
+            query_id = ctx.get_testodrome_id().as_deref(),
+            sslmode = ?compute.ssl_mode,
+            "connected to compute node at {} ({}) latency={}",
+            compute.hostname,
+            compute.socket_addr,
+            ctx.get_proxy_latency(),
+        );
+
         let RawConnection {
             stream: _,
             parameters,
@@ -252,8 +285,6 @@ impl AuthInfo {
             secret_key,
         } = connection;
 
-        tracing::Span::current().record("pid", tracing::field::display(process_id));
-
         // 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(
@@ -280,6 +311,7 @@ impl ConnectInfo {
     async fn connect_raw(
         &self,
         config: &ComputeConfig,
+        direct: bool,
     ) -> Result<(SocketAddr, MaybeTlsStream<TcpStream, RustlsStream>), TlsError> {
         let timeout = config.timeout;
 
@@ -322,7 +354,7 @@ impl ConnectInfo {
         match connect_once(&*addrs).await {
             Ok((sockaddr, stream)) => Ok((
                 sockaddr,
-                tls::connect_tls(stream, self.ssl_mode, config, host).await?,
+                tls::connect_tls(stream, self.ssl_mode, config, host, direct).await?,
             )),
             Err(err) => {
                 warn!("couldn't connect to compute node at {host}:{port}: {err}");
@@ -349,7 +381,7 @@ pub struct PostgresSettings {
 
 pub struct ComputeConnection {
     /// Socket connected to a compute node.
-    pub stream: MaybeTlsStream<tokio::net::TcpStream, RustlsStream>,
+    pub stream: MaybeRustlsStream,
     /// Labels for proxy's metrics.
     pub aux: MetricsAuxInfo,
     pub hostname: Host,
@@ -365,23 +397,12 @@ impl ConnectInfo {
         ctx: &RequestContext,
         aux: &MetricsAuxInfo,
         config: &ComputeConfig,
+        direct: bool,
     ) -> Result<ComputeConnection, ConnectionError> {
         let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
-        let (socket_addr, stream) = self.connect_raw(config).await?;
+        let (socket_addr, stream) = self.connect_raw(config, direct).await?;
         drop(pause);
 
-        tracing::Span::current().record("compute_id", tracing::field::display(&aux.compute_id));
-
-        // TODO: lots of useful info but maybe we can move it elsewhere (eg traces?)
-        info!(
-            cold_start_info = ctx.cold_start_info().as_str(),
-            "connected to compute node at {} ({socket_addr}) sslmode={:?}, latency={}, query_id={}",
-            self.host,
-            self.ssl_mode,
-            ctx.get_proxy_latency(),
-            ctx.get_testodrome_id().unwrap_or_default(),
-        );
-
         let connection = ComputeConnection {
             stream,
             socket_addr,

proxy/src/compute/tls.rs

@@ -11,8 +11,6 @@ use crate::proxy::retry::CouldRetry;
 
 #[derive(Debug, Error)]
 pub enum TlsError {
-    #[error(transparent)]
-    Dns(#[from] InvalidDnsNameError),
     #[error(transparent)]
     Connection(#[from] std::io::Error),
     #[error("TLS required but not provided")]
@@ -22,7 +20,6 @@ pub enum TlsError {
 impl CouldRetry for TlsError {
     fn could_retry(&self) -> bool {
         match self {
-            TlsError::Dns(_) => false,
             TlsError::Connection(err) => err.could_retry(),
             // perhaps compute didn't realise it supports TLS?
             TlsError::Required => true,
@@ -35,6 +32,7 @@ pub async fn connect_tls<S, T>(
     mode: SslMode,
     tls: &T,
     host: &str,
+    direct: bool,
 ) -> Result<MaybeTlsStream<S, T::Stream>, TlsError>
 where
     S: AsyncRead + AsyncWrite + Unpin + Send,
@@ -49,7 +47,7 @@ where
         SslMode::Prefer | SslMode::Require => {}
     }
 
-    if !request_tls(&mut stream).await? {
+    if !direct && !request_tls(&mut stream).await? {
         if SslMode::Require == mode {
             return Err(TlsError::Required);
         }
@@ -57,7 +55,6 @@ where
         return Ok(MaybeTlsStream::Raw(stream));
     }
 
-    Ok(MaybeTlsStream::Tls(
-        tls.make_tls_connect(host)?.connect(stream).boxed().await?,
-    ))
+    let c = tls.make_tls_connect(host).map_err(std::io::Error::other)?;
+    Ok(MaybeTlsStream::Tls(c.connect(stream).boxed().await?))
 }

proxy/src/console_redirect_proxy.rs

@@ -222,6 +222,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
         ctx,
         &TcpMechanism {
             locks: &config.connect_compute_locks,
+            direct: false,
         },
         &node_info,
         config.wake_compute_retry_config,

proxy/src/control_plane/client/cplane_proxy_v1.rs

@@ -263,7 +263,12 @@ impl NeonControlPlaneClient {
                 None => SslMode::Disable,
             };
             let host = match body.server_name {
-                Some(host) => host.into(),
+                Some(host) => {
+                    if rustls::pki_types::DnsName::try_from_str(&host).is_err() {
+                        return Err(WakeComputeError::BadComputeAddress(host.into_boxed_str()));
+                    }
+                    host.into()
+                }
                 None => host.into(),
             };
 

proxy/src/control_plane/mod.rs

@@ -77,8 +77,9 @@ impl NodeInfo {
         &self,
         ctx: &RequestContext,
         config: &ComputeConfig,
+        direct: bool,
     ) -> Result<compute::ComputeConnection, compute::ConnectionError> {
-        self.conn_info.connect(ctx, &self.aux, config).await
+        self.conn_info.connect(ctx, &self.aux, config, direct).await
     }
 }
 

proxy/src/proxy/connect_compute.rs

@@ -1,18 +1,15 @@
-use async_trait::async_trait;
 use tokio::time;
 use tracing::{debug, info, warn};
 
 use crate::compute::{self, COULD_NOT_CONNECT, ComputeConnection};
 use crate::config::{ComputeConfig, RetryConfig};
 use crate::context::RequestContext;
-use crate::control_plane::errors::WakeComputeError;
 use crate::control_plane::locks::ApiLocks;
 use crate::control_plane::{self, NodeInfo};
-use crate::error::ReportableError;
 use crate::metrics::{
     ConnectOutcome, ConnectionFailureKind, Metrics, RetriesMetricGroup, RetryType,
 };
-use crate::proxy::retry::{CouldRetry, ShouldRetryWakeCompute, retry_after, should_retry};
+use crate::proxy::retry::{ShouldRetryWakeCompute, retry_after, should_retry};
 use crate::proxy::wake_compute::{WakeComputeBackend, wake_compute};
 use crate::types::Host;
 
@@ -35,42 +32,34 @@ pub(crate) fn invalidate_cache(node_info: control_plane::CachedNodeInfo) -> Node
     node_info.invalidate()
 }
 
-#[async_trait]
 pub(crate) trait ConnectMechanism {
     type Connection;
-    type ConnectError: ReportableError;
-    type Error: From<Self::ConnectError>;
     async fn connect_once(
         &self,
         ctx: &RequestContext,
         node_info: &control_plane::CachedNodeInfo,
         config: &ComputeConfig,
-    ) -> Result<Self::Connection, Self::ConnectError>;
+    ) -> Result<Self::Connection, compute::ConnectionError>;
 }
 
 pub(crate) struct TcpMechanism {
     /// connect_to_compute concurrency lock
     pub(crate) locks: &'static ApiLocks<Host>,
+    // whether to negotiate TLS for postgres protocol.
+    pub(crate) direct: bool,
 }
 
-#[async_trait]
 impl ConnectMechanism for TcpMechanism {
     type Connection = ComputeConnection;
-    type ConnectError = compute::ConnectionError;
-    type Error = compute::ConnectionError;
 
-    #[tracing::instrument(skip_all, fields(
-        pid = tracing::field::Empty,
-        compute_id = tracing::field::Empty
-    ))]
     async fn connect_once(
         &self,
         ctx: &RequestContext,
         node_info: &control_plane::CachedNodeInfo,
         config: &ComputeConfig,
-    ) -> Result<ComputeConnection, Self::Error> {
+    ) -> Result<ComputeConnection, compute::ConnectionError> {
         let permit = self.locks.get_permit(&node_info.conn_info.host).await?;
-        permit.release_result(node_info.connect(ctx, config).await)
+        permit.release_result(node_info.connect(ctx, config, self.direct).await)
     }
 }
 
@@ -82,11 +71,7 @@ pub(crate) async fn connect_to_compute<M: ConnectMechanism, B: WakeComputeBacken
     user_info: &B,
     wake_compute_retry_config: RetryConfig,
     compute: &ComputeConfig,
-) -> Result<M::Connection, M::Error>
-where
-    M::ConnectError: CouldRetry + ShouldRetryWakeCompute + std::fmt::Debug,
-    M::Error: From<WakeComputeError>,
-{
+) -> Result<M::Connection, compute::ConnectionError> {
     let mut num_retries = 0;
     let node_info =
         wake_compute(&mut num_retries, ctx, user_info, wake_compute_retry_config).await?;
@@ -120,7 +105,7 @@ where
                 },
                 num_retries.into(),
             );
-            return Err(err.into());
+            return Err(err);
         }
         node_info
     } else {
@@ -161,7 +146,7 @@ where
                         },
                         num_retries.into(),
                     );
-                    return Err(e.into());
+                    return Err(e);
                 }
 
                 warn!(error = ?e, num_retries, retriable = true, COULD_NOT_CONNECT);

proxy/src/proxy/mod.rs

@@ -358,6 +358,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
         ctx,
         &TcpMechanism {
             locks: &config.connect_compute_locks,
+            direct: false,
         },
         &auth::Backend::ControlPlane(cplane, creds.info.clone()),
         config.wake_compute_retry_config,

proxy/src/proxy/retry.rs

@@ -1,4 +1,3 @@
-use std::error::Error;
 use std::io;
 
 use tokio::time;
@@ -31,85 +30,25 @@ impl CouldRetry for io::Error {
     }
 }
 
-impl CouldRetry for postgres_client::error::DbError {
-    fn could_retry(&self) -> bool {
-        use postgres_client::error::SqlState;
-        matches!(
-            self.code(),
-            &SqlState::CONNECTION_FAILURE
-                | &SqlState::CONNECTION_EXCEPTION
-                | &SqlState::CONNECTION_DOES_NOT_EXIST
-                | &SqlState::SQLCLIENT_UNABLE_TO_ESTABLISH_SQLCONNECTION,
-        )
-    }
-}
-impl ShouldRetryWakeCompute for postgres_client::error::DbError {
-    fn should_retry_wake_compute(&self) -> bool {
-        use postgres_client::error::SqlState;
-        // Here are errors that happens after the user successfully authenticated to the database.
-        // TODO: there are pgbouncer errors that should be retried, but they are not listed here.
-        let non_retriable_pg_errors = matches!(
-            self.code(),
-            &SqlState::TOO_MANY_CONNECTIONS
-                | &SqlState::OUT_OF_MEMORY
-                | &SqlState::SYNTAX_ERROR
-                | &SqlState::T_R_SERIALIZATION_FAILURE
-                | &SqlState::INVALID_CATALOG_NAME
-                | &SqlState::INVALID_SCHEMA_NAME
-                | &SqlState::INVALID_PARAMETER_VALUE,
-        );
-        if non_retriable_pg_errors {
-            return false;
-        }
-        // PGBouncer errors that should not trigger a wake_compute retry.
-        if self.code() == &SqlState::PROTOCOL_VIOLATION {
-            // Source for the error message:
-            // https://github.com/pgbouncer/pgbouncer/blob/f15997fe3effe3a94ba8bcc1ea562e6117d1a131/src/client.c#L1070
-            return !self
-                .message()
-                .contains("no more connections allowed (max_client_conn)");
-        }
-        true
-    }
-}
-
-impl CouldRetry for postgres_client::Error {
-    fn could_retry(&self) -> bool {
-        if let Some(io_err) = self.source().and_then(|x| x.downcast_ref()) {
-            io::Error::could_retry(io_err)
-        } else if let Some(db_err) = self.source().and_then(|x| x.downcast_ref()) {
-            postgres_client::error::DbError::could_retry(db_err)
-        } else {
-            false
-        }
-    }
-}
-impl ShouldRetryWakeCompute for postgres_client::Error {
-    fn should_retry_wake_compute(&self) -> bool {
-        if let Some(db_err) = self.source().and_then(|x| x.downcast_ref()) {
-            postgres_client::error::DbError::should_retry_wake_compute(db_err)
-        } else {
-            // likely an IO error. Possible the compute has shutdown and the
-            // cache is stale.
-            true
-        }
-    }
-}
-
 impl CouldRetry for compute::ConnectionError {
     fn could_retry(&self) -> bool {
         match self {
             compute::ConnectionError::TlsError(err) => err.could_retry(),
             compute::ConnectionError::WakeComputeError(err) => err.could_retry(),
             compute::ConnectionError::TooManyConnectionAttempts(_) => false,
+            #[cfg(test)]
+            compute::ConnectionError::TestError { retryable, .. } => *retryable,
         }
     }
 }
+
 impl ShouldRetryWakeCompute for compute::ConnectionError {
     fn should_retry_wake_compute(&self) -> bool {
         match self {
             // the cache entry was not checked for validity
             compute::ConnectionError::TooManyConnectionAttempts(_) => false,
+            #[cfg(test)]
+            compute::ConnectionError::TestError { wakeable, .. } => *wakeable,
             _ => true,
         }
     }
@@ -120,56 +59,3 @@ pub(crate) fn retry_after(num_retries: u32, config: RetryConfig) -> time::Durati
         .base_delay
         .mul_f64(config.backoff_factor.powi((num_retries as i32) - 1))
 }
-
-#[cfg(test)]
-mod tests {
-    use postgres_client::error::{DbError, SqlState};
-
-    use super::ShouldRetryWakeCompute;
-
-    #[test]
-    fn should_retry_wake_compute_for_db_error() {
-        // These SQLStates should NOT trigger a wake_compute retry.
-        let non_retry_states = [
-            SqlState::TOO_MANY_CONNECTIONS,
-            SqlState::OUT_OF_MEMORY,
-            SqlState::SYNTAX_ERROR,
-            SqlState::T_R_SERIALIZATION_FAILURE,
-            SqlState::INVALID_CATALOG_NAME,
-            SqlState::INVALID_SCHEMA_NAME,
-            SqlState::INVALID_PARAMETER_VALUE,
-        ];
-        for state in non_retry_states {
-            let err = DbError::new_test_error(state.clone(), "oops".to_string());
-            assert!(
-                !err.should_retry_wake_compute(),
-                "State {state:?} unexpectedly retried"
-            );
-        }
-
-        // Errors coming from pgbouncer should not trigger a wake_compute retry
-        let non_retry_pgbouncer_errors = ["no more connections allowed (max_client_conn)"];
-        for error in non_retry_pgbouncer_errors {
-            let err = DbError::new_test_error(SqlState::PROTOCOL_VIOLATION, error.to_string());
-            assert!(
-                !err.should_retry_wake_compute(),
-                "PGBouncer error {error:?} unexpectedly retried"
-            );
-        }
-
-        // These SQLStates should trigger a wake_compute retry.
-        let retry_states = [
-            SqlState::CONNECTION_FAILURE,
-            SqlState::CONNECTION_EXCEPTION,
-            SqlState::CONNECTION_DOES_NOT_EXIST,
-            SqlState::SQLCLIENT_UNABLE_TO_ESTABLISH_SQLCONNECTION,
-        ];
-        for state in retry_states {
-            let err = DbError::new_test_error(state.clone(), "oops".to_string());
-            assert!(
-                err.should_retry_wake_compute(),
-                "State {state:?} unexpectedly skipped retry"
-            );
-        }
-    }
-}

proxy/src/proxy/tests/mitm.rs

@@ -169,7 +169,7 @@ async fn scram_auth_disable_channel_binding() -> anyhow::Result<()> {
         .dbname("db")
         .password("password")
         .ssl_mode(SslMode::Require)
-        .connect_raw(server, client_config.make_tls_connect()?)
+        .tls_and_authenticate(server, client_config.make_tls_connect()?)
         .await?;
 
     proxy.await?
@@ -252,7 +252,7 @@ async fn connect_failure(
         .dbname("db")
         .password("password")
         .ssl_mode(SslMode::Require)
-        .connect_raw(server, client_config.make_tls_connect()?)
+        .tls_and_authenticate(server, client_config.make_tls_connect()?)
         .await
         .err()
         .context("client shouldn't be able to connect")?;

proxy/src/proxy/tests/mod.rs

@@ -10,7 +10,7 @@ use async_trait::async_trait;
 use http::StatusCode;
 use postgres_client::config::SslMode;
 use postgres_client::tls::{MakeTlsConnect, NoTls};
-use retry::{ShouldRetryWakeCompute, retry_after};
+use retry::retry_after;
 use rstest::rstest;
 use rustls::crypto::ring;
 use rustls::pki_types;
@@ -20,6 +20,7 @@ use tracing_test::traced_test;
 use super::retry::CouldRetry;
 use super::*;
 use crate::auth::backend::{ComputeUserInfo, MaybeOwned};
+use crate::compute::ConnectionError;
 use crate::config::{ComputeConfig, RetryConfig};
 use crate::control_plane::client::{ControlPlaneClient, TestControlPlaneClient};
 use crate::control_plane::messages::{ControlPlaneErrorMessage, Details, MetricsAuxInfo, Status};
@@ -199,7 +200,7 @@ async fn handshake_tls_is_enforced_by_proxy() -> anyhow::Result<()> {
         .user("john_doe")
         .dbname("earth")
         .ssl_mode(SslMode::Disable)
-        .connect_raw(server, NoTls)
+        .tls_and_authenticate(server, NoTls)
         .await
         .err() // -> Option<E>
         .context("client shouldn't be able to connect")?;
@@ -228,7 +229,7 @@ async fn handshake_tls() -> anyhow::Result<()> {
         .user("john_doe")
         .dbname("earth")
         .ssl_mode(SslMode::Require)
-        .connect_raw(server, client_config.make_tls_connect()?)
+        .tls_and_authenticate(server, client_config.make_tls_connect()?)
         .await?;
 
     proxy.await?
@@ -245,7 +246,7 @@ async fn handshake_raw() -> anyhow::Result<()> {
         .dbname("earth")
         .set_param("options", "project=generic-project-name")
         .ssl_mode(SslMode::Prefer)
-        .connect_raw(server, NoTls)
+        .tls_and_authenticate(server, NoTls)
         .await?;
 
     proxy.await?
@@ -293,7 +294,7 @@ async fn scram_auth_good(#[case] password: &str) -> anyhow::Result<()> {
         .dbname("db")
         .password(password)
         .ssl_mode(SslMode::Require)
-        .connect_raw(server, client_config.make_tls_connect()?)
+        .tls_and_authenticate(server, client_config.make_tls_connect()?)
         .await?;
 
     proxy.await?
@@ -317,7 +318,7 @@ async fn scram_auth_disable_channel_binding() -> anyhow::Result<()> {
         .dbname("db")
         .password("password")
         .ssl_mode(SslMode::Require)
-        .connect_raw(server, client_config.make_tls_connect()?)
+        .tls_and_authenticate(server, client_config.make_tls_connect()?)
         .await?;
 
     proxy.await?
@@ -344,7 +345,7 @@ async fn scram_auth_mock() -> anyhow::Result<()> {
         .dbname("db")
         .password(&password) // no password will match the mocked secret
         .ssl_mode(SslMode::Require)
-        .connect_raw(server, client_config.make_tls_connect()?)
+        .tls_and_authenticate(server, client_config.make_tls_connect()?)
         .await
         .err() // -> Option<E>
         .context("client shouldn't be able to connect")?;
@@ -423,71 +424,36 @@ impl TestConnectMechanism {
 #[derive(Debug)]
 struct TestConnection;
 
-#[derive(Debug)]
-struct TestConnectError {
-    retryable: bool,
-    wakeable: bool,
-    kind: crate::error::ErrorKind,
-}
-
-impl ReportableError for TestConnectError {
-    fn get_error_kind(&self) -> crate::error::ErrorKind {
-        self.kind
-    }
-}
-
-impl std::fmt::Display for TestConnectError {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{self:?}")
-    }
-}
-
-impl std::error::Error for TestConnectError {}
-
-impl CouldRetry for TestConnectError {
-    fn could_retry(&self) -> bool {
-        self.retryable
-    }
-}
-impl ShouldRetryWakeCompute for TestConnectError {
-    fn should_retry_wake_compute(&self) -> bool {
-        self.wakeable
-    }
-}
-
-#[async_trait]
 impl ConnectMechanism for TestConnectMechanism {
     type Connection = TestConnection;
-    type ConnectError = TestConnectError;
-    type Error = anyhow::Error;
 
     async fn connect_once(
         &self,
         _ctx: &RequestContext,
         _node_info: &control_plane::CachedNodeInfo,
         _config: &ComputeConfig,
-    ) -> Result<Self::Connection, Self::ConnectError> {
+    ) -> Result<Self::Connection, ConnectionError> {
         let mut counter = self.counter.lock().unwrap();
         let action = self.sequence[*counter];
         *counter += 1;
         match action {
             ConnectAction::Connect => Ok(TestConnection),
-            ConnectAction::Retry => Err(TestConnectError {
+            ConnectAction::Retry => Err(ConnectionError::TestError {
                 retryable: true,
                 wakeable: true,
                 kind: ErrorKind::Compute,
             }),
-            ConnectAction::RetryNoWake => Err(TestConnectError {
+            ConnectAction::RetryNoWake => Err(ConnectionError::TestError {
                 retryable: true,
                 wakeable: false,
                 kind: ErrorKind::Compute,
             }),
-            ConnectAction::Fail => Err(TestConnectError {
+            ConnectAction::Fail => Err(ConnectionError::TestError {
                 retryable: false,
                 wakeable: true,
                 kind: ErrorKind::Compute,
             }),
-            ConnectAction::FailNoWake => Err(TestConnectError {
+            ConnectAction::FailNoWake => Err(ConnectionError::TestError {
                 retryable: false,
                 wakeable: false,
                 kind: ErrorKind::Compute,

proxy/src/redis/keys.rs

@@ -23,9 +23,8 @@ impl KeyPrefix {
 
 #[cfg(test)]
 mod tests {
-    use crate::pqproto::id_to_cancel_key;
-
     use super::*;
+    use crate::pqproto::id_to_cancel_key;
 
     #[test]
     fn test_build_redis_key() {

proxy/src/serverless/backend.rs

@@ -1,17 +1,12 @@
-use std::io;
-use std::net::{IpAddr, SocketAddr};
 use std::sync::Arc;
 use std::time::Duration;
 
-use async_trait::async_trait;
 use ed25519_dalek::SigningKey;
 use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
 use jose_jwk::jose_b64;
-use postgres_client::config::SslMode;
+use postgres_client::SocketConfig;
+use postgres_client::maybe_tls_stream::MaybeTlsStream;
 use rand::rngs::OsRng;
-use rustls::pki_types::{DnsName, ServerName};
-use tokio::net::{TcpStream, lookup_host};
-use tokio_rustls::TlsConnector;
 use tracing::field::display;
 use tracing::{debug, info};
 
@@ -23,21 +18,19 @@ use super::local_conn_pool::{self, EXT_NAME, EXT_SCHEMA, EXT_VERSION, LocalConnP
 use crate::auth::backend::local::StaticAuthRules;
 use crate::auth::backend::{ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo};
 use crate::auth::{self, AuthError};
+use crate::compute::{self, ComputeConnection};
 use crate::compute_ctl::{
     ComputeCtlError, ExtensionInstallRequest, Privilege, SetRoleGrantsRequest,
 };
-use crate::config::{ComputeConfig, ProxyConfig};
+use crate::config::ProxyConfig;
 use crate::context::RequestContext;
-use crate::control_plane::CachedNodeInfo;
 use crate::control_plane::client::ApiLockError;
 use crate::control_plane::errors::{GetAuthInfoError, WakeComputeError};
-use crate::control_plane::locks::ApiLocks;
 use crate::error::{ErrorKind, ReportableError, UserFacingError};
 use crate::intern::EndpointIdInt;
-use crate::proxy::connect_compute::ConnectMechanism;
-use crate::proxy::retry::{CouldRetry, ShouldRetryWakeCompute};
+use crate::proxy::connect_compute::TcpMechanism;
 use crate::rate_limiter::EndpointRateLimiter;
-use crate::types::{EndpointId, Host, LOCAL_PROXY_SUFFIX};
+use crate::types::{EndpointId, LOCAL_PROXY_SUFFIX};
 
 pub(crate) struct PoolingBackend {
     pub(crate) http_conn_pool: Arc<GlobalConnPool<Send, HttpConnPool<Send>>>,
@@ -157,11 +150,6 @@ impl PoolingBackend {
     // Wake up the destination if needed. Code here is a bit involved because
     // we reuse the code from the usual proxy and we need to prepare few structures
     // that this code expects.
-    #[tracing::instrument(skip_all, fields(
-        pid = tracing::field::Empty,
-        compute_id = tracing::field::Empty,
-        conn_id = tracing::field::Empty,
-    ))]
     pub(crate) async fn connect_to_compute(
         &self,
         ctx: &RequestContext,
@@ -181,30 +169,24 @@ impl PoolingBackend {
             return Ok(client);
         }
         let conn_id = uuid::Uuid::new_v4();
-        tracing::Span::current().record("conn_id", display(conn_id));
         info!(%conn_id, "pool: opening a new connection '{conn_info}'");
         let backend = self.auth_backend.as_ref().map(|()| keys.info);
-        crate::proxy::connect_compute::connect_to_compute(
+        let connection = crate::proxy::connect_compute::connect_to_compute(
             ctx,
-            &TokioMechanism {
-                conn_id,
-                conn_info,
-                pool: self.pool.clone(),
+            &TcpMechanism {
                 locks: &self.config.connect_compute_locks,
-                keys: keys.keys,
+                direct: false,
             },
             &backend,
             self.config.wake_compute_retry_config,
             &self.config.connect_to_compute,
         )
-        .await
+        .await?;
+
+        authenticate(ctx, &self.pool, &conn_info, keys.keys, connection, conn_id).await
     }
 
     // Wake up the destination if needed
-    #[tracing::instrument(skip_all, fields(
-        compute_id = tracing::field::Empty,
-        conn_id = tracing::field::Empty,
-    ))]
     pub(crate) async fn connect_to_local_proxy(
         &self,
         ctx: &RequestContext,
@@ -216,7 +198,6 @@ impl PoolingBackend {
         }
 
         let conn_id = uuid::Uuid::new_v4();
-        tracing::Span::current().record("conn_id", display(conn_id));
         debug!(%conn_id, "pool: opening a new connection '{conn_info}'");
         let backend = self.auth_backend.as_ref().map(|()| ComputeUserInfo {
             user: conn_info.user_info.user.clone(),
@@ -226,19 +207,19 @@ impl PoolingBackend {
             )),
             options: conn_info.user_info.options.clone(),
         });
-        crate::proxy::connect_compute::connect_to_compute(
+        let connection = crate::proxy::connect_compute::connect_to_compute(
             ctx,
-            &HyperMechanism {
-                conn_id,
-                conn_info,
-                pool: self.http_conn_pool.clone(),
+            &TcpMechanism {
                 locks: &self.config.connect_compute_locks,
+                direct: true,
             },
             &backend,
             self.config.wake_compute_retry_config,
             &self.config.connect_to_compute,
         )
-        .await
+        .await?;
+
+        h2handshake(ctx, &self.http_conn_pool, &conn_info, connection, conn_id).await
     }
 
     /// Connect to postgres over localhost.
@@ -248,10 +229,6 @@ impl PoolingBackend {
     /// # Panics
     ///
     /// Panics if called with a non-local_proxy backend.
-    #[tracing::instrument(skip_all, fields(
-        pid = tracing::field::Empty,
-        conn_id = tracing::field::Empty,
-    ))]
     pub(crate) async fn connect_to_local_postgres(
         &self,
         ctx: &RequestContext,
@@ -373,6 +350,8 @@ fn create_random_jwk() -> (SigningKey, jose_jwk::Key) {
 pub(crate) enum HttpConnError {
     #[error("pooled connection closed at inconsistent state")]
     ConnectionClosedAbruptly(#[from] tokio::sync::watch::error::SendError<uuid::Uuid>),
+    #[error("could not connect to compute")]
+    ConnectError(#[from] compute::ConnectionError),
     #[error("could not connect to postgres in compute")]
     PostgresConnectionError(#[from] postgres_client::Error),
     #[error("could not connect to local-proxy in compute")]
@@ -394,8 +373,6 @@ pub(crate) enum HttpConnError {
 
 #[derive(Debug, thiserror::Error)]
 pub(crate) enum LocalProxyConnError {
-    #[error("error with connection to local-proxy")]
-    Io(#[source] std::io::Error),
     #[error("could not establish h2 connection")]
     H2(#[from] hyper::Error),
 }
@@ -403,6 +380,7 @@ pub(crate) enum LocalProxyConnError {
 impl ReportableError for HttpConnError {
     fn get_error_kind(&self) -> ErrorKind {
         match self {
+            HttpConnError::ConnectError(_) => ErrorKind::Compute,
             HttpConnError::ConnectionClosedAbruptly(_) => ErrorKind::Compute,
             HttpConnError::PostgresConnectionError(p) => p.get_error_kind(),
             HttpConnError::LocalProxyConnectionError(_) => ErrorKind::Compute,
@@ -419,6 +397,7 @@ impl ReportableError for HttpConnError {
 impl UserFacingError for HttpConnError {
     fn to_string_client(&self) -> String {
         match self {
+            HttpConnError::ConnectError(p) => p.to_string_client(),
             HttpConnError::ConnectionClosedAbruptly(_) => self.to_string(),
             HttpConnError::PostgresConnectionError(p) => p.to_string(),
             HttpConnError::LocalProxyConnectionError(p) => p.to_string(),
@@ -434,36 +413,9 @@ impl UserFacingError for HttpConnError {
     }
 }
 
-impl CouldRetry for HttpConnError {
-    fn could_retry(&self) -> bool {
-        match self {
-            HttpConnError::PostgresConnectionError(e) => e.could_retry(),
-            HttpConnError::LocalProxyConnectionError(e) => e.could_retry(),
-            HttpConnError::ComputeCtl(_) => false,
-            HttpConnError::ConnectionClosedAbruptly(_) => false,
-            HttpConnError::JwtPayloadError(_) => false,
-            HttpConnError::GetAuthInfo(_) => false,
-            HttpConnError::AuthError(_) => false,
-            HttpConnError::WakeCompute(_) => false,
-            HttpConnError::TooManyConnectionAttempts(_) => false,
-        }
-    }
-}
-impl ShouldRetryWakeCompute for HttpConnError {
-    fn should_retry_wake_compute(&self) -> bool {
-        match self {
-            HttpConnError::PostgresConnectionError(e) => e.should_retry_wake_compute(),
-            // we never checked cache validity
-            HttpConnError::TooManyConnectionAttempts(_) => false,
-            _ => true,
-        }
-    }
-}
-
 impl ReportableError for LocalProxyConnError {
     fn get_error_kind(&self) -> ErrorKind {
         match self {
-            LocalProxyConnError::Io(_) => ErrorKind::Compute,
             LocalProxyConnError::H2(_) => ErrorKind::Compute,
         }
     }
@@ -475,208 +427,106 @@ impl UserFacingError for LocalProxyConnError {
     }
 }
 
-impl CouldRetry for LocalProxyConnError {
-    fn could_retry(&self) -> bool {
-        match self {
-            LocalProxyConnError::Io(_) => false,
-            LocalProxyConnError::H2(_) => false,
-        }
-    }
-}
-impl ShouldRetryWakeCompute for LocalProxyConnError {
-    fn should_retry_wake_compute(&self) -> bool {
-        match self {
-            LocalProxyConnError::Io(_) => false,
-            LocalProxyConnError::H2(_) => false,
-        }
-    }
-}
-
-struct TokioMechanism {
-    pool: Arc<GlobalConnPool<postgres_client::Client, EndpointConnPool<postgres_client::Client>>>,
-    conn_info: ConnInfo,
-    conn_id: uuid::Uuid,
+async fn authenticate(
+    ctx: &RequestContext,
+    pool: &Arc<GlobalConnPool<postgres_client::Client, EndpointConnPool<postgres_client::Client>>>,
+    conn_info: &ConnInfo,
     keys: ComputeCredentialKeys,
-
-    /// connect_to_compute concurrency lock
-    locks: &'static ApiLocks<Host>,
-}
-
-#[async_trait]
-impl ConnectMechanism for TokioMechanism {
-    type Connection = Client<postgres_client::Client>;
-    type ConnectError = HttpConnError;
-    type Error = HttpConnError;
-
-    async fn connect_once(
-        &self,
-        ctx: &RequestContext,
-        node_info: &CachedNodeInfo,
-        compute_config: &ComputeConfig,
-    ) -> Result<Self::Connection, Self::ConnectError> {
-        let permit = self.locks.get_permit(&node_info.conn_info.host).await?;
-
-        let mut config = node_info.conn_info.to_postgres_client_config();
-        let config = config
-            .user(&self.conn_info.user_info.user)
-            .dbname(&self.conn_info.dbname)
-            .connect_timeout(compute_config.timeout);
-
-        if let ComputeCredentialKeys::AuthKeys(auth_keys) = self.keys {
-            config.auth_keys(auth_keys);
-        }
-
-        let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
-        let res = config.connect(compute_config).await;
-        drop(pause);
-        let (client, connection) = permit.release_result(res)?;
-
-        tracing::Span::current().record("pid", tracing::field::display(client.get_process_id()));
-        tracing::Span::current().record(
-            "compute_id",
-            tracing::field::display(&node_info.aux.compute_id),
-        );
-
-        if let Some(query_id) = ctx.get_testodrome_id() {
-            info!("latency={}, query_id={}", ctx.get_proxy_latency(), query_id);
-        }
-
-        Ok(poll_client(
-            self.pool.clone(),
-            ctx,
-            self.conn_info.clone(),
-            client,
-            connection,
-            self.conn_id,
-            node_info.aux.clone(),
-        ))
-    }
-}
-
-struct HyperMechanism {
-    pool: Arc<GlobalConnPool<Send, HttpConnPool<Send>>>,
-    conn_info: ConnInfo,
+    compute: ComputeConnection,
     conn_id: uuid::Uuid,
+) -> Result<Client<postgres_client::Client>, HttpConnError> {
+    // client config with stubbed connect info.
+    let mut config = postgres_client::Config::new(String::new(), 0);
+    config
+        .user(&conn_info.user_info.user)
+        .dbname(&conn_info.dbname);
 
-    /// connect_to_compute concurrency lock
-    locks: &'static ApiLocks<Host>,
-}
-
-#[async_trait]
-impl ConnectMechanism for HyperMechanism {
-    type Connection = http_conn_pool::Client<Send>;
-    type ConnectError = HttpConnError;
-    type Error = HttpConnError;
-
-    async fn connect_once(
-        &self,
-        ctx: &RequestContext,
-        node_info: &CachedNodeInfo,
-        config: &ComputeConfig,
-    ) -> Result<Self::Connection, Self::ConnectError> {
-        let host_addr = node_info.conn_info.host_addr;
-        let host = &node_info.conn_info.host;
-        let permit = self.locks.get_permit(host).await?;
-
-        let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
-
-        let tls = if node_info.conn_info.ssl_mode == SslMode::Disable {
-            None
-        } else {
-            Some(&config.tls)
-        };
-
-        let port = node_info.conn_info.port;
-        let res = connect_http2(host_addr, host, port, config.timeout, tls).await;
-        drop(pause);
-        let (client, connection) = permit.release_result(res)?;
-
-        tracing::Span::current().record(
-            "compute_id",
-            tracing::field::display(&node_info.aux.compute_id),
-        );
-
-        if let Some(query_id) = ctx.get_testodrome_id() {
-            info!("latency={}, query_id={}", ctx.get_proxy_latency(), query_id);
-        }
-
-        Ok(poll_http2_client(
-            self.pool.clone(),
-            ctx,
-            &self.conn_info,
-            client,
-            connection,
-            self.conn_id,
-            node_info.aux.clone(),
-        ))
+    if let ComputeCredentialKeys::AuthKeys(auth_keys) = keys {
+        config.auth_keys(auth_keys);
     }
+
+    let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
+    let connection = config.authenticate(compute.stream).await?;
+    drop(pause);
+
+    // TODO: lots of useful info but maybe we can move it elsewhere (eg traces?)
+    info!(
+        compute_id = %compute.aux.compute_id,
+        pid = connection.process_id,
+        cold_start_info = ctx.cold_start_info().as_str(),
+        query_id = ctx.get_testodrome_id().as_deref(),
+        sslmode = ?compute.ssl_mode,
+        %conn_id,
+        "connected to compute node at {} ({}) latency={}",
+        compute.hostname,
+        compute.socket_addr,
+        ctx.get_proxy_latency(),
+    );
+
+    let (client, connection) = connection.into_managed_conn(
+        SocketConfig {
+            host_addr: Some(compute.socket_addr.ip()),
+            host: postgres_client::config::Host::Tcp(compute.hostname.to_string()),
+            port: compute.socket_addr.port(),
+            connect_timeout: None,
+        },
+        compute.ssl_mode,
+    );
+
+    Ok(poll_client(
+        pool.clone(),
+        ctx,
+        conn_info.clone(),
+        client,
+        connection,
+        conn_id,
+        compute.aux,
+    ))
 }
 
-async fn connect_http2(
-    host_addr: Option<IpAddr>,
-    host: &str,
-    port: u16,
-    timeout: Duration,
-    tls: Option<&Arc<rustls::ClientConfig>>,
-) -> Result<(http_conn_pool::Send, http_conn_pool::Connect), LocalProxyConnError> {
-    let addrs = match host_addr {
-        Some(addr) => vec![SocketAddr::new(addr, port)],
-        None => lookup_host((host, port))
-            .await
-            .map_err(LocalProxyConnError::Io)?
-            .collect(),
-    };
-    let mut last_err = None;
-
-    let mut addrs = addrs.into_iter();
-    let stream = loop {
-        let Some(addr) = addrs.next() else {
-            return Err(last_err.unwrap_or_else(|| {
-                LocalProxyConnError::Io(io::Error::new(
-                    io::ErrorKind::InvalidInput,
-                    "could not resolve any addresses",
-                ))
-            }));
-        };
-
-        match tokio::time::timeout(timeout, TcpStream::connect(addr)).await {
-            Ok(Ok(stream)) => {
-                stream.set_nodelay(true).map_err(LocalProxyConnError::Io)?;
-                break stream;
-            }
-            Ok(Err(e)) => {
-                last_err = Some(LocalProxyConnError::Io(e));
-            }
-            Err(e) => {
-                last_err = Some(LocalProxyConnError::Io(io::Error::new(
-                    io::ErrorKind::TimedOut,
-                    e,
-                )));
-            }
-        }
-    };
-
-    let stream = if let Some(tls) = tls {
-        let host = DnsName::try_from(host)
-            .map_err(io::Error::other)
-            .map_err(LocalProxyConnError::Io)?
-            .to_owned();
-        let stream = TlsConnector::from(tls.clone())
-            .connect(ServerName::DnsName(host), stream)
-            .await
-            .map_err(LocalProxyConnError::Io)?;
-        Box::pin(stream) as AsyncRW
-    } else {
-        Box::pin(stream) as AsyncRW
+async fn h2handshake(
+    ctx: &RequestContext,
+    pool: &Arc<GlobalConnPool<Send, HttpConnPool<Send>>>,
+    conn_info: &ConnInfo,
+    compute: ComputeConnection,
+    conn_id: uuid::Uuid,
+) -> Result<http_conn_pool::Client<Send>, HttpConnError> {
+    let stream = match compute.stream {
+        MaybeTlsStream::Raw(tcp) => Box::pin(tcp) as AsyncRW,
+        MaybeTlsStream::Tls(tls) => Box::into_pin(tls.0) as AsyncRW,
     };
 
+    let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
     let (client, connection) = hyper::client::conn::http2::Builder::new(TokioExecutor::new())
         .timer(TokioTimer::new())
         .keep_alive_interval(Duration::from_secs(20))
         .keep_alive_while_idle(true)
         .keep_alive_timeout(Duration::from_secs(5))
         .handshake(TokioIo::new(stream))
-        .await?;
+        .await
+        .map_err(LocalProxyConnError::H2)?;
+    drop(pause);
 
-    Ok((client, connection))
+    // TODO: lots of useful info but maybe we can move it elsewhere (eg traces?)
+    info!(
+        compute_id = %compute.aux.compute_id,
+        cold_start_info = ctx.cold_start_info().as_str(),
+        query_id = ctx.get_testodrome_id().as_deref(),
+        sslmode = ?compute.ssl_mode,
+        %conn_id,
+        "connected to compute node at {} ({}) latency={}",
+        compute.hostname,
+        compute.socket_addr,
+        ctx.get_proxy_latency(),
+    );
+
+    Ok(poll_http2_client(
+        pool.clone(),
+        ctx,
+        conn_info,
+        client,
+        connection,
+        conn_id,
+        compute.aux,
+    ))
 }

proxy/src/serverless/conn_pool.rs

@@ -69,7 +69,7 @@ pub(crate) fn poll_client<C: ClientInnerExt>(
     let mut session_id = ctx.session_id();
     let (tx, mut rx) = tokio::sync::watch::channel(session_id);
 
-    let span = info_span!(parent: None, "connection", %conn_id);
+    let span = info_span!(parent: None, "connection", %conn_id, pid=client.get_process_id(), compute_id=%aux.compute_id);
     let cold_start_info = ctx.cold_start_info();
     span.in_scope(|| {
         info!(cold_start_info = cold_start_info.as_str(), %conn_info, %session_id, "new connection");

proxy/src/serverless/conn_pool_lib.rs

@@ -518,15 +518,14 @@ impl<C: ClientInnerExt> GlobalConnPool<C, EndpointConnPool<C>> {
                 info!("pool: cached connection '{conn_info}' is closed, opening a new one");
                 return Ok(None);
             }
-            tracing::Span::current()
-                .record("conn_id", tracing::field::display(client.get_conn_id()));
-            tracing::Span::current().record(
-                "pid",
-                tracing::field::display(client.inner.get_process_id()),
-            );
-            debug!(
+            info!(
+                conn_id = %client.get_conn_id(),
+                pid = client.inner.get_process_id(),
+                compute_id = &*client.aux.compute_id,
                 cold_start_info = ColdStartInfo::HttpPoolHit.as_str(),
-                "pool: reusing connection '{conn_info}'"
+                query_id = ctx.get_testodrome_id().as_deref(),
+                "reusing connection: latency={}",
+                ctx.get_proxy_latency(),
             );
 
             match client.get_data() {

proxy/src/serverless/http_conn_pool.rs

@@ -6,7 +6,7 @@ use hyper::client::conn::http2;
 use hyper_util::rt::{TokioExecutor, TokioIo};
 use parking_lot::RwLock;
 use smol_str::ToSmolStr;
-use tracing::{Instrument, debug, error, info, info_span};
+use tracing::{Instrument, error, info, info_span};
 
 use super::AsyncRW;
 use super::backend::HttpConnError;
@@ -115,7 +115,6 @@ impl<C: ClientInnerExt + Clone> Drop for HttpConnPool<C> {
 }
 
 impl<C: ClientInnerExt + Clone> GlobalConnPool<C, HttpConnPool<C>> {
-    #[expect(unused_results)]
     pub(crate) fn get(
         self: &Arc<Self>,
         ctx: &RequestContext,
@@ -132,10 +131,13 @@ impl<C: ClientInnerExt + Clone> GlobalConnPool<C, HttpConnPool<C>> {
             return result;
         };
 
-        tracing::Span::current().record("conn_id", tracing::field::display(client.conn.conn_id));
-        debug!(
+        info!(
+            conn_id = %client.conn.conn_id,
+            compute_id = &*client.conn.aux.compute_id,
             cold_start_info = ColdStartInfo::HttpPoolHit.as_str(),
-            "pool: reusing connection '{conn_info}'"
+            query_id = ctx.get_testodrome_id().as_deref(),
+            "reusing connection: latency={}",
+            ctx.get_proxy_latency(),
         );
         ctx.set_cold_start_info(ColdStartInfo::HttpPoolHit);
         ctx.success();
@@ -197,7 +199,7 @@ pub(crate) fn poll_http2_client(
     let conn_gauge = Metrics::get().proxy.db_connections.guard(ctx.protocol());
     let session_id = ctx.session_id();
 
-    let span = info_span!(parent: None, "connection", %conn_id);
+    let span = info_span!(parent: None, "connection", %conn_id, compute_id=%aux.compute_id);
     let cold_start_info = ctx.cold_start_info();
     span.in_scope(|| {
         info!(cold_start_info = cold_start_info.as_str(), %conn_info, %session_id, "new connection");
@@ -229,6 +231,8 @@ pub(crate) fn poll_http2_client(
 
     tokio::spawn(
         async move {
+            info!("new local proxy connection");
+
             let _conn_gauge = conn_gauge;
             let res = connection.await;
             match res {

proxy/src/serverless/local_conn_pool.rs

@@ -30,7 +30,7 @@ use serde_json::value::RawValue;
 use tokio::net::TcpStream;
 use tokio::time::Instant;
 use tokio_util::sync::CancellationToken;
-use tracing::{Instrument, debug, error, info, info_span, warn};
+use tracing::{Instrument, error, info, info_span, warn};
 
 use super::backend::HttpConnError;
 use super::conn_pool_lib::{
@@ -107,15 +107,13 @@ impl<C: ClientInnerExt> LocalConnPool<C> {
                 return Ok(None);
             }
 
-            tracing::Span::current()
-                .record("conn_id", tracing::field::display(client.get_conn_id()));
-            tracing::Span::current().record(
-                "pid",
-                tracing::field::display(client.inner.get_process_id()),
-            );
-            debug!(
+            info!(
+                pid = client.inner.get_process_id(),
+                conn_id = %client.get_conn_id(),
                 cold_start_info = ColdStartInfo::HttpPoolHit.as_str(),
-                "local_pool: reusing connection '{conn_info}'"
+                query_id = ctx.get_testodrome_id().as_deref(),
+                "reusing connection: latency={}",
+                ctx.get_proxy_latency(),
             );
 
             match client.get_data() {

proxy/src/tls/postgres_rustls.rs

@@ -60,7 +60,7 @@ mod private {
         }
     }
 
-    pub struct RustlsStream<S>(Box<TlsStream<S>>);
+    pub struct RustlsStream<S>(pub Box<TlsStream<S>>);
 
     impl<S> postgres_client::tls::TlsStream for RustlsStream<S>
     where

safekeeper/src/copy_timeline.rs

@@ -9,7 +9,7 @@ use tokio::fs::OpenOptions;
 use tokio::io::{AsyncSeekExt, AsyncWriteExt};
 use tracing::{info, warn};
 use utils::id::TenantTimelineId;
-use utils::lsn::{Lsn, WalSegmentSize};
+use utils::lsn::Lsn;
 
 use crate::GlobalTimelines;
 use crate::control_file::FileStorage;
@@ -100,7 +100,7 @@ pub async fn handle_request(
         }
     }
 
-    let wal_seg_size = state.server.wal_seg_size;
+    let wal_seg_size = state.server.wal_seg_size as usize;
     if wal_seg_size == 0 {
         bail!("wal_seg_size is not set");
     }
@@ -171,7 +171,7 @@ pub async fn handle_request(
 
 async fn copy_disk_segments(
     tli: &WalResidentTimeline,
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
     start_lsn: Lsn,
     end_lsn: Lsn,
     tli_dir_path: &Utf8PathBuf,

safekeeper/src/http/routes.rs

@@ -103,7 +103,7 @@ async fn timeline_create_handler(mut request: Request<Body>) -> Result<Response<
     let server_info = ServerInfo {
         pg_version: request_data.pg_version,
         system_id: request_data.system_id.unwrap_or(0),
-        wal_seg_size: request_data.wal_seg_size.unwrap_or(WAL_SEGMENT_SIZE),
+        wal_seg_size: request_data.wal_seg_size.unwrap_or(WAL_SEGMENT_SIZE as u32),
     };
     let global_timelines = get_global_timelines(&request);
     global_timelines

safekeeper/src/metrics.rs

@@ -831,7 +831,7 @@ impl Collector for TimelineCollector {
             if tli.last_removed_segno != 0 {
                 let segno_count = tli
                     .flush_lsn
-                    .segment_number(tli.persisted_state.server.wal_seg_size)
+                    .segment_number(tli.persisted_state.server.wal_seg_size as usize)
                     - tli.last_removed_segno;
                 let disk_usage_bytes = segno_count * tli.persisted_state.server.wal_seg_size as u64;
                 self.disk_usage

safekeeper/src/pull_timeline.rs

@@ -27,7 +27,7 @@ use tracing::{error, info, instrument};
 use utils::crashsafe::fsync_async_opt;
 use utils::id::{NodeId, TenantTimelineId};
 use utils::logging::SecretString;
-use utils::lsn::{Lsn, WalSegmentSize};
+use utils::lsn::Lsn;
 use utils::pausable_failpoint;
 
 use crate::control_file::CONTROL_FILE_NAME;
@@ -100,7 +100,7 @@ pub struct SnapshotContext {
     pub term: Term,
     pub last_log_term: Term,
     pub flush_lsn: Lsn,
-    pub wal_seg_size: WalSegmentSize,
+    pub wal_seg_size: usize,
     // used to remove WAL hold off in Drop.
     pub tli: WalResidentTimeline,
 }

safekeeper/src/safekeeper.rs

@@ -1439,7 +1439,7 @@ mod tests {
 
     fn test_sk_state() -> TimelinePersistentState {
         let mut state = TimelinePersistentState::empty();
-        state.server.wal_seg_size = WAL_SEGMENT_SIZE;
+        state.server.wal_seg_size = WAL_SEGMENT_SIZE as u32;
         state.tenant_id = TenantId::from([1u8; 16]);
         state.timeline_id = TimelineId::from([1u8; 16]);
         state

safekeeper/src/state.rs

@@ -152,7 +152,7 @@ impl TimelinePersistentState {
             ServerInfo {
                 pg_version: PgVersionId::from(PgMajorVersion::PG17),
                 system_id: 0, /* Postgres system identifier */
-                wal_seg_size: WAL_SEGMENT_SIZE,
+                wal_seg_size: WAL_SEGMENT_SIZE as u32,
             },
             Lsn::INVALID,
             Lsn::INVALID,

safekeeper/src/timeline.rs

@@ -23,7 +23,7 @@ use tokio::time::Instant;
 use tokio_util::sync::CancellationToken;
 use tracing::*;
 use utils::id::{NodeId, TenantId, TenantTimelineId};
-use utils::lsn::{Lsn, WalSegmentSize};
+use utils::lsn::Lsn;
 use utils::sync::gate::Gate;
 
 use crate::metrics::{FullTimelineInfo, MISC_OPERATION_SECONDS, WalStorageMetrics};
@@ -338,8 +338,8 @@ impl SharedState {
         Ok(Self::new(sk))
     }
 
-    pub(crate) fn get_wal_seg_size(&self) -> WalSegmentSize {
-        self.sk.state().server.wal_seg_size
+    pub(crate) fn get_wal_seg_size(&self) -> usize {
+        self.sk.state().server.wal_seg_size as usize
     }
 
     fn get_safekeeper_info(
@@ -747,7 +747,7 @@ impl Timeline {
     }
 
     /// Returns wal_seg_size.
-    pub async fn get_wal_seg_size(&self) -> WalSegmentSize {
+    pub async fn get_wal_seg_size(&self) -> usize {
         self.read_shared_state().await.get_wal_seg_size()
     }
 

safekeeper/src/timeline_eviction.rs

@@ -11,7 +11,6 @@ use tokio::fs::File;
 use tokio::io::{AsyncRead, AsyncWriteExt};
 use tracing::{debug, info, instrument, warn};
 use utils::crashsafe::durable_rename;
-use utils::lsn::WalSegmentSize;
 
 use crate::metrics::{
     EVICTION_EVENTS_COMPLETED, EVICTION_EVENTS_STARTED, EvictionEvent, NUM_EVICTED_TIMELINES,
@@ -277,12 +276,12 @@ async fn compare_local_segment_with_remote(
 async fn do_validation(
     mgr: &Manager,
     file: &mut File,
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
     partial: &PartialRemoteSegment,
     storage: &GenericRemoteStorage,
 ) -> anyhow::Result<()> {
-    let local_size = file.metadata().await?.len();
-    if WalSegmentSize::try_from(local_size) != Ok(wal_seg_size) {
+    let local_size = file.metadata().await?.len() as usize;
+    if local_size != wal_seg_size {
         anyhow::bail!(
             "local segment size is invalid: found {}, expected {}",
             local_size,
@@ -297,12 +296,12 @@ async fn do_validation(
     // remote segment should have bytes excatly up to `flush_lsn`
     let expected_remote_size = partial.flush_lsn.segment_offset(mgr.wal_seg_size);
     // let's compare the first `expected_remote_size` bytes
-    compare_n_bytes(&mut remote_reader, file, expected_remote_size as usize).await?;
+    compare_n_bytes(&mut remote_reader, file, expected_remote_size).await?;
     // and check that the remote segment ends here
     check_end(&mut remote_reader).await?;
 
     // if local segment is longer, the rest should be zeroes
-    read_n_zeroes(file, (mgr.wal_seg_size - expected_remote_size) as usize).await?;
+    read_n_zeroes(file, mgr.wal_seg_size - expected_remote_size).await?;
     // and check that the local segment ends here
     check_end(file).await?;
 

safekeeper/src/timeline_manager.rs

@@ -20,7 +20,7 @@ use tokio::task::{JoinError, JoinHandle};
 use tokio::time::Instant;
 use tokio_util::sync::CancellationToken;
 use tracing::{Instrument, debug, info, info_span, instrument, warn};
-use utils::lsn::{Lsn, WalSegmentSize};
+use utils::lsn::Lsn;
 
 use crate::SafeKeeperConf;
 use crate::control_file::{FileStorage, Storage};
@@ -198,7 +198,7 @@ pub(crate) struct Manager {
     // configuration & dependencies
     pub(crate) tli: ManagerTimeline,
     pub(crate) conf: SafeKeeperConf,
-    pub(crate) wal_seg_size: WalSegmentSize,
+    pub(crate) wal_seg_size: usize,
     pub(crate) walsenders: Arc<WalSenders>,
     pub(crate) wal_backup: Arc<WalBackup>,
 

safekeeper/src/wal_backup.rs

@@ -23,7 +23,7 @@ use tokio::task::JoinHandle;
 use tokio_util::sync::CancellationToken;
 use tracing::*;
 use utils::id::{NodeId, TenantTimelineId};
-use utils::lsn::{Lsn, WalSegmentSize};
+use utils::lsn::Lsn;
 use utils::{backoff, pausable_failpoint};
 
 use crate::metrics::{BACKED_UP_SEGMENTS, BACKUP_ERRORS, WAL_BACKUP_TASKS};
@@ -52,7 +52,7 @@ impl WalBackupTaskHandle {
 
 /// Do we have anything to upload to S3, i.e. should safekeepers run backup activity?
 pub(crate) fn is_wal_backup_required(
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
     num_computes: usize,
     state: &StateSnapshot,
 ) -> bool {
@@ -210,7 +210,7 @@ impl WalBackup {
 struct WalBackupTask {
     timeline: WalResidentTimeline,
     timeline_dir: Utf8PathBuf,
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
     parallel_jobs: usize,
     commit_lsn_watch_rx: watch::Receiver<Lsn>,
     storage: Arc<GenericRemoteStorage>,
@@ -338,7 +338,7 @@ async fn backup_lsn_range(
     storage: Arc<GenericRemoteStorage>,
     backup_lsn: &mut Lsn,
     end_lsn: Lsn,
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
     timeline_dir: &Utf8Path,
     parallel_jobs: usize,
 ) -> Result<()> {
@@ -461,12 +461,12 @@ impl Segment {
         remote_timeline_path.join(self.object_name())
     }
 
-    pub fn size(self) -> WalSegmentSize {
-        (u64::from(self.end_lsn) - u64::from(self.start_lsn)) as WalSegmentSize
+    pub fn size(self) -> usize {
+        (u64::from(self.end_lsn) - u64::from(self.start_lsn)) as usize
     }
 }
 
-fn get_segments(start: Lsn, end: Lsn, seg_size: WalSegmentSize) -> Vec<Segment> {
+fn get_segments(start: Lsn, end: Lsn, seg_size: usize) -> Vec<Segment> {
     let first_seg = start.segment_number(seg_size);
     let last_seg = end.segment_number(seg_size);
 
@@ -484,7 +484,7 @@ async fn backup_object(
     storage: &GenericRemoteStorage,
     source_file: &Utf8Path,
     target_file: &RemotePath,
-    size: WalSegmentSize,
+    size: usize,
 ) -> Result<()> {
     let file = File::open(&source_file)
         .await
@@ -495,7 +495,7 @@ async fn backup_object(
     let cancel = CancellationToken::new();
 
     storage
-        .upload_storage_object(file, size as usize, target_file, &cancel)
+        .upload_storage_object(file, size, target_file, &cancel)
         .await
 }
 
@@ -503,7 +503,7 @@ pub(crate) async fn backup_partial_segment(
     storage: &GenericRemoteStorage,
     source_file: &Utf8Path,
     target_file: &RemotePath,
-    size: WalSegmentSize,
+    size: usize,
 ) -> Result<()> {
     let file = File::open(&source_file)
         .await
@@ -519,7 +519,7 @@ pub(crate) async fn backup_partial_segment(
     storage
         .upload(
             file,
-            size as usize,
+            size,
             target_file,
             Some(StorageMetadata::from([("sk_type", "partial_segment")])),
             &cancel,
@@ -647,7 +647,7 @@ pub async fn delete_objects(storage: &GenericRemoteStorage, paths: &[RemotePath]
 /// Copy segments from one timeline to another. Used in copy_timeline.
 pub async fn copy_s3_segments(
     storage: &GenericRemoteStorage,
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
     src_ttid: &TenantTimelineId,
     dst_ttid: &TenantTimelineId,
     from_segment: XLogSegNo,

safekeeper/src/wal_backup_partial.rs

@@ -29,7 +29,7 @@ use serde::{Deserialize, Serialize};
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, info, instrument, warn};
 use utils::id::NodeId;
-use utils::lsn::{Lsn, WalSegmentSize};
+use utils::lsn::Lsn;
 
 use crate::SafeKeeperConf;
 use crate::metrics::{
@@ -151,7 +151,7 @@ impl State {
 }
 
 pub struct PartialBackup {
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
     tli: WalResidentTimeline,
     conf: SafeKeeperConf,
     local_prefix: Utf8PathBuf,

safekeeper/src/wal_storage.rs

@@ -28,7 +28,7 @@ use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
 use tracing::*;
 use utils::crashsafe::durable_rename;
 use utils::id::TenantTimelineId;
-use utils::lsn::{Lsn, WalSegmentSize};
+use utils::lsn::Lsn;
 
 use crate::metrics::{
     REMOVED_WAL_SEGMENTS, WAL_STORAGE_OPERATION_SECONDS, WalStorageMetrics, time_io_closure,
@@ -92,7 +92,7 @@ pub struct PhysicalStorage {
     no_sync: bool,
 
     /// Size of WAL segment in bytes.
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
     pg_version: PgVersionId,
     system_id: u64,
 
@@ -170,7 +170,7 @@ impl PhysicalStorage {
         state: &TimelinePersistentState,
         no_sync: bool,
     ) -> Result<PhysicalStorage> {
-        let wal_seg_size = state.server.wal_seg_size;
+        let wal_seg_size = state.server.wal_seg_size as usize;
 
         // Find out where stored WAL ends, starting at commit_lsn which is a
         // known recent record boundary (unless we don't have WAL at all).
@@ -315,12 +315,7 @@ impl PhysicalStorage {
 
     /// Write WAL bytes, which are known to be located in a single WAL segment. Returns true if the
     /// segment was completed, closed, and flushed to disk.
-    async fn write_in_segment(
-        &mut self,
-        segno: u64,
-        xlogoff: WalSegmentSize,
-        buf: &[u8],
-    ) -> Result<bool> {
+    async fn write_in_segment(&mut self, segno: u64, xlogoff: usize, buf: &[u8]) -> Result<bool> {
         let mut file = if let Some(file) = self.file.take() {
             file
         } else {
@@ -336,7 +331,7 @@ impl PhysicalStorage {
         // syscall, but needed in case of async). It does *not* fsyncs the file.
         file.flush().await?;
 
-        if xlogoff as usize + buf.len() == self.wal_seg_size as usize {
+        if xlogoff + buf.len() == self.wal_seg_size {
             // If we reached the end of a WAL segment, flush and close it.
             self.fdatasync_file(&file).await?;
 
@@ -377,8 +372,8 @@ impl PhysicalStorage {
             let segno = self.write_lsn.segment_number(self.wal_seg_size);
 
             // If crossing a WAL boundary, only write up until we reach wal segment size.
-            let bytes_write = if xlogoff as usize + buf.len() > self.wal_seg_size as usize {
-                (self.wal_seg_size - xlogoff) as usize
+            let bytes_write = if xlogoff + buf.len() > self.wal_seg_size {
+                self.wal_seg_size - xlogoff
             } else {
                 buf.len()
             };
@@ -609,7 +604,7 @@ impl Storage for PhysicalStorage {
 /// Remove all WAL segments in timeline_dir that match the given predicate.
 async fn remove_segments_from_disk(
     timeline_dir: &Utf8Path,
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
     remove_predicate: impl Fn(XLogSegNo) -> bool,
 ) -> Result<()> {
     let _timer = WAL_STORAGE_OPERATION_SECONDS
@@ -650,7 +645,7 @@ async fn remove_segments_from_disk(
 pub struct WalReader {
     remote_path: RemotePath,
     timeline_dir: Utf8PathBuf,
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
     pos: Lsn,
     wal_segment: Option<Pin<Box<dyn AsyncRead + Send + Sync>>>,
 
@@ -688,7 +683,7 @@ impl WalReader {
         if start_pos
             < state
                 .timeline_start_lsn
-                .segment_lsn(state.server.wal_seg_size)
+                .segment_lsn(state.server.wal_seg_size as usize)
         {
             bail!(
                 "Requested streaming from {}, which is before the start of the timeline {}, and also doesn't start at the first segment of that timeline",
@@ -700,7 +695,7 @@ impl WalReader {
         Ok(Self {
             remote_path: remote_timeline_path(ttid)?,
             timeline_dir,
-            wal_seg_size: state.server.wal_seg_size,
+            wal_seg_size: state.server.wal_seg_size as usize,
             pos: start_pos,
             wal_segment: None,
             wal_backup,
@@ -748,14 +743,12 @@ impl WalReader {
             // How many bytes may we consume in total?
             let tl_start_seg_offset = self.timeline_start_lsn.segment_offset(self.wal_seg_size);
 
-            debug_assert!(seg_bytes.len() > pos_seg_offset as usize);
-            debug_assert!(seg_bytes.len() > tl_start_seg_offset as usize);
+            debug_assert!(seg_bytes.len() > pos_seg_offset);
+            debug_assert!(seg_bytes.len() > tl_start_seg_offset);
 
             // Copy as many bytes as possible into the buffer
-            let len = ((tl_start_seg_offset - pos_seg_offset) as usize).min(buf.len());
-            buf[0..len].copy_from_slice(
-                &seg_bytes[pos_seg_offset as usize..pos_seg_offset as usize + len],
-            );
+            let len = (tl_start_seg_offset - pos_seg_offset).min(buf.len());
+            buf[0..len].copy_from_slice(&seg_bytes[pos_seg_offset..pos_seg_offset + len]);
 
             self.pos += len as u64;
 
@@ -777,7 +770,7 @@ impl WalReader {
         // How much to read and send in message? We cannot cross the WAL file
         // boundary, and we don't want send more than provided buffer.
         let xlogoff = self.pos.segment_offset(self.wal_seg_size);
-        let send_size = min(buf.len(), (self.wal_seg_size - xlogoff) as usize);
+        let send_size = min(buf.len(), self.wal_seg_size - xlogoff);
 
         // Read some data from the file.
         let buf = &mut buf[0..send_size];
@@ -838,7 +831,7 @@ impl WalReader {
 pub(crate) async fn open_wal_file(
     timeline_dir: &Utf8Path,
     segno: XLogSegNo,
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
 ) -> Result<(tokio::fs::File, bool)> {
     let (wal_file_path, wal_file_partial_path) = wal_file_paths(timeline_dir, segno, wal_seg_size);
 
@@ -865,7 +858,7 @@ pub(crate) async fn open_wal_file(
 pub fn wal_file_paths(
     timeline_dir: &Utf8Path,
     segno: XLogSegNo,
-    wal_seg_size: WalSegmentSize,
+    wal_seg_size: usize,
 ) -> (Utf8PathBuf, Utf8PathBuf) {
     let wal_file_name = XLogFileName(PG_TLI, segno, wal_seg_size);
     let wal_file_path = timeline_dir.join(wal_file_name.clone());

storage_scrubber/src/scan_safekeeper_metadata.rs

@@ -13,7 +13,7 @@ use serde::Serialize;
 use tokio_postgres::types::PgLsn;
 use tracing::{debug, error, info};
 use utils::id::{TenantId, TenantTimelineId, TimelineId};
-use utils::lsn::{Lsn, WalSegmentSize};
+use utils::lsn::Lsn;
 
 use crate::cloud_admin_api::CloudAdminApiClient;
 use crate::metadata_stream::stream_listing;
@@ -22,7 +22,7 @@ use crate::{
 };
 
 /// Generally we should ask safekeepers, but so far we use everywhere default 16MB.
-const WAL_SEGSIZE: WalSegmentSize = 16 * 1024 * 1024;
+const WAL_SEGSIZE: usize = 16 * 1024 * 1024;
 
 #[derive(Serialize)]
 pub struct MetadataSummary {

test_runner/regress/test_ondemand_download.py

@@ -671,12 +671,6 @@ def test_layer_download_cancelled_by_config_location(neon_env_builder: NeonEnvBu
     """
     neon_env_builder.enable_pageserver_remote_storage(s3_storage())
 
-    # On the new mode, the test runs into a cancellation issue, i.e. the walproposer can't shut down
-    # as it is hang-waiting on the timeline_checkpoint call in WalIngest::new.
-    neon_env_builder.storage_controller_config = {
-        "timelines_onto_safekeepers": False,
-    }
-
     # turn off background tasks so that they don't interfere with the downloads
     env = neon_env_builder.init_start(
         initial_tenant_conf={