Add forgotten files.

Add fake_timeline endoint creating timeline + some WAL.

curl -X POST http://127.0.0.1:7676/v1/fake_timeline
Set in pg_receivewal.c:
  stream.startpos = 0x1493AC8;
pg_install/v15/bin/pg_receivewal -v -d "host=localhost port=5454 options='-c tenant_id=deadbeefdeadbeefdeadbeefdeadbeef timeline_id=deadbeefdeadbeefdeadbeefdeadbeef'" -D ~/tmp/tmp/tmp

control_plane/src/compute.rs

@@ -14,7 +14,7 @@ use anyhow::{Context, Result};
 use utils::{
     id::{TenantId, TimelineId},
     lsn::Lsn,
-    postgres_backend::AuthType,
+    postgres_backend_async::AuthType,
 };
 
 use crate::local_env::{LocalEnv, DEFAULT_PG_VERSION};

control_plane/src/local_env.rs

@@ -19,7 +19,7 @@ use std::process::{Command, Stdio};
 use utils::{
     auth::{encode_from_key_file, Claims, Scope},
     id::{NodeId, TenantId, TenantTimelineId, TimelineId},
-    postgres_backend::AuthType,
+    postgres_backend_async::AuthType,
 };
 
 use crate::safekeeper::SafekeeperNode;

libs/pq_proto/Cargo.toml

@@ -7,11 +7,13 @@ license = "Apache-2.0"
 [dependencies]
 anyhow = "1.0"
 bytes = "1.0.1"
+byteorder = "1.4.3"
 pin-project-lite = "0.2.7"
 postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
 rand = "0.8.3"
 serde = { version = "1.0", features = ["derive"] }
 tokio = { version = "1.17", features = ["macros"] }
+tokio-util = { version = "0.7.3" }
 tracing = "0.1"
 thiserror = "1.0"
 

libs/pq_proto/src/codec.rs

@@ -0,0 +1,62 @@
+//! Provides `PostgresCodec` defining how to serilize/deserialize Postgres
+//! messages to/from the wire, to be used with `tokio_util::codec::Framed`.
+use std::io;
+
+use bytes::BytesMut;
+use tokio_util::codec::{Decoder, Encoder};
+
+use crate::{BeMessage, FeMessage, FeStartupPacket, ProtocolError};
+
+// Defines how to serilize/deserialize Postgres messages to/from the wire, to be
+// used with `tokio_util::codec::Framed`.
+pub struct PostgresCodec {
+    // Have we already decoded startup message? All further should start with
+    // message type byte then.
+    startup_read: bool,
+}
+
+impl PostgresCodec {
+    pub fn new() -> Self {
+        PostgresCodec {
+            startup_read: false,
+        }
+    }
+}
+
+/// Error on postgres connection: either IO (physical transport error) or
+/// protocol violation.
+#[derive(thiserror::Error, Debug)]
+pub enum ConnectionError {
+    #[error(transparent)]
+    Io(#[from] io::Error),
+    #[error(transparent)]
+    Protocol(#[from] ProtocolError),
+}
+
+impl Encoder<&BeMessage<'_>> for PostgresCodec {
+    type Error = ConnectionError;
+
+    fn encode(&mut self, item: &BeMessage, dst: &mut BytesMut) -> Result<(), ConnectionError> {
+        BeMessage::write(dst, &item)?;
+        Ok(())
+    }
+}
+
+impl Decoder for PostgresCodec {
+    type Item = FeMessage;
+    type Error = ConnectionError;
+
+    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<FeMessage>, ConnectionError> {
+        let msg = if !self.startup_read {
+            let msg = FeStartupPacket::parse(src);
+            if let Ok(Some(FeMessage::StartupPacket(FeStartupPacket::StartupMessage { .. }))) = msg
+            {
+                self.startup_read = true;
+            }
+            msg?
+        } else {
+            FeMessage::parse(src)?
+        };
+        Ok(msg)
+    }
+}

libs/pq_proto/src/lib.rs

@@ -3,9 +3,11 @@
 //! on message formats.
 
 // Tools for calling certain async methods in sync contexts.
+pub mod codec;
 pub mod sync;
 
-use anyhow::{ensure, Context, Result};
+use anyhow::{anyhow, bail, ensure, Context, Result};
+use byteorder::{BigEndian, ByteOrder, ReadBytesExt};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use postgres_protocol::PG_EPOCH;
 use serde::{Deserialize, Serialize};
@@ -19,7 +21,7 @@ use std::{
     time::{Duration, SystemTime},
 };
 use sync::{AsyncishRead, SyncFuture};
-use tokio::io::AsyncReadExt;
+// use tokio::io::AsyncReadExt;
 use tracing::{trace, warn};
 
 pub type Oid = u32;
@@ -194,36 +196,108 @@ macro_rules! retry_read {
     };
 }
 
-/// An error occured during connection being open.
+/// An error occured while parsing or serializing raw stream into Postgres
+/// messages.
 #[derive(thiserror::Error, Debug)]
-pub enum ConnectionError {
+pub enum ProtocolError {
     /// IO error during writing to or reading from the connection socket.
+    /// removeme
     #[error("Socket IO error: {0}")]
     Socket(std::io::Error),
-    /// Invalid packet was received from client
+    /// Invalid packet was received from the client (e.g. unexpected message
+    /// type or broken len).
     #[error("Protocol error: {0}")]
     Protocol(String),
-    /// Failed to parse a protocol mesage
+    /// Failed to parse or, (unlikely), serialize a protocol message.
     #[error("Message parse error: {0}")]
     MessageParse(anyhow::Error),
 }
 
-impl From<anyhow::Error> for ConnectionError {
+// Allows to return anyhow error from msg parsing routines, meaning less typing.
+impl From<anyhow::Error> for ProtocolError {
     fn from(e: anyhow::Error) -> Self {
         Self::MessageParse(e)
     }
 }
 
-impl ConnectionError {
+impl ProtocolError {
     pub fn into_io_error(self) -> io::Error {
         match self {
-            ConnectionError::Socket(io) => io,
+            ProtocolError::Socket(io) => io,
             other => io::Error::new(io::ErrorKind::Other, other.to_string()),
         }
     }
 }
 
 impl FeMessage {
+    /// Read and parse one message from the `buf` input buffer. If there is at
+    /// least one valid message, returns it, advancing `buf`; redundant copies
+    /// are avoided, as thanks to `bytes` crate ptrs in parsed message point
+    /// directly into the `buf` (processed data is garbage collected after
+    /// parsed message is dropped).
+    ///
+    /// Returns None if `buf` doesn't contain enough data for a single message.
+    /// For efficiency, tries to reserve large enough space in `buf` for the
+    /// next message in this case.
+    ///
+    /// Returns Error if message is malformed, the only possible ErrorKind is
+    /// InvalidInput.
+    //
+    // Inspired by rust-postgres Message::parse.
+    pub fn parse(buf: &mut BytesMut) -> Result<Option<FeMessage>, ProtocolError> {
+        // Every message contains message type byte and 4 bytes len; can't do
+        // much without them.
+        if buf.len() < 5 {
+            let to_read = 5 - buf.len();
+            buf.reserve(to_read);
+            return Ok(None);
+        }
+
+        // We shouldn't advance `buf` as probably full message is not there yet,
+        // so can't directly use Bytes::get_u32 etc.
+        let tag = buf[0];
+        let len = (&buf[1..5]).read_u32::<BigEndian>().unwrap();
+        if len < 4 {
+            return Err(ProtocolError::Protocol(format!(
+                "invalid message length {}",
+                len
+            )));
+        }
+
+        // lenth field includes itself, but not message type.
+        let total_len = len as usize + 1;
+        if buf.len() < total_len {
+            // Don't have full message yet.
+            let to_read = total_len - buf.len();
+            buf.reserve(to_read);
+            return Ok(None);
+        }
+
+        // got the message, advance buffer
+        let mut msg = buf.split_to(total_len).freeze();
+        msg.advance(5); // consume message type and len
+
+        match tag {
+            b'Q' => Ok(Some(FeMessage::Query(msg))),
+            b'P' => Ok(Some(FeParseMessage::parse(msg)?)),
+            b'D' => Ok(Some(FeDescribeMessage::parse(msg)?)),
+            b'E' => Ok(Some(FeExecuteMessage::parse(msg)?)),
+            b'B' => Ok(Some(FeBindMessage::parse(msg)?)),
+            b'C' => Ok(Some(FeCloseMessage::parse(msg)?)),
+            b'S' => Ok(Some(FeMessage::Sync)),
+            b'X' => Ok(Some(FeMessage::Terminate)),
+            b'd' => Ok(Some(FeMessage::CopyData(msg))),
+            b'c' => Ok(Some(FeMessage::CopyDone)),
+            b'f' => Ok(Some(FeMessage::CopyFail)),
+            b'p' => Ok(Some(FeMessage::PasswordMessage(msg))),
+            tag => {
+                return Err(ProtocolError::Protocol(format!(
+                    "unknown message tag: {tag},'{msg:?}'"
+                )))
+            }
+        }
+    }
+
     /// Read one message from the stream.
     /// This function returns `Ok(None)` in case of EOF.
     /// One way to handle this properly:
@@ -245,68 +319,8 @@ impl FeMessage {
     /// }
     /// ```
     #[inline(never)]
-    pub fn read(
-        stream: &mut (impl io::Read + Unpin),
-    ) -> Result<Option<FeMessage>, ConnectionError> {
-        Self::read_fut(&mut AsyncishRead(stream)).wait()
-    }
-
-    /// Read one message from the stream.
-    /// See documentation for `Self::read`.
-    pub fn read_fut<Reader>(
-        stream: &mut Reader,
-    ) -> SyncFuture<Reader, impl Future<Output = Result<Option<FeMessage>, ConnectionError>> + '_>
-    where
-        Reader: tokio::io::AsyncRead + Unpin,
-    {
-        // We return a Future that's sync (has a `wait` method) if and only if the provided stream is SyncProof.
-        // SyncFuture contract: we are only allowed to await on sync-proof futures, the AsyncRead and
-        // AsyncReadExt methods of the stream.
-        SyncFuture::new(async move {
-            // Each libpq message begins with a message type byte, followed by message length
-            // If the client closes the connection, return None. But if the client closes the
-            // connection in the middle of a message, we will return an error.
-            let tag = match retry_read!(stream.read_u8().await) {
-                Ok(b) => b,
-                Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
-                Err(e) => return Err(ConnectionError::Socket(e)),
-            };
-
-            // The message length includes itself, so it better be at least 4.
-            let len = retry_read!(stream.read_u32().await)
-                .map_err(ConnectionError::Socket)?
-                .checked_sub(4)
-                .ok_or_else(|| ConnectionError::Protocol("invalid message length".to_string()))?;
-
-            let body = {
-                let mut buffer = vec![0u8; len as usize];
-                stream
-                    .read_exact(&mut buffer)
-                    .await
-                    .map_err(ConnectionError::Socket)?;
-                Bytes::from(buffer)
-            };
-
-            match tag {
-                b'Q' => Ok(Some(FeMessage::Query(body))),
-                b'P' => Ok(Some(FeParseMessage::parse(body)?)),
-                b'D' => Ok(Some(FeDescribeMessage::parse(body)?)),
-                b'E' => Ok(Some(FeExecuteMessage::parse(body)?)),
-                b'B' => Ok(Some(FeBindMessage::parse(body)?)),
-                b'C' => Ok(Some(FeCloseMessage::parse(body)?)),
-                b'S' => Ok(Some(FeMessage::Sync)),
-                b'X' => Ok(Some(FeMessage::Terminate)),
-                b'd' => Ok(Some(FeMessage::CopyData(body))),
-                b'c' => Ok(Some(FeMessage::CopyDone)),
-                b'f' => Ok(Some(FeMessage::CopyFail)),
-                b'p' => Ok(Some(FeMessage::PasswordMessage(body))),
-                tag => {
-                    return Err(ConnectionError::Protocol(format!(
-                        "unknown message tag: {tag},'{body:?}'"
-                    )))
-                }
-            }
-        })
+    pub fn read(_stream: &mut (impl io::Read + Unpin)) -> Result<Option<FeMessage>, ProtocolError> {
+        Ok(None) // removeme
     }
 }
 
@@ -314,21 +328,124 @@ impl FeStartupPacket {
     /// Read startup message from the stream.
     // XXX: It's tempting yet undesirable to accept `stream` by value,
     // since such a change will cause user-supplied &mut references to be consumed
-    pub fn read(
-        stream: &mut (impl io::Read + Unpin),
-    ) -> Result<Option<FeMessage>, ConnectionError> {
+    pub fn read(stream: &mut (impl io::Read + Unpin)) -> Result<Option<FeMessage>, ProtocolError> {
         Self::read_fut(&mut AsyncishRead(stream)).wait()
     }
 
+    /// Read and parse startup message from the `buf` input buffer. It is
+    /// different from [`FeMessage::parse`] because startup messages don't have
+    /// message type byte; otherwise, its comments apply.
+    pub fn parse(buf: &mut BytesMut) -> Result<Option<FeMessage>, ProtocolError> {
+        const MAX_STARTUP_PACKET_LENGTH: usize = 10000;
+        const RESERVED_INVALID_MAJOR_VERSION: u32 = 1234;
+        const CANCEL_REQUEST_CODE: u32 = 5678;
+        const NEGOTIATE_SSL_CODE: u32 = 5679;
+        const NEGOTIATE_GSS_CODE: u32 = 5680;
+
+        if buf.len() < 4 {
+            let to_read = 5 - buf.len();
+            buf.reserve(to_read);
+            return Ok(None);
+        }
+
+        // We shouldn't advance `buf` as probably full message is not there yet,
+        // so can't directly use Bytes::get_u32 etc.
+        let len = (&buf[0..4]).read_u32::<BigEndian>().unwrap() as usize;
+        if len < 8 || len > MAX_STARTUP_PACKET_LENGTH {
+            return Err(ProtocolError::Protocol(format!(
+                "invalid startup packet message length {}",
+                len
+            )));
+        }
+
+        if buf.len() < len {
+            // Don't have full message yet.
+            let to_read = len - buf.len();
+            buf.reserve(to_read);
+            return Ok(None);
+        }
+
+        // got the message, advance buffer
+        let mut msg = buf.split_to(len).freeze();
+        msg.advance(4); // consume len
+
+        let request_code = msg.get_u32();
+        let req_hi = request_code >> 16;
+        let req_lo = request_code & ((1 << 16) - 1);
+        // StartupMessage, CancelRequest, SSLRequest etc are differentiated by request code.
+        let message = match (req_hi, req_lo) {
+            (RESERVED_INVALID_MAJOR_VERSION, CANCEL_REQUEST_CODE) => {
+                if msg.remaining() < 8 {
+                    return Err(ProtocolError::MessageParse(anyhow!(
+                        "CancelRequest message is malformed, backend PID / secret key missing"
+                    )));
+                }
+                FeStartupPacket::CancelRequest(CancelKeyData {
+                    backend_pid: msg.get_i32(),
+                    cancel_key: msg.get_i32(),
+                })
+            }
+            (RESERVED_INVALID_MAJOR_VERSION, NEGOTIATE_SSL_CODE) => {
+                // Requested upgrade to SSL (aka TLS)
+                FeStartupPacket::SslRequest
+            }
+            (RESERVED_INVALID_MAJOR_VERSION, NEGOTIATE_GSS_CODE) => {
+                // Requested upgrade to GSSAPI
+                FeStartupPacket::GssEncRequest
+            }
+            (RESERVED_INVALID_MAJOR_VERSION, unrecognized_code) => {
+                return Err(ProtocolError::Protocol(format!(
+                    "Unrecognized request code {unrecognized_code}"
+                )));
+            }
+            // TODO bail if protocol major_version is not 3?
+            (major_version, minor_version) => {
+                // StartupMessage
+
+                // Parse pairs of null-terminated strings (key, value).
+                // See `postgres: ProcessStartupPacket, build_startup_packet`.
+                let mut tokens = str::from_utf8(&msg)
+                    .context("StartupMessage params: invalid utf-8")?
+                    .strip_suffix('\0') // drop packet's own null
+                    .ok_or_else(|| {
+                        ProtocolError::Protocol(
+                            "StartupMessage params: missing null terminator".to_string(),
+                        )
+                    })?
+                    .split_terminator('\0');
+
+                let mut params = HashMap::new();
+                while let Some(name) = tokens.next() {
+                    let value = tokens.next().ok_or_else(|| {
+                        ProtocolError::Protocol(
+                            "StartupMessage params: key without value".to_string(),
+                        )
+                    })?;
+
+                    params.insert(name.to_owned(), value.to_owned());
+                }
+
+                FeStartupPacket::StartupMessage {
+                    major_version,
+                    minor_version,
+                    params: StartupMessageParams { params },
+                }
+            }
+        };
+        Ok(Some(FeMessage::StartupPacket(message)))
+    }
+
     /// Read startup message from the stream.
     // XXX: It's tempting yet undesirable to accept `stream` by value,
     // since such a change will cause user-supplied &mut references to be consumed
     pub fn read_fut<Reader>(
         stream: &mut Reader,
-    ) -> SyncFuture<Reader, impl Future<Output = Result<Option<FeMessage>, ConnectionError>> + '_>
+    ) -> SyncFuture<Reader, impl Future<Output = Result<Option<FeMessage>, ProtocolError>> + '_>
     where
         Reader: tokio::io::AsyncRead + Unpin,
     {
+        use tokio::io::AsyncReadExt;
+
         const MAX_STARTUP_PACKET_LENGTH: usize = 10000;
         const RESERVED_INVALID_MAJOR_VERSION: u32 = 1234;
         const CANCEL_REQUEST_CODE: u32 = 5678;
@@ -343,18 +460,18 @@ impl FeStartupPacket {
             let len = match retry_read!(stream.read_u32().await) {
                 Ok(len) => len as usize,
                 Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
-                Err(e) => return Err(ConnectionError::Socket(e)),
+                Err(e) => return Err(ProtocolError::Socket(e)),
             };
 
             #[allow(clippy::manual_range_contains)]
             if len < 4 || len > MAX_STARTUP_PACKET_LENGTH {
-                return Err(ConnectionError::Protocol(format!(
+                return Err(ProtocolError::Protocol(format!(
                     "invalid message length {len}"
                 )));
             }
 
             let request_code =
-                retry_read!(stream.read_u32().await).map_err(ConnectionError::Socket)?;
+                retry_read!(stream.read_u32().await).map_err(ProtocolError::Socket)?;
 
             // the rest of startup packet are params
             let params_len = len - 8;
@@ -362,7 +479,7 @@ impl FeStartupPacket {
             stream
                 .read_exact(params_bytes.as_mut())
                 .await
-                .map_err(ConnectionError::Socket)?;
+                .map_err(ProtocolError::Socket)?;
 
             // Parse params depending on request code
             let req_hi = request_code >> 16;
@@ -370,14 +487,16 @@ impl FeStartupPacket {
             let message = match (req_hi, req_lo) {
                 (RESERVED_INVALID_MAJOR_VERSION, CANCEL_REQUEST_CODE) => {
                     if params_len != 8 {
-                        return Err(ConnectionError::Protocol(
+                        return Err(ProtocolError::Protocol(
                             "expected 8 bytes for CancelRequest params".to_string(),
                         ));
                     }
                     let mut cursor = Cursor::new(params_bytes);
                     FeStartupPacket::CancelRequest(CancelKeyData {
-                        backend_pid: cursor.read_i32().await.map_err(ConnectionError::Socket)?,
-                        cancel_key: cursor.read_i32().await.map_err(ConnectionError::Socket)?,
+                        backend_pid: 2,
+                        cancel_key: 2,
+                        // backend_pid: cursor.read_i32().await.map_err(ConnectionError::Socket)?,
+                        // cancel_key: cursor.read_i32().await.map_err(ConnectionError::Socket)?,
                     })
                 }
                 (RESERVED_INVALID_MAJOR_VERSION, NEGOTIATE_SSL_CODE) => {
@@ -389,7 +508,7 @@ impl FeStartupPacket {
                     FeStartupPacket::GssEncRequest
                 }
                 (RESERVED_INVALID_MAJOR_VERSION, unrecognized_code) => {
-                    return Err(ConnectionError::Protocol(format!(
+                    return Err(ProtocolError::Protocol(format!(
                         "Unrecognized request code {unrecognized_code}"
                     )));
                 }
@@ -401,7 +520,7 @@ impl FeStartupPacket {
                         .context("StartupMessage params: invalid utf-8")?
                         .strip_suffix('\0') // drop packet's own null
                         .ok_or_else(|| {
-                            ConnectionError::Protocol(
+                            ProtocolError::Protocol(
                                 "StartupMessage params: missing null terminator".to_string(),
                             )
                         })?
@@ -410,7 +529,7 @@ impl FeStartupPacket {
                     let mut params = HashMap::new();
                     while let Some(name) = tokens.next() {
                         let value = tokens.next().ok_or_else(|| {
-                            ConnectionError::Protocol(
+                            ProtocolError::Protocol(
                                 "StartupMessage params: key without value".to_string(),
                             )
                         })?;
@@ -440,6 +559,9 @@ impl FeParseMessage {
 
         let _pstmt_name = read_cstr(&mut buf)?;
         let query_string = read_cstr(&mut buf)?;
+        if buf.remaining() < 2 {
+            bail!("Parse message is malformed, nparams missing");
+        }
         let nparams = buf.get_i16();
 
         ensure!(nparams == 0, "query params not implemented");
@@ -466,6 +588,9 @@ impl FeDescribeMessage {
 impl FeExecuteMessage {
     fn parse(mut buf: Bytes) -> anyhow::Result<FeMessage> {
         let portal_name = read_cstr(&mut buf)?;
+        if buf.remaining() < 4 {
+            bail!("FeExecuteMessage message is malformed, maxrows missing");
+        }
         let maxrows = buf.get_i32();
 
         ensure!(portal_name.is_empty(), "named portals not implemented");
@@ -547,6 +672,11 @@ impl<'a> BeMessage<'a> {
         value: b"UTF8",
     };
 
+    pub const INTEGER_DATETIMES: Self = Self::ParameterStatus {
+        name: b"integer_datetimes",
+        value: b"on",
+    };
+
     /// Build a [`BeMessage::ParameterStatus`] holding the server version.
     pub fn server_version(version: &'a str) -> Self {
         Self::ParameterStatus {
@@ -665,13 +795,12 @@ fn write_body<R>(buf: &mut BytesMut, f: impl FnOnce(&mut BytesMut) -> R) -> R {
 }
 
 /// Safe write of s into buf as cstring (String in the protocol).
-fn write_cstr(s: impl AsRef<[u8]>, buf: &mut BytesMut) -> io::Result<()> {
+fn write_cstr(s: impl AsRef<[u8]>, buf: &mut BytesMut) -> Result<(), ProtocolError> {
     let bytes = s.as_ref();
     if bytes.contains(&0) {
-        return Err(io::Error::new(
-            io::ErrorKind::InvalidInput,
-            "string contains embedded null",
-        ));
+        return Err(ProtocolError::MessageParse(anyhow!(
+            "string contains embedded null"
+        )));
     }
     buf.put_slice(bytes);
     buf.put_u8(0);
@@ -680,7 +809,7 @@ fn write_cstr(s: impl AsRef<[u8]>, buf: &mut BytesMut) -> io::Result<()> {
 
 fn read_cstr(buf: &mut Bytes) -> anyhow::Result<Bytes> {
     let pos = buf.iter().position(|x| *x == 0);
-    let result = buf.split_to(pos.context("missing terminator")?);
+    let result = buf.split_to(pos.context("missing cstring terminator")?);
     buf.advance(1); // drop the null terminator
     Ok(result)
 }
@@ -688,12 +817,12 @@ fn read_cstr(buf: &mut Bytes) -> anyhow::Result<Bytes> {
 pub const SQLSTATE_INTERNAL_ERROR: &[u8; 5] = b"XX000";
 
 impl<'a> BeMessage<'a> {
-    /// Write message to the given buf.
-    // Unlike the reading side, we use BytesMut
-    // here as msg len precedes its body and it is handy to write it down first
-    // and then fill the length. With Write we would have to either calc it
-    // manually or have one more buffer.
-    pub fn write(buf: &mut BytesMut, message: &BeMessage) -> io::Result<()> {
+    /// Serialize `message` to the given `buf`.
+    /// Apart from smart memory managemet, BytesMut is good here as msg len
+    /// precedes its body and it is handy to write it down first and then fill
+    /// the length. With Write we would have to either calc it manually or have
+    /// one more buffer.
+    pub fn write(buf: &mut BytesMut, message: &BeMessage) -> Result<(), ProtocolError> {
         match message {
             BeMessage::AuthenticationOk => {
                 buf.put_u8(b'R');
@@ -719,7 +848,7 @@ impl<'a> BeMessage<'a> {
 
             BeMessage::AuthenticationSasl(msg) => {
                 buf.put_u8(b'R');
-                write_body(buf, |buf| {
+                write_body(buf, |buf| -> Result<(), ProtocolError> {
                     use BeAuthenticationSaslMessage::*;
                     match msg {
                         Methods(methods) => {
@@ -738,7 +867,7 @@ impl<'a> BeMessage<'a> {
                             buf.put_slice(extra);
                         }
                     }
-                    Ok::<_, io::Error>(())
+                    Ok(())
                 })?;
             }
 
@@ -829,7 +958,7 @@ impl<'a> BeMessage<'a> {
             BeMessage::ErrorResponse(error_msg, pg_error_code) => {
                 // 'E' signalizes ErrorResponse messages
                 buf.put_u8(b'E');
-                write_body(buf, |buf| {
+                write_body(buf, |buf| -> Result<(), ProtocolError> {
                     buf.put_u8(b'S'); // severity
                     buf.put_slice(b"ERROR\0");
 
@@ -842,7 +971,7 @@ impl<'a> BeMessage<'a> {
                     write_cstr(error_msg, buf)?;
 
                     buf.put_u8(0); // terminator
-                    Ok::<_, io::Error>(())
+                    Ok(())
                 })?;
             }
 
@@ -854,7 +983,7 @@ impl<'a> BeMessage<'a> {
 
                 // 'N' signalizes NoticeResponse messages
                 buf.put_u8(b'N');
-                write_body(buf, |buf| {
+                write_body(buf, |buf| -> Result<(), ProtocolError> {
                     buf.put_u8(b'S'); // severity
                     buf.put_slice(b"NOTICE\0");
 
@@ -865,7 +994,7 @@ impl<'a> BeMessage<'a> {
                     write_cstr(error_msg.as_bytes(), buf)?;
 
                     buf.put_u8(0); // terminator
-                    Ok::<_, io::Error>(())
+                    Ok(())
                 })?;
             }
 
@@ -909,7 +1038,7 @@ impl<'a> BeMessage<'a> {
 
             BeMessage::RowDescription(rows) => {
                 buf.put_u8(b'T');
-                write_body(buf, |buf| {
+                write_body(buf, |buf| -> Result<(), ProtocolError> {
                     buf.put_i16(rows.len() as i16); // # of fields
                     for row in rows.iter() {
                         write_cstr(row.name, buf)?;
@@ -920,7 +1049,7 @@ impl<'a> BeMessage<'a> {
                         buf.put_i32(-1); /* typmod */
                         buf.put_i16(0); /* format code */
                     }
-                    Ok::<_, io::Error>(())
+                    Ok(())
                 })?;
             }
 

libs/remote_storage/src/lib.rs

@@ -111,7 +111,7 @@ pub trait RemoteStorage: Send + Sync + 'static {
 }
 
 pub struct Download {
-    pub download_stream: Pin<Box<dyn io::AsyncRead + Unpin + Send>>,
+    pub download_stream: Pin<Box<dyn io::AsyncRead + Unpin + Send + Sync>>,
     /// Extra key-value data, associated with the current remote file.
     pub metadata: Option<StorageMetadata>,
 }

libs/utils/Cargo.toml

@@ -10,13 +10,16 @@ async-trait = "0.1"
 anyhow = "1.0"
 bincode = "1.3"
 bytes = "1.0.1"
+futures = "0.3"
 hyper = { version = "0.14.7", features = ["full"] }
+pin-utils = "0.1"
 routerify = "3"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1"
 thiserror = "1.0"
 tokio = { version = "1.17", features = ["macros"]}
 tokio-rustls = "0.23"
+tokio-util = { version = "0.7.3" }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
 nix = "0.25"

libs/utils/src/lib.rs

@@ -13,7 +13,7 @@ pub mod simple_rcu;
 pub mod vec_map;
 
 pub mod bin_ser;
-pub mod postgres_backend;
+// pub mod postgres_backend;
 pub mod postgres_backend_async;
 
 // helper functions for creating and fsyncing
@@ -52,6 +52,8 @@ pub mod signals;
 
 pub mod fs_ext;
 
+pub mod send_rc;
+
 /// use with fail::cfg("$name", "return(2000)")
 #[macro_export]
 macro_rules! failpoint_sleep_millis_async {

libs/utils/src/postgres_backend_async.rs

@@ -2,29 +2,24 @@
 //! To use, create PostgresBackend and run() it, passing the Handler
 //! implementation determining how to process the queries. Currently its API
 //! is rather narrow, but we can extend it once required.
-
-use crate::postgres_backend::AuthType;
 use anyhow::Context;
 use bytes::{Buf, Bytes, BytesMut};
-use pq_proto::{BeMessage, ConnectionError, FeMessage, FeStartupPacket, SQLSTATE_INTERNAL_ERROR};
-use std::future::Future;
-use std::io;
+use futures::stream::StreamExt;
+use futures::{pin_mut, Sink, SinkExt};
+use serde::{Deserialize, Serialize};
 use std::net::SocketAddr;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::Poll;
-use tracing::{debug, error, info, trace};
-
+use std::{fmt, io};
+use std::{future::Future, str::FromStr};
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
 use tokio_rustls::TlsAcceptor;
+use tokio_util::codec::Framed;
+use tracing::{debug, error, info, trace};
 
-pub fn is_expected_io_error(e: &io::Error) -> bool {
-    use io::ErrorKind::*;
-    matches!(
-        e.kind(),
-        ConnectionRefused | ConnectionAborted | ConnectionReset
-    )
-}
+use pq_proto::codec::{ConnectionError, PostgresCodec};
+use pq_proto::{BeMessage, FeMessage, FeStartupPacket, SQLSTATE_INTERNAL_ERROR};
 
 /// An error, occurred during query processing:
 /// either during the connection ([`ConnectionError`]) or before/after it.
@@ -40,7 +35,7 @@ pub enum QueryError {
 
 impl From<io::Error> for QueryError {
     fn from(e: io::Error) -> Self {
-        Self::Disconnected(ConnectionError::Socket(e))
+        Self::Disconnected(ConnectionError::Io(e))
     }
 }
 
@@ -53,6 +48,14 @@ impl QueryError {
     }
 }
 
+pub fn is_expected_io_error(e: &io::Error) -> bool {
+    use io::ErrorKind::*;
+    matches!(
+        e.kind(),
+        ConnectionRefused | ConnectionAborted | ConnectionReset
+    )
+}
+
 #[async_trait::async_trait]
 pub trait Handler {
     /// Handle single query.
@@ -93,6 +96,7 @@ pub trait Handler {
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd)]
 pub enum ProtoState {
     Initialization,
+    // Encryption handshake is done; waiting for encrypted Startup message.
     Encrypted,
     Authentication,
     Established,
@@ -105,15 +109,14 @@ pub enum ProcessMsgResult {
     Break,
 }
 
-/// Always-writeable sock_split stream.
-/// May not be readable. See [`PostgresBackend::take_stream_in`]
-pub enum Stream {
-    Unencrypted(BufReader<tokio::net::TcpStream>),
-    Tls(Box<tokio_rustls::server::TlsStream<BufReader<tokio::net::TcpStream>>>),
-    Broken,
+/// Either plain TCP stream or encrypted one, implementing AsyncRead + AsyncWrite.
+pub enum MaybeTlsStream {
+    Unencrypted(tokio::net::TcpStream),
+    Tls(Box<tokio_rustls::server::TlsStream<tokio::net::TcpStream>>),
+    Broken, // temporary value for switch to TLS
 }
 
-impl AsyncWrite for Stream {
+impl AsyncWrite for MaybeTlsStream {
     fn poll_write(
         self: Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
@@ -122,14 +125,14 @@ impl AsyncWrite for Stream {
         match self.get_mut() {
             Self::Unencrypted(stream) => Pin::new(stream).poll_write(cx, buf),
             Self::Tls(stream) => Pin::new(stream).poll_write(cx, buf),
-            Self::Broken => unreachable!(),
+            _ => unreachable!(),
         }
     }
     fn poll_flush(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<io::Result<()>> {
         match self.get_mut() {
             Self::Unencrypted(stream) => Pin::new(stream).poll_flush(cx),
             Self::Tls(stream) => Pin::new(stream).poll_flush(cx),
-            Self::Broken => unreachable!(),
+            _ => unreachable!(),
         }
     }
     fn poll_shutdown(
@@ -139,11 +142,11 @@ impl AsyncWrite for Stream {
         match self.get_mut() {
             Self::Unencrypted(stream) => Pin::new(stream).poll_shutdown(cx),
             Self::Tls(stream) => Pin::new(stream).poll_shutdown(cx),
-            Self::Broken => unreachable!(),
+            _ => unreachable!(),
         }
     }
 }
-impl AsyncRead for Stream {
+impl AsyncRead for MaybeTlsStream {
     fn poll_read(
         self: Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
@@ -152,18 +155,49 @@ impl AsyncRead for Stream {
         match self.get_mut() {
             Self::Unencrypted(stream) => Pin::new(stream).poll_read(cx, buf),
             Self::Tls(stream) => Pin::new(stream).poll_read(cx, buf),
-            Self::Broken => unreachable!(),
+            _ => unreachable!(),
         }
     }
 }
 
-pub struct PostgresBackend {
-    stream: Stream,
+#[derive(Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize)]
+pub enum AuthType {
+    Trust,
+    // This mimics postgres's AuthenticationCleartextPassword but instead of password expects JWT
+    NeonJWT,
+}
 
-    // Output buffer. c.f. BeMessage::write why we are using BytesMut here.
-    // The data between 0 and "current position" as tracked by the bytes::Buf
-    // implementation of BytesMut, have already been written.
-    buf_out: BytesMut,
+impl FromStr for AuthType {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "Trust" => Ok(Self::Trust),
+            "NeonJWT" => Ok(Self::NeonJWT),
+            _ => anyhow::bail!("invalid value \"{s}\" for auth type"),
+        }
+    }
+}
+
+impl fmt::Display for AuthType {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(match self {
+            AuthType::Trust => "Trust",
+            AuthType::NeonJWT => "NeonJWT",
+        })
+    }
+}
+
+pub struct PostgresBackend {
+    // Provides serialization/deserialization to the underlying transport backed
+    // with buffers; implements Sink consuming messages and Stream reading them.
+    //
+    // Sink::start_send only queues message to the interal buffer.
+    // SinkExt::flush flushes buffer to the stream.
+    //
+    // StreamExt::read reads next message. In case of EOF without partial
+    // message it returns None.
+    stream: Framed<MaybeTlsStream, PostgresCodec>,
 
     pub state: ProtoState,
 
@@ -196,10 +230,10 @@ impl PostgresBackend {
         tls_config: Option<Arc<rustls::ServerConfig>>,
     ) -> io::Result<Self> {
         let peer_addr = socket.peer_addr()?;
+        let stream = MaybeTlsStream::Unencrypted(socket);
 
         Ok(Self {
-            stream: Stream::Unencrypted(BufReader::new(socket)),
-            buf_out: BytesMut::with_capacity(10 * 1024),
+            stream: Framed::new(stream, PostgresCodec::new()),
             state: ProtoState::Initialization,
             auth_type,
             tls_config,
@@ -212,29 +246,60 @@ impl PostgresBackend {
     }
 
     /// Read full message or return None if connection is closed.
-    pub async fn read_message(&mut self) -> Result<Option<FeMessage>, QueryError> {
-        use ProtoState::*;
-        match self.state {
-            Initialization | Encrypted => FeStartupPacket::read_fut(&mut self.stream).await,
-            Authentication | Established => FeMessage::read_fut(&mut self.stream).await,
-            Closed => Ok(None),
+    pub async fn read_message(&mut self) -> Result<Option<FeMessage>, ConnectionError> {
+        if let ProtoState::Closed = self.state {
+            Ok(None)
+        } else {
+            let msg = self.stream.next().await;
+            // Option<Result<...>>, so swap.
+            msg.map_or(Ok(None), |res| res.map(Some))
         }
-        .map_err(QueryError::from)
+    }
+
+    /// Polling version of read_message, saves the caller need to pin.
+    pub fn poll_read_message(
+        &mut self,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<Option<FeMessage>, ConnectionError>> {
+        let read_fut = self.read_message();
+        pin_mut!(read_fut);
+        read_fut.poll(cx)
     }
 
     /// Flush output buffer into the socket.
     pub async fn flush(&mut self) -> io::Result<()> {
-        while self.buf_out.has_remaining() {
-            let bytes_written = self.stream.write(self.buf_out.chunk()).await?;
-            self.buf_out.advance(bytes_written);
-        }
-        self.buf_out.clear();
-        Ok(())
+        self.stream.flush().await.map_err(|e| match e {
+            ConnectionError::Io(e) => e,
+            // the only error we can get from flushing is IO
+            _ => unreachable!(),
+        })
     }
 
-    /// Write message into internal output buffer.
-    pub fn write_message(&mut self, message: &BeMessage<'_>) -> io::Result<&mut Self> {
-        BeMessage::write(&mut self.buf_out, message)?;
+    /// Polling version of `flush()`, saves the caller need to pin.
+    pub fn poll_flush(
+        &mut self,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        let flush_fut = self.flush();
+        pin_mut!(flush_fut);
+        flush_fut.poll(cx)
+    }
+
+    /// Write message into internal output buffer. Technically error type can be
+    /// only ProtocolError here (if, unlikely, serialization fails), but callers
+    /// typically wrap it anyway.
+    pub fn write_message(&mut self, message: &BeMessage<'_>) -> Result<&mut Self, ConnectionError> {
+        Pin::new(&mut self.stream).start_send(message)?;
+        Ok(self)
+    }
+
+    /// Write message into internal output buffer and flush it to the stream.
+    pub async fn write_message_flush(
+        &mut self,
+        message: &BeMessage<'_>,
+    ) -> Result<&mut Self, ConnectionError> {
+        self.write_message(message)?;
+        self.flush().await?;
         Ok(self)
     }
 
@@ -246,28 +311,6 @@ impl PostgresBackend {
         CopyDataWriter { pgb: self }
     }
 
-    /// A polling function that tries to write all the data from 'buf_out' to the
-    /// underlying stream.
-    fn poll_write_buf(
-        &mut self,
-        cx: &mut std::task::Context<'_>,
-    ) -> Poll<Result<(), std::io::Error>> {
-        while self.buf_out.has_remaining() {
-            match Pin::new(&mut self.stream).poll_write(cx, self.buf_out.chunk()) {
-                Poll::Ready(Ok(bytes_written)) => {
-                    self.buf_out.advance(bytes_written);
-                }
-                Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
-                Poll::Pending => return Poll::Pending,
-            }
-        }
-        Poll::Ready(Ok(()))
-    }
-
-    fn poll_flush(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), std::io::Error>> {
-        Pin::new(&mut self.stream).poll_flush(cx)
-    }
-
     // Wrapper for run_message_loop() that shuts down socket when we are done
     pub async fn run<F, S>(
         mut self,
@@ -279,7 +322,7 @@ impl PostgresBackend {
         S: Future,
     {
         let ret = self.run_message_loop(handler, shutdown_watcher).await;
-        let _ = self.stream.shutdown();
+        let _ = self.stream.get_mut().shutdown();
         ret
     }
 
@@ -359,14 +402,22 @@ impl PostgresBackend {
     }
 
     async fn start_tls(&mut self) -> anyhow::Result<()> {
-        if let Stream::Unencrypted(plain_stream) =
-            std::mem::replace(&mut self.stream, Stream::Broken)
+        if let MaybeTlsStream::Unencrypted(plain_stream) =
+            // temporary replace stream with fake broken to prepare TLS one
+            std::mem::replace(self.stream.get_mut(), MaybeTlsStream::Broken)
         {
             let acceptor = TlsAcceptor::from(self.tls_config.clone().unwrap());
-            let tls_stream = acceptor.accept(plain_stream).await?;
-
-            self.stream = Stream::Tls(Box::new(tls_stream));
-            return Ok(());
+            match acceptor.accept(plain_stream).await {
+                Ok(tls_stream) => {
+                    // push back ready TLS stream
+                    *self.stream.get_mut() = MaybeTlsStream::Tls(Box::new(tls_stream));
+                    return Ok(());
+                }
+                Err(e) => {
+                    self.state = ProtoState::Closed;
+                    return Err(e.into());
+                }
+            }
         };
         anyhow::bail!("TLS already started");
     }
@@ -380,13 +431,12 @@ impl PostgresBackend {
         let have_tls = self.tls_config.is_some();
         match msg {
             FeMessage::StartupPacket(m) => {
-                trace!("got startup message {m:?}");
-
                 match m {
                     FeStartupPacket::SslRequest => {
                         debug!("SSL requested");
 
                         self.write_message(&BeMessage::EncryptionResponse(have_tls))?;
+
                         if have_tls {
                             self.start_tls().await?;
                             self.state = ProtoState::Encrypted;
@@ -415,6 +465,7 @@ impl PostgresBackend {
                             AuthType::Trust => {
                                 self.write_message(&BeMessage::AuthenticationOk)?
                                     .write_message(&BeMessage::CLIENT_ENCODING)?
+                                    .write_message(&BeMessage::INTEGER_DATETIMES)?
                                     // The async python driver requires a valid server_version
                                     .write_message(&BeMessage::server_version("14.1"))?
                                     .write_message(&BeMessage::ReadyForQuery)?;
@@ -454,6 +505,7 @@ impl PostgresBackend {
                 }
                 self.write_message(&BeMessage::AuthenticationOk)?
                     .write_message(&BeMessage::CLIENT_ENCODING)?
+                    .write_message(&BeMessage::INTEGER_DATETIMES)?
                     .write_message(&BeMessage::ReadyForQuery)?;
                 self.state = ProtoState::Established;
             }
@@ -573,7 +625,7 @@ impl<'a> AsyncWrite for CopyDataWriter<'a> {
         // It's not strictly required to flush between each message, but makes it easier
         // to view in wireshark, and usually the messages that the callers write are
         // decently-sized anyway.
-        match this.pgb.poll_write_buf(cx) {
+        match this.pgb.poll_flush(cx) {
             Poll::Ready(Ok(())) => {}
             Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
             Poll::Pending => return Poll::Pending,
@@ -583,7 +635,11 @@ impl<'a> AsyncWrite for CopyDataWriter<'a> {
         // XXX: if the input is large, we should split it into multiple messages.
         // Not sure what the threshold should be, but the ultimate hard limit is that
         // the length cannot exceed u32.
-        this.pgb.write_message(&BeMessage::CopyData(buf))?;
+        this.pgb
+            .write_message(&BeMessage::CopyData(buf))
+            // write_message only writes to buffer, so can fail iff message is
+            // invaid, but CopyData can't be invalid.
+            .expect("failed to serialize CopyData");
 
         Poll::Ready(Ok(buf.len()))
     }
@@ -593,23 +649,14 @@ impl<'a> AsyncWrite for CopyDataWriter<'a> {
         cx: &mut std::task::Context<'_>,
     ) -> Poll<Result<(), std::io::Error>> {
         let this = self.get_mut();
-        match this.pgb.poll_write_buf(cx) {
-            Poll::Ready(Ok(())) => {}
-            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
-            Poll::Pending => return Poll::Pending,
-        }
         this.pgb.poll_flush(cx)
     }
+
     fn poll_shutdown(
         self: Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
     ) -> Poll<Result<(), std::io::Error>> {
         let this = self.get_mut();
-        match this.pgb.poll_write_buf(cx) {
-            Poll::Ready(Ok(())) => {}
-            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
-            Poll::Pending => return Poll::Pending,
-        }
         this.pgb.poll_flush(cx)
     }
 }
@@ -623,7 +670,7 @@ pub fn short_error(e: &QueryError) -> String {
 
 pub(super) fn log_query_error(query: &str, e: &QueryError) {
     match e {
-        QueryError::Disconnected(ConnectionError::Socket(io_error)) => {
+        QueryError::Disconnected(ConnectionError::Io(io_error)) => {
             if is_expected_io_error(io_error) {
                 info!("query handler for '{query}' failed with expected io error: {io_error}");
             } else {

libs/utils/src/send_rc.rs

@@ -0,0 +1,116 @@
+/// Provides Send wrappers of Rc and RefMut.
+use std::{
+    borrow::Borrow,
+    cell::{Ref, RefCell, RefMut},
+    ops::{Deref, DerefMut},
+    rc::Rc,
+};
+
+/// Rc wrapper which is Send.
+/// This is useful to allow transferring a group of Rcs pointing to the same
+/// object between threads, e.g. in self referential struct.
+#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
+pub struct SendRc<T>
+where
+    T: ?Sized,
+{
+    rc: Rc<T>,
+}
+
+// SAFETY: Passing Rc(s)<T: Send> between threads is fine as long as there is no
+// concurrent access to the object they point to, so you must move all such Rcs
+// together. This appears to be impossible to express in rust type system and
+// SendRc doesn't provide any additional protection -- but unlike sendable
+// crate, neither it requires any additional actions before/after move. Ensuring
+// that sending conforms to the above is the responsibility of the type user.
+unsafe impl<T: ?Sized + Send> Send for SendRc<T> {}
+
+impl<T> SendRc<T> {
+    /// Constructs a new SendRc<T>
+    pub fn new(value: T) -> SendRc<T> {
+        SendRc { rc: Rc::new(value) }
+    }
+}
+
+// https://stegosaurusdormant.com/understanding-derive-clone/ explains in detail
+// why derive Clone doesn't work here.
+impl<T> Clone for SendRc<T> {
+    fn clone(&self) -> Self {
+        SendRc {
+            rc: self.rc.clone(),
+        }
+    }
+}
+
+// Deref into inner rc.
+impl<T> Deref for SendRc<T> {
+    type Target = Rc<T>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.rc
+    }
+}
+
+/// Extends RefCell with borrow[_mut] variants which return Sendable Ref[Mut]
+/// wrappers.
+pub trait RefCellSend<T: ?Sized> {
+    fn borrow_mut_send(&self) -> RefMutSend<'_, T>;
+}
+
+impl<T: Sized> RefCellSend<T> for RefCell<T> {
+    fn borrow_mut_send(&self) -> RefMutSend<'_, T> {
+        RefMutSend {
+            ref_mut: self.borrow_mut(),
+        }
+    }
+}
+
+/// RefMut wrapper which is Send. See impl Send for safety. Allows to move a
+/// RefMut along with RefCell it originates from between threads, e.g. have Send
+/// Future containing RefMut.
+#[derive(Debug)]
+pub struct RefMutSend<'b, T>
+where
+    T: 'b + ?Sized,
+{
+    ref_mut: RefMut<'b, T>,
+}
+
+// SAFETY: Similar to SendRc, this is safe as long as RefMut stays in the same
+// thread with original RefCell, so they should be passed together.
+// Actually, since this is a referential type violating this is not
+// straightforward; examples of unsafe usage could be
+// - Passing a RefMut to different thread without source RefCell. Seems only
+//   possible with std::thread::scope.
+// - Somehow multiple threads get access to single RefCell concurrently,
+//   violating its !Sync requirement. Improper usage of SendRc can do that.
+unsafe impl<'b, T: ?Sized + Send> Send for RefMutSend<'b, T> {}
+
+impl<'b, T> RefMutSend<'b, T> {
+    /// Constructs a new RefMutSend<T>
+    pub fn new(ref_mut: RefMut<'b, T>) -> RefMutSend<'b, T> {
+        RefMutSend { ref_mut }
+    }
+}
+
+// Deref into inner RefMut.
+impl<'b, T> Deref for RefMutSend<'b, T>
+where
+    T: 'b + ?Sized,
+{
+    type Target = RefMut<'b, T>;
+
+    fn deref<'a>(&'a self) -> &'a RefMut<'b, T> {
+        &self.ref_mut
+    }
+}
+
+// DerefMut into inner RefMut.
+impl<'b, T> DerefMut for RefMutSend<'b, T>
+where
+    T: 'b + ?Sized,
+{
+    fn deref_mut<'a>(&'a mut self) -> &'a mut RefMut<'b, T> {
+        &mut self.ref_mut
+    }
+}

pageserver/Cargo.toml

@@ -67,7 +67,6 @@ storage_broker = { version = "0.1", path = "../storage_broker" }
 tenant_size_model = { path = "../libs/tenant_size_model" }
 utils = { path = "../libs/utils" }
 workspace_hack = { version = "0.1", path = "../workspace_hack" }
-im = "15.1.0"
 reqwest = { version = "0.11", default-features = false, features = ["rustls-tls"] }
 
 [dev-dependencies]

pageserver/benches/bench_layer_map.rs

@@ -1,12 +1,13 @@
-use pageserver::keyspace::{KeyPartitioning, KeySpace};
+use anyhow::Result;
 use pageserver::repository::Key;
 use pageserver::tenant::layer_map::LayerMap;
-use pageserver::tenant::storage_layer::Layer;
-use pageserver::tenant::storage_layer::{DeltaFileName, ImageFileName, LayerDescriptor};
+use pageserver::tenant::storage_layer::{DeltaFileName, ImageFileName, ValueReconstructState};
+use pageserver::tenant::storage_layer::{Layer, ValueReconstructResult};
 use rand::prelude::{SeedableRng, SliceRandom, StdRng};
 use std::cmp::{max, min};
 use std::fs::File;
 use std::io::{BufRead, BufReader};
+use std::ops::Range;
 use std::path::PathBuf;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -16,8 +17,80 @@ use utils::lsn::Lsn;
 
 use criterion::{criterion_group, criterion_main, Criterion};
 
-fn build_layer_map(filename_dump: PathBuf) -> LayerMap<LayerDescriptor> {
-    let mut layer_map = LayerMap::<LayerDescriptor>::default();
+struct DummyDelta {
+    key_range: Range<Key>,
+    lsn_range: Range<Lsn>,
+}
+
+impl Layer for DummyDelta {
+    fn get_key_range(&self) -> Range<Key> {
+        self.key_range.clone()
+    }
+
+    fn get_lsn_range(&self) -> Range<Lsn> {
+        self.lsn_range.clone()
+    }
+    fn get_value_reconstruct_data(
+        &self,
+        _key: Key,
+        _lsn_range: Range<Lsn>,
+        _reconstruct_data: &mut ValueReconstructState,
+    ) -> Result<ValueReconstructResult> {
+        panic!()
+    }
+
+    fn is_incremental(&self) -> bool {
+        true
+    }
+
+    fn dump(&self, _verbose: bool) -> Result<()> {
+        unimplemented!()
+    }
+
+    fn short_id(&self) -> String {
+        unimplemented!()
+    }
+}
+
+struct DummyImage {
+    key_range: Range<Key>,
+    lsn: Lsn,
+}
+
+impl Layer for DummyImage {
+    fn get_key_range(&self) -> Range<Key> {
+        self.key_range.clone()
+    }
+
+    fn get_lsn_range(&self) -> Range<Lsn> {
+        // End-bound is exclusive
+        self.lsn..(self.lsn + 1)
+    }
+
+    fn get_value_reconstruct_data(
+        &self,
+        _key: Key,
+        _lsn_range: Range<Lsn>,
+        _reconstruct_data: &mut ValueReconstructState,
+    ) -> Result<ValueReconstructResult> {
+        panic!()
+    }
+
+    fn is_incremental(&self) -> bool {
+        false
+    }
+
+    fn dump(&self, _verbose: bool) -> Result<()> {
+        unimplemented!()
+    }
+
+    fn short_id(&self) -> String {
+        unimplemented!()
+    }
+}
+
+fn build_layer_map(filename_dump: PathBuf) -> LayerMap<dyn Layer> {
+    let mut layer_map = LayerMap::<dyn Layer>::default();
 
     let mut min_lsn = Lsn(u64::MAX);
     let mut max_lsn = Lsn(0);
@@ -27,21 +100,17 @@ fn build_layer_map(filename_dump: PathBuf) -> LayerMap<LayerDescriptor> {
     for fname in filenames {
         let fname = &fname.unwrap();
         if let Some(imgfilename) = ImageFileName::parse_str(fname) {
-            let layer = LayerDescriptor {
-                key: imgfilename.key_range,
-                lsn: imgfilename.lsn..(imgfilename.lsn + 1),
-                is_incremental: false,
-                short_id: fname.to_string(),
+            let layer = DummyImage {
+                key_range: imgfilename.key_range,
+                lsn: imgfilename.lsn,
             };
             layer_map.insert_historic(Arc::new(layer));
             min_lsn = min(min_lsn, imgfilename.lsn);
             max_lsn = max(max_lsn, imgfilename.lsn);
         } else if let Some(deltafilename) = DeltaFileName::parse_str(fname) {
-            let layer = LayerDescriptor {
-                key: deltafilename.key_range.clone(),
-                lsn: deltafilename.lsn_range.clone(),
-                is_incremental: true,
-                short_id: fname.to_string(),
+            let layer = DummyDelta {
+                key_range: deltafilename.key_range,
+                lsn_range: deltafilename.lsn_range.clone(),
             };
             layer_map.insert_historic(Arc::new(layer));
             min_lsn = min(min_lsn, deltafilename.lsn_range.start);
@@ -50,7 +119,6 @@ fn build_layer_map(filename_dump: PathBuf) -> LayerMap<LayerDescriptor> {
             panic!("unexpected filename {fname}");
         }
     }
-    layer_map.rebuild_index();
 
     println!("min: {min_lsn}, max: {max_lsn}");
 
@@ -58,7 +126,7 @@ fn build_layer_map(filename_dump: PathBuf) -> LayerMap<LayerDescriptor> {
 }
 
 /// Construct a layer map query pattern for benchmarks
-fn uniform_query_pattern(layer_map: &LayerMap<LayerDescriptor>) -> Vec<(Key, Lsn)> {
+fn uniform_query_pattern(layer_map: &LayerMap<dyn Layer>) -> Vec<(Key, Lsn)> {
     // For each image layer we query one of the pages contained, at LSN right
     // before the image layer was created. This gives us a somewhat uniform
     // coverage of both the lsn and key space because image layers have
@@ -82,41 +150,6 @@ fn uniform_query_pattern(layer_map: &LayerMap<LayerDescriptor>) -> Vec<(Key, Lsn
         .collect()
 }
 
-// Construct a partitioning for testing get_difficulty map when we
-// don't have an exact result of `collect_keyspace` to work with.
-fn uniform_key_partitioning(layer_map: &LayerMap<LayerDescriptor>, _lsn: Lsn) -> KeyPartitioning {
-    let mut parts = Vec::new();
-
-    // We add a partition boundary at the start of each image layer,
-    // no matter what lsn range it covers. This is just the easiest
-    // thing to do. A better thing to do would be to get a real
-    // partitioning from some database. Even better, remove the need
-    // for key partitions by deciding where to create image layers
-    // directly based on a coverage-based difficulty map.
-    let mut keys: Vec<_> = layer_map
-        .iter_historic_layers()
-        .filter_map(|l| {
-            if l.is_incremental() {
-                None
-            } else {
-                let kr = l.get_key_range();
-                Some(kr.start.next())
-            }
-        })
-        .collect();
-    keys.sort();
-
-    let mut current_key = Key::from_hex("000000000000000000000000000000000000").unwrap();
-    for key in keys {
-        parts.push(KeySpace {
-            ranges: vec![current_key..key],
-        });
-        current_key = key;
-    }
-
-    KeyPartitioning { parts }
-}
-
 // Benchmark using metadata extracted from our performance test environment, from
 // a project where we have run pgbench many timmes. The pgbench database was initialized
 // between each test run.
@@ -150,71 +183,24 @@ fn bench_from_captest_env(c: &mut Criterion) {
 // Benchmark using metadata extracted from a real project that was taknig
 // too long processing layer map queries.
 fn bench_from_real_project(c: &mut Criterion) {
-    // Init layer map
-    let now = Instant::now();
+    // TODO consider compressing this file
     let layer_map = build_layer_map(PathBuf::from("benches/odd-brook-layernames.txt"));
-    println!("Finished layer map init in {:?}", now.elapsed());
-
-    // Choose uniformly distributed queries
     let queries: Vec<(Key, Lsn)> = uniform_query_pattern(&layer_map);
 
-    // Choose inputs for get_difficulty_map
-    let latest_lsn = layer_map
-        .iter_historic_layers()
-        .map(|l| l.get_lsn_range().end)
-        .max()
-        .unwrap();
-    let partitioning = uniform_key_partitioning(&layer_map, latest_lsn);
-
-    // Check correctness of get_difficulty_map
-    // TODO put this in a dedicated test outside of this mod
-    {
-        println!("running correctness check");
-
-        let now = Instant::now();
-        let result_bruteforce = layer_map.get_difficulty_map_bruteforce(latest_lsn, &partitioning);
-        assert!(result_bruteforce.len() == partitioning.parts.len());
-        println!("Finished bruteforce in {:?}", now.elapsed());
-
-        let now = Instant::now();
-        let result_fast = layer_map.get_difficulty_map(latest_lsn, &partitioning, None);
-        assert!(result_fast.len() == partitioning.parts.len());
-        println!("Finished fast in {:?}", now.elapsed());
-
-        // Assert results are equal. Manually iterate for easier debugging.
-        //
-        // TODO The difficulty numbers are in the 50-80 range. That's pretty bad.
-        //      We should be monitoring it in prod.
-        let zip = std::iter::zip(
-            &partitioning.parts,
-            std::iter::zip(result_bruteforce, result_fast),
-        );
-        for (_part, (bruteforce, fast)) in zip {
-            assert_eq!(bruteforce, fast);
-        }
-
-        println!("No issues found");
-    }
-
-    // Define and name the benchmark function
-    let mut group = c.benchmark_group("real_map");
-    group.bench_function("uniform_queries", |b| {
+    // Test with uniform query pattern
+    c.bench_function("real_map_uniform_queries", |b| {
         b.iter(|| {
             for q in queries.clone().into_iter() {
                 layer_map.search(q.0, q.1);
             }
         });
     });
-    group.bench_function("get_difficulty_map", |b| {
-        b.iter(|| {
-            layer_map.get_difficulty_map(latest_lsn, &partitioning, Some(3));
-        });
-    });
-    group.finish();
 }
 
 // Benchmark using synthetic data. Arrange image layers on stacked diagonal lines.
 fn bench_sequential(c: &mut Criterion) {
+    let mut layer_map: LayerMap<dyn Layer> = LayerMap::default();
+
     // Init layer map. Create 100_000 layers arranged in 1000 diagonal lines.
     //
     // TODO This code is pretty slow and runs even if we're only running other
@@ -222,38 +208,39 @@ fn bench_sequential(c: &mut Criterion) {
     //      Putting it inside the `bench_function` closure is not a solution
     //      because then it runs multiple times during warmup.
     let now = Instant::now();
-    let mut layer_map = LayerMap::default();
     for i in 0..100_000 {
+        // TODO try inserting a super-wide layer in between every 10 to reflect
+        //      what often happens with L1 layers that include non-rel changes.
+        //      Maybe do that as a separate test.
         let i32 = (i as u32) % 100;
         let zero = Key::from_hex("000000000000000000000000000000000000").unwrap();
-        let layer = LayerDescriptor {
-            key: zero.add(10 * i32)..zero.add(10 * i32 + 1),
-            lsn: Lsn(i)..Lsn(i + 1),
-            is_incremental: false,
-            short_id: format!("Layer {}", i),
+        let layer = DummyImage {
+            key_range: zero.add(10 * i32)..zero.add(10 * i32 + 1),
+            lsn: Lsn(10 * i),
         };
         layer_map.insert_historic(Arc::new(layer));
     }
-    layer_map.rebuild_index();
-    println!("Finished layer map init in {:?}", now.elapsed());
+
+    // Manually measure runtime without criterion because criterion
+    // has a minimum sample size of 10 and I don't want to run it 10 times.
+    println!("Finished init in {:?}", now.elapsed());
 
     // Choose 100 uniformly random queries
     let rng = &mut StdRng::seed_from_u64(1);
     let queries: Vec<(Key, Lsn)> = uniform_query_pattern(&layer_map)
-        .choose_multiple(rng, 100)
+        .choose_multiple(rng, 1)
         .copied()
         .collect();
 
     // Define and name the benchmark function
-    let mut group = c.benchmark_group("sequential");
-    group.bench_function("uniform_queries", |b| {
+    c.bench_function("sequential_uniform_queries", |b| {
+        // Run the search queries
         b.iter(|| {
             for q in queries.clone().into_iter() {
                 layer_map.search(q.0, q.1);
             }
         });
     });
-    group.finish();
 }
 
 criterion_group!(group_1, bench_from_captest_env);

pageserver/src/bin/pageserver.rs

@@ -24,7 +24,7 @@ use pageserver::{
 use utils::{
     auth::JwtAuth,
     logging,
-    postgres_backend::AuthType,
+    postgres_backend_async::AuthType,
     project_git_version,
     sentry_init::{init_sentry, release_name},
     signals::{self, Signal},

pageserver/src/config.rs

@@ -24,7 +24,7 @@ use toml_edit::{Document, Item};
 use utils::{
     id::{NodeId, TenantId, TimelineId},
     logging::LogFormat,
-    postgres_backend::AuthType,
+    postgres_backend_async::AuthType,
 };
 
 use crate::tenant::config::TenantConf;

pageserver/src/page_service.rs

@@ -19,7 +19,7 @@ use pageserver_api::models::{
     PagestreamFeMessage, PagestreamGetPageRequest, PagestreamGetPageResponse,
     PagestreamNblocksRequest, PagestreamNblocksResponse,
 };
-use pq_proto::ConnectionError;
+use pq_proto::codec::ConnectionError;
 use pq_proto::FeStartupPacket;
 use pq_proto::{BeMessage, FeMessage, RowDescriptor};
 use std::io;
@@ -35,7 +35,7 @@ use utils::{
     auth::{Claims, JwtAuth, Scope},
     id::{TenantId, TimelineId},
     lsn::Lsn,
-    postgres_backend::AuthType,
+    postgres_backend_async::AuthType,
     postgres_backend_async::{self, PostgresBackend},
     simple_rcu::RcuReadGuard,
 };
@@ -67,7 +67,7 @@ fn copyin_stream(pgb: &mut PostgresBackend) -> impl Stream<Item = io::Result<Byt
                     Err(QueryError::Other(anyhow::anyhow!(msg)))
                 }
 
-                msg = pgb.read_message() => { msg }
+                msg = pgb.read_message() => { msg.map_err(QueryError::from)}
             };
 
             match msg {
@@ -78,14 +78,16 @@ fn copyin_stream(pgb: &mut PostgresBackend) -> impl Stream<Item = io::Result<Byt
                         FeMessage::Sync => continue,
                         FeMessage::Terminate => {
                             let msg = "client terminated connection with Terminate message during COPY";
-                            let query_error_error = QueryError::Disconnected(ConnectionError::Socket(io::Error::new(io::ErrorKind::ConnectionReset, msg)));
-                            pgb.write_message(&BeMessage::ErrorResponse(msg, Some(query_error_error.pg_error_code())))?;
+                            let query_error = QueryError::Disconnected(ConnectionError::Io(io::Error::new(io::ErrorKind::ConnectionReset, msg)));
+                            pgb.write_message(&BeMessage::ErrorResponse(msg, Some(query_error.pg_error_code())))
+                                .expect("failed to serialize ErrorResponse");
                             Err(io::Error::new(io::ErrorKind::ConnectionReset, msg))?;
                             break;
                         }
                         m => {
                             let msg = format!("unexpected message {m:?}");
-                            pgb.write_message(&BeMessage::ErrorResponse(&msg, None))?;
+                            pgb.write_message(&BeMessage::ErrorResponse(&msg, None))
+                                .expect("failed to serialize ErrorResponse");
                             Err(io::Error::new(io::ErrorKind::Other, msg))?;
                             break;
                         }
@@ -95,16 +97,17 @@ fn copyin_stream(pgb: &mut PostgresBackend) -> impl Stream<Item = io::Result<Byt
                 }
                 Ok(None) => {
                     let msg = "client closed connection during COPY";
-                    let query_error_error = QueryError::Disconnected(ConnectionError::Socket(io::Error::new(io::ErrorKind::ConnectionReset, msg)));
-                    pgb.write_message(&BeMessage::ErrorResponse(msg, Some(query_error_error.pg_error_code())))?;
+                    let query_error = QueryError::Disconnected(ConnectionError::Io(io::Error::new(io::ErrorKind::ConnectionReset, msg)));
+                    pgb.write_message(&BeMessage::ErrorResponse(msg, Some(query_error.pg_error_code())))
+                        .expect("failed to serialize ErrorResponse");
                     pgb.flush().await?;
                     Err(io::Error::new(io::ErrorKind::ConnectionReset, msg))?;
                 }
-                Err(QueryError::Disconnected(ConnectionError::Socket(io_error))) => {
+                Err(QueryError::Disconnected(ConnectionError::Io(io_error))) => {
                     Err(io_error)?;
                 }
                 Err(other) => {
-                    Err(io::Error::new(io::ErrorKind::Other, other))?;
+                    Err(io::Error::new(io::ErrorKind::Other, other.to_string()))?;
                 }
             };
         }
@@ -202,7 +205,7 @@ async fn page_service_conn_main(
             // we've been requested to shut down
             Ok(())
         }
-        Err(QueryError::Disconnected(ConnectionError::Socket(io_error))) => {
+        Err(QueryError::Disconnected(ConnectionError::Io(io_error))) => {
             // `ConnectionReset` error happens when the Postgres client closes the connection.
             // As this disconnection happens quite often and is expected,
             // we decided to downgrade the logging level to `INFO`.

pageserver/src/repository.rs

@@ -37,17 +37,6 @@ impl Key {
             | self.field6 as i128
     }
 
-    pub fn from_i128(x: i128) -> Self {
-        Key {
-            field1: ((x >> 120) & 0xf) as u8,
-            field2: ((x >> 104) & 0xFFFF) as u32,
-            field3: (x >> 72) as u32,
-            field4: (x >> 40) as u32,
-            field5: (x >> 32) as u8,
-            field6: x as u32,
-        }
-    }
-
     pub fn next(&self) -> Key {
         self.add(1)
     }

pageserver/src/tenant/layer_map.rs

@@ -10,21 +10,23 @@
 //! corresponding files are written to disk.
 //!
 
-mod historic_layer_coverage;
-mod layer_coverage;
-
-use crate::keyspace::KeyPartitioning;
 use crate::metrics::NUM_ONDISK_LAYERS;
 use crate::repository::Key;
-use crate::tenant::storage_layer::InMemoryLayer;
-use crate::tenant::storage_layer::Layer;
+use crate::tenant::storage_layer::{range_eq, range_overlaps};
+use amplify_num::i256;
 use anyhow::Result;
+use num_traits::identities::{One, Zero};
+use num_traits::{Bounded, Num, Signed};
+use rstar::{RTree, RTreeObject, AABB};
+use std::cmp::Ordering;
 use std::collections::VecDeque;
 use std::ops::Range;
+use std::ops::{Add, Div, Mul, Neg, Rem, Sub};
 use std::sync::Arc;
+use tracing::*;
 use utils::lsn::Lsn;
 
-use historic_layer_coverage::BufferedHistoricLayerCoverage;
+use super::storage_layer::{InMemoryLayer, Layer};
 
 ///
 /// LayerMap tracks what layers exist on a timeline.
@@ -49,8 +51,8 @@ pub struct LayerMap<L: ?Sized> {
     ///
     pub frozen_layers: VecDeque<Arc<InMemoryLayer>>,
 
-    /// Index of the historic layers optimized for search
-    historic: BufferedHistoricLayerCoverage<Arc<L>>,
+    /// All the historic layers are kept here
+    historic_layers: RTree<LayerRTreeObject<L>>,
 
     /// L0 layers have key range Key::MIN..Key::MAX, and locating them using R-Tree search is very inefficient.
     /// So L0 layers are held in l0_delta_layers vector, in addition to the R-tree.
@@ -63,12 +65,180 @@ impl<L: ?Sized> Default for LayerMap<L> {
             open_layer: None,
             next_open_layer_at: None,
             frozen_layers: VecDeque::default(),
+            historic_layers: RTree::default(),
             l0_delta_layers: Vec::default(),
-            historic: BufferedHistoricLayerCoverage::default(),
         }
     }
 }
 
+struct LayerRTreeObject<L: ?Sized> {
+    layer: Arc<L>,
+
+    envelope: AABB<[IntKey; 2]>,
+}
+
+// Representation of Key as numeric type.
+// We can not use native implementation of i128, because rstar::RTree
+// doesn't handle properly integer overflow during area calculation: sum(Xi*Yi).
+// Overflow will cause panic in debug mode and incorrect area calculation in release mode,
+// which leads to non-optimally balanced R-Tree (but doesn't fit correctness of R-Tree work).
+// By using i256 as the type, even though all the actual values would fit in i128, we can be
+// sure that multiplication doesn't overflow.
+//
+
+#[derive(Clone, PartialEq, Eq, PartialOrd, Debug)]
+struct IntKey(i256);
+
+impl Copy for IntKey {}
+
+impl IntKey {
+    fn from(i: i128) -> Self {
+        IntKey(i256::from(i))
+    }
+}
+
+impl Bounded for IntKey {
+    fn min_value() -> Self {
+        IntKey(i256::MIN)
+    }
+    fn max_value() -> Self {
+        IntKey(i256::MAX)
+    }
+}
+
+impl Signed for IntKey {
+    fn is_positive(&self) -> bool {
+        self.0 > i256::ZERO
+    }
+    fn is_negative(&self) -> bool {
+        self.0 < i256::ZERO
+    }
+    fn signum(&self) -> Self {
+        match self.0.cmp(&i256::ZERO) {
+            Ordering::Greater => IntKey(i256::ONE),
+            Ordering::Less => IntKey(-i256::ONE),
+            Ordering::Equal => IntKey(i256::ZERO),
+        }
+    }
+    fn abs(&self) -> Self {
+        IntKey(self.0.abs())
+    }
+    fn abs_sub(&self, other: &Self) -> Self {
+        if self.0 <= other.0 {
+            IntKey(i256::ZERO)
+        } else {
+            IntKey(self.0 - other.0)
+        }
+    }
+}
+
+impl Neg for IntKey {
+    type Output = Self;
+    fn neg(self) -> Self::Output {
+        IntKey(-self.0)
+    }
+}
+
+impl Rem for IntKey {
+    type Output = Self;
+    fn rem(self, rhs: Self) -> Self::Output {
+        IntKey(self.0 % rhs.0)
+    }
+}
+
+impl Div for IntKey {
+    type Output = Self;
+    fn div(self, rhs: Self) -> Self::Output {
+        IntKey(self.0 / rhs.0)
+    }
+}
+
+impl Add for IntKey {
+    type Output = Self;
+    fn add(self, rhs: Self) -> Self::Output {
+        IntKey(self.0 + rhs.0)
+    }
+}
+
+impl Sub for IntKey {
+    type Output = Self;
+    fn sub(self, rhs: Self) -> Self::Output {
+        IntKey(self.0 - rhs.0)
+    }
+}
+
+impl Mul for IntKey {
+    type Output = Self;
+    fn mul(self, rhs: Self) -> Self::Output {
+        IntKey(self.0 * rhs.0)
+    }
+}
+
+impl One for IntKey {
+    fn one() -> Self {
+        IntKey(i256::ONE)
+    }
+}
+
+impl Zero for IntKey {
+    fn zero() -> Self {
+        IntKey(i256::ZERO)
+    }
+    fn is_zero(&self) -> bool {
+        self.0 == i256::ZERO
+    }
+}
+
+impl Num for IntKey {
+    type FromStrRadixErr = <i128 as Num>::FromStrRadixErr;
+    fn from_str_radix(str: &str, radix: u32) -> Result<Self, Self::FromStrRadixErr> {
+        Ok(IntKey(i256::from(i128::from_str_radix(str, radix)?)))
+    }
+}
+
+impl<T: ?Sized> PartialEq for LayerRTreeObject<T> {
+    fn eq(&self, other: &Self) -> bool {
+        // FIXME: ptr_eq might fail to return true for 'dyn'
+        // references.  Clippy complains about this. In practice it
+        // seems to work, the assertion below would be triggered
+        // otherwise but this ought to be fixed.
+        #[allow(clippy::vtable_address_comparisons)]
+        Arc::ptr_eq(&self.layer, &other.layer)
+    }
+}
+
+impl<L> RTreeObject for LayerRTreeObject<L>
+where
+    L: ?Sized,
+{
+    type Envelope = AABB<[IntKey; 2]>;
+    fn envelope(&self) -> Self::Envelope {
+        self.envelope
+    }
+}
+
+impl<L> LayerRTreeObject<L>
+where
+    L: ?Sized + Layer,
+{
+    fn new(layer: Arc<L>) -> Self {
+        let key_range = layer.get_key_range();
+        let lsn_range = layer.get_lsn_range();
+
+        let envelope = AABB::from_corners(
+            [
+                IntKey::from(key_range.start.to_i128()),
+                IntKey::from(lsn_range.start.0 as i128),
+            ],
+            [
+                IntKey::from(key_range.end.to_i128() - 1),
+                IntKey::from(lsn_range.end.0 as i128 - 1),
+            ], // AABB::upper is inclusive, while `key_range.end` and `lsn_range.end` are exclusive
+        );
+        LayerRTreeObject { layer, envelope }
+    }
+}
+
 /// Return value of LayerMap::search
 pub struct SearchResult<L: ?Sized> {
     pub layer: Arc<L>,
@@ -80,117 +250,138 @@ where
     L: ?Sized + Layer,
 {
     ///
-    /// Find the latest layer (by lsn.end) that covers the given
-    /// 'key', with lsn.start < 'end_lsn'.
+    /// Find the latest layer that covers the given 'key', with lsn <
+    /// 'end_lsn'.
     ///
-    /// The caller of this function is the page reconstruction
-    /// algorithm looking for the next relevant delta layer, or
-    /// the terminal image layer. The caller will pass the lsn_floor
-    /// value as end_lsn in the next call to search.
-    ///
-    /// If there's an image layer exactly below the given end_lsn,
-    /// search should return that layer regardless if there are
-    /// overlapping deltas.
-    ///
-    /// If the latest layer is a delta and there is an overlapping
-    /// image with it below, the lsn_floor returned should be right
-    /// above that image so we don't skip it in the search. Otherwise
-    /// the lsn_floor returned should be the bottom of the delta layer
-    /// because we should make as much progress down the lsn axis
-    /// as possible. It's fine if this way we skip some overlapping
-    /// deltas, because the delta we returned would contain the same
-    /// wal content.
-    ///
-    /// TODO: This API is convoluted and inefficient. If the caller
-    /// makes N search calls, we'll end up finding the same latest
-    /// image layer N times. We should either cache the latest image
-    /// layer result, or simplify the api to `get_latest_image` and
-    /// `get_latest_delta`, and only call `get_latest_image` once.
+    /// Returns the layer, if any, and an 'lsn_floor' value that
+    /// indicates which portion of the layer the caller should
+    /// check. 'lsn_floor' is normally the start-LSN of the layer, but
+    /// can be greater if there is an overlapping layer that might
+    /// contain the version, even if it's missing from the returned
+    /// layer.
     ///
     /// NOTE: This only searches the 'historic' layers, *not* the
     /// 'open' and 'frozen' layers!
     ///
     pub fn search(&self, key: Key, end_lsn: Lsn) -> Option<SearchResult<L>> {
-        let version = self.historic.get().unwrap().get_version(end_lsn.0 - 1)?;
-        let latest_delta = version.delta_coverage.query(key.to_i128());
-        let latest_image = version.image_coverage.query(key.to_i128());
+        // Find the latest image layer that covers the given key
+        let mut latest_img: Option<Arc<L>> = None;
+        let mut latest_img_lsn: Option<Lsn> = None;
+        let envelope = AABB::from_corners(
+            [IntKey::from(key.to_i128()), IntKey::from(0i128)],
+            [
+                IntKey::from(key.to_i128()),
+                IntKey::from(end_lsn.0 as i128 - 1),
+            ],
+        );
+        for e in self
+            .historic_layers
+            .locate_in_envelope_intersecting(&envelope)
+        {
+            let l = &e.layer;
+            if l.is_incremental() {
+                continue;
+            }
+            assert!(l.get_key_range().contains(&key));
+            let img_lsn = l.get_lsn_range().start;
+            assert!(img_lsn < end_lsn);
+            if Lsn(img_lsn.0 + 1) == end_lsn {
+                // found exact match
+                return Some(SearchResult {
+                    layer: Arc::clone(l),
+                    lsn_floor: img_lsn,
+                });
+            }
+            if img_lsn > latest_img_lsn.unwrap_or(Lsn(0)) {
+                latest_img = Some(Arc::clone(l));
+                latest_img_lsn = Some(img_lsn);
+            }
+        }
 
-        match (latest_delta, latest_image) {
-            (None, None) => None,
-            (None, Some(image)) => {
-                let lsn_floor = image.get_lsn_range().start;
-                Some(SearchResult {
-                    layer: image,
-                    lsn_floor,
-                })
+        // Search the delta layers
+        let mut latest_delta: Option<Arc<L>> = None;
+        for e in self
+            .historic_layers
+            .locate_in_envelope_intersecting(&envelope)
+        {
+            let l = &e.layer;
+            if !l.is_incremental() {
+                continue;
             }
-            (Some(delta), None) => {
-                let lsn_floor = delta.get_lsn_range().start;
-                Some(SearchResult {
-                    layer: delta,
-                    lsn_floor,
-                })
+            assert!(l.get_key_range().contains(&key));
+            if l.get_lsn_range().start >= end_lsn {
+                info!(
+                    "Candidate delta layer {}..{} is too new for lsn {}",
+                    l.get_lsn_range().start,
+                    l.get_lsn_range().end,
+                    end_lsn
+                );
             }
-            (Some(delta), Some(image)) => {
-                let img_lsn = image.get_lsn_range().start;
-                let image_is_newer = image.get_lsn_range().end > delta.get_lsn_range().end;
-                let image_exact_match = Lsn(img_lsn.0 + 1) == end_lsn;
-                if image_is_newer || image_exact_match {
-                    Some(SearchResult {
-                        layer: image,
-                        lsn_floor: img_lsn,
-                    })
+            assert!(l.get_lsn_range().start < end_lsn);
+            if l.get_lsn_range().end >= end_lsn {
+                // this layer contains the requested point in the key/lsn space.
+                // No need to search any further
+                trace!(
+                    "found layer {} for request on {key} at {end_lsn}",
+                    l.short_id(),
+                );
+                latest_delta.replace(Arc::clone(l));
+                break;
+            }
+            if l.get_lsn_range().end > latest_img_lsn.unwrap_or(Lsn(0)) {
+                // this layer's end LSN is smaller than the requested point. If there's
+                // nothing newer, this is what we need to return. Remember this.
+                if let Some(old_candidate) = &latest_delta {
+                    if l.get_lsn_range().end > old_candidate.get_lsn_range().end {
+                        latest_delta.replace(Arc::clone(l));
+                    }
                 } else {
-                    let lsn_floor =
-                        std::cmp::max(delta.get_lsn_range().start, image.get_lsn_range().start + 1);
-                    Some(SearchResult {
-                        layer: delta,
-                        lsn_floor,
-                    })
+                    latest_delta.replace(Arc::clone(l));
                 }
             }
         }
+        if let Some(l) = latest_delta {
+            trace!(
+                "found (old) layer {} for request on {key} at {end_lsn}",
+                l.short_id(),
+            );
+            let lsn_floor = std::cmp::max(
+                Lsn(latest_img_lsn.unwrap_or(Lsn(0)).0 + 1),
+                l.get_lsn_range().start,
+            );
+            Some(SearchResult {
+                lsn_floor,
+                layer: l,
+            })
+        } else if let Some(l) = latest_img {
+            trace!("found img layer and no deltas for request on {key} at {end_lsn}");
+            Some(SearchResult {
+                lsn_floor: latest_img_lsn.unwrap(),
+                layer: l,
+            })
+        } else {
+            trace!("no layer found for request on {key} at {end_lsn}");
+            None
+        }
     }
 
     ///
     /// Insert an on-disk layer
     ///
     pub fn insert_historic(&mut self, layer: Arc<L>) {
-        let kr = layer.get_key_range();
-        let lr = layer.get_lsn_range();
-        self.historic.insert(
-            kr.start.to_i128()..kr.end.to_i128(),
-            lr.start.0..lr.end.0,
-            Arc::clone(&layer),
-            !layer.is_incremental(),
-        );
-
         if layer.get_key_range() == (Key::MIN..Key::MAX) {
-            self.l0_delta_layers.push(layer);
+            self.l0_delta_layers.push(layer.clone());
         }
-
+        self.historic_layers.insert(LayerRTreeObject::new(layer));
         NUM_ONDISK_LAYERS.inc();
     }
 
-    /// Must be called after a batch of insert_historic calls, before querying
-    pub fn rebuild_index(&mut self) {
-        self.historic.rebuild();
-    }
-
     ///
     /// Remove an on-disk layer from the map.
     ///
     /// This should be called when the corresponding file on disk has been deleted.
     ///
     pub fn remove_historic(&mut self, layer: Arc<L>) {
-        let kr = layer.get_key_range();
-        let lr = layer.get_lsn_range();
-        self.historic.remove(
-            kr.start.to_i128()..kr.end.to_i128(),
-            lr.start.0..lr.end.0,
-            !layer.is_incremental(),
-        );
-
         if layer.get_key_range() == (Key::MIN..Key::MAX) {
             let len_before = self.l0_delta_layers.len();
 
@@ -203,51 +394,96 @@ where
                 .retain(|other| !Arc::ptr_eq(other, &layer));
             assert_eq!(self.l0_delta_layers.len(), len_before - 1);
         }
-
+        assert!(self
+            .historic_layers
+            .remove(&LayerRTreeObject::new(layer))
+            .is_some());
         NUM_ONDISK_LAYERS.dec();
     }
 
-    /// Is there a newer image layer for given key- and LSN-range? Or a set
-    /// of image layers within the specified lsn range that cover the entire
-    /// specified key range?
+    /// Is there a newer image layer for given key- and LSN-range?
     ///
     /// This is used for garbage collection, to determine if an old layer can
     /// be deleted.
-    pub fn image_layer_exists(&self, key: &Range<Key>, lsn: &Range<Lsn>) -> Result<bool> {
-        if key.is_empty() {
-            return Ok(true);
-        }
+    pub fn image_layer_exists(
+        &self,
+        key_range: &Range<Key>,
+        lsn_range: &Range<Lsn>,
+    ) -> Result<bool> {
+        let mut range_remain = key_range.clone();
 
-        let version = match self.historic.get().unwrap().get_version(lsn.end.0) {
-            Some(v) => v,
-            None => return Ok(false),
-        };
+        loop {
+            let mut made_progress = false;
+            let envelope = AABB::from_corners(
+                [
+                    IntKey::from(range_remain.start.to_i128()),
+                    IntKey::from(lsn_range.start.0 as i128),
+                ],
+                [
+                    IntKey::from(range_remain.end.to_i128() - 1),
+                    IntKey::from(lsn_range.end.0 as i128 - 1),
+                ],
+            );
+            for e in self
+                .historic_layers
+                .locate_in_envelope_intersecting(&envelope)
+            {
+                let l = &e.layer;
+                if l.is_incremental() {
+                    continue;
+                }
+                let img_lsn = l.get_lsn_range().start;
+                if l.get_key_range().contains(&range_remain.start) && lsn_range.contains(&img_lsn) {
+                    made_progress = true;
+                    let img_key_end = l.get_key_range().end;
 
-        let start = key.start.to_i128();
-        let end = key.end.to_i128();
+                    if img_key_end >= range_remain.end {
+                        return Ok(true);
+                    }
+                    range_remain.start = img_key_end;
+                }
+            }
 
-        let layer_covers = |layer: Option<Arc<L>>| match layer {
-            Some(layer) => layer.get_lsn_range().start >= lsn.start,
-            None => false,
-        };
-
-        // Check the start is covered
-        if !layer_covers(version.image_coverage.query(start)) {
-            return Ok(false);
-        }
-
-        // Check after all changes of coverage
-        for (_, change_val) in version.image_coverage.range(start..end) {
-            if !layer_covers(change_val) {
+            if !made_progress {
                 return Ok(false);
             }
         }
-
-        Ok(true)
     }
 
     pub fn iter_historic_layers(&self) -> impl '_ + Iterator<Item = Arc<L>> {
-        self.historic.iter()
+        self.historic_layers.iter().map(|e| e.layer.clone())
+    }
+
+    /// Find the last image layer that covers 'key', ignoring any image layers
+    /// newer than 'lsn'.
+    fn find_latest_image(&self, key: Key, lsn: Lsn) -> Option<Arc<L>> {
+        let mut candidate_lsn = Lsn(0);
+        let mut candidate = None;
+        let envelope = AABB::from_corners(
+            [IntKey::from(key.to_i128()), IntKey::from(0)],
+            [IntKey::from(key.to_i128()), IntKey::from(lsn.0 as i128)],
+        );
+        for e in self
+            .historic_layers
+            .locate_in_envelope_intersecting(&envelope)
+        {
+            let l = &e.layer;
+            if l.is_incremental() {
+                continue;
+            }
+
+            assert!(l.get_key_range().contains(&key));
+            let this_lsn = l.get_lsn_range().start;
+            assert!(this_lsn <= lsn);
+            if this_lsn < candidate_lsn {
+                // our previous candidate was better
+                continue;
+            }
+            candidate_lsn = this_lsn;
+            candidate = Some(Arc::clone(l));
+        }
+
+        candidate
     }
 
     ///
@@ -263,239 +499,87 @@ where
         key_range: &Range<Key>,
         lsn: Lsn,
     ) -> Result<Vec<(Range<Key>, Option<Arc<L>>)>> {
-        // TODO I'm 80% sure the lsn bound is inclusive. Double-check that
-        //      and document it. Do the same for image_layer_exists and count_deltas
-        let version = match self.historic.get().unwrap().get_version(lsn.0) {
-            Some(v) => v,
-            None => return Ok(vec![]),
-        };
-
-        let start = key_range.start.to_i128();
-        let end = key_range.end.to_i128();
-
-        // Initialize loop variables
-        let mut coverage: Vec<(Range<Key>, Option<Arc<L>>)> = vec![];
-        let mut current_key = start;
-        let mut current_val = version.image_coverage.query(start);
-
-        // Loop through the change events and push intervals
-        for (change_key, change_val) in version.image_coverage.range(start..end) {
-            let kr = Key::from_i128(current_key)..Key::from_i128(change_key);
-            coverage.push((kr, current_val.take()));
-            current_key = change_key;
-            current_val = change_val.clone();
+        let mut points = vec![key_range.start];
+        let envelope = AABB::from_corners(
+            [IntKey::from(key_range.start.to_i128()), IntKey::from(0)],
+            [
+                IntKey::from(key_range.end.to_i128()),
+                IntKey::from(lsn.0 as i128),
+            ],
+        );
+        for e in self
+            .historic_layers
+            .locate_in_envelope_intersecting(&envelope)
+        {
+            let l = &e.layer;
+            assert!(l.get_lsn_range().start <= lsn);
+            let range = l.get_key_range();
+            if key_range.contains(&range.start) {
+                points.push(l.get_key_range().start);
+            }
+            if key_range.contains(&range.end) {
+                points.push(l.get_key_range().end);
+            }
         }
+        points.push(key_range.end);
 
-        // Add the final interval
-        let kr = Key::from_i128(current_key)..Key::from_i128(end);
-        coverage.push((kr, current_val.take()));
+        points.sort();
+        points.dedup();
 
-        Ok(coverage)
+        // Ok, we now have a list of "interesting" points in the key space
+
+        // For each range between the points, find the latest image
+        let mut start = *points.first().unwrap();
+        let mut ranges = Vec::new();
+        for end in points[1..].iter() {
+            let img = self.find_latest_image(start, lsn);
+
+            ranges.push((start..*end, img));
+
+            start = *end;
+        }
+        Ok(ranges)
     }
 
-    /// Count the height of the tallest stack of deltas in this 2d region.
-    ///
-    /// If `limit` is provided we don't try to count above that number.
-    ///
-    /// This number is used to compute the largest number of deltas that
-    /// we'll need to visit for any page reconstruction in this region.
-    /// We use this heuristic to decide whether to create an image layer.
-    pub fn count_deltas(
-        &self,
-        key: &Range<Key>,
-        lsn: &Range<Lsn>,
-        limit: Option<usize>,
-    ) -> Result<usize> {
-        // We get the delta coverage of the region, and for each part of the coverage
-        // we recurse right underneath the delta. The recursion depth is limited by
-        // the largest result this function could return, which is in practice between
-        // 3 and 10 (since we usually try to create an image when the number gets larger).
-
-        if lsn.is_empty() || key.is_empty() || limit == Some(0) {
+    /// Count how many L1 delta layers there are that overlap with the
+    /// given key and LSN range.
+    pub fn count_deltas(&self, key_range: &Range<Key>, lsn_range: &Range<Lsn>) -> Result<usize> {
+        let mut result = 0;
+        if lsn_range.start >= lsn_range.end {
             return Ok(0);
         }
+        let envelope = AABB::from_corners(
+            [
+                IntKey::from(key_range.start.to_i128()),
+                IntKey::from(lsn_range.start.0 as i128),
+            ],
+            [
+                IntKey::from(key_range.end.to_i128() - 1),
+                IntKey::from(lsn_range.end.0 as i128 - 1),
+            ],
+        );
+        for e in self
+            .historic_layers
+            .locate_in_envelope_intersecting(&envelope)
+        {
+            let l = &e.layer;
+            if !l.is_incremental() {
+                continue;
+            }
+            assert!(range_overlaps(&l.get_lsn_range(), lsn_range));
+            assert!(range_overlaps(&l.get_key_range(), key_range));
 
-        let version = match self.historic.get().unwrap().get_version(lsn.end.0 - 1) {
-            Some(v) => v,
-            None => return Ok(0),
-        };
-
-        let start = key.start.to_i128();
-        let end = key.end.to_i128();
-
-        // Initialize loop variables
-        let mut max_stacked_deltas = 0;
-        let mut current_key = start;
-        let mut current_val = version.delta_coverage.query(start);
-
-        // Loop through the delta coverage and recurse on each part
-        for (change_key, change_val) in version.delta_coverage.range(start..end) {
-            // If there's a relevant delta in this part, add 1 and recurse down
-            if let Some(val) = current_val {
-                if val.get_lsn_range().end > lsn.start {
-                    let kr = Key::from_i128(current_key)..Key::from_i128(change_key);
-                    let lr = lsn.start..val.get_lsn_range().start;
-                    if !kr.is_empty() {
-                        let new_limit = limit.map(|l| l - 1);
-                        let max_stacked_deltas_underneath =
-                            self.count_deltas(&kr, &lr, new_limit)?;
-                        max_stacked_deltas =
-                            std::cmp::max(max_stacked_deltas, 1 + max_stacked_deltas_underneath);
-                    }
-                }
+            // We ignore level0 delta layers. Unless the whole keyspace fits
+            // into one partition
+            if !range_eq(key_range, &(Key::MIN..Key::MAX))
+                && range_eq(&l.get_key_range(), &(Key::MIN..Key::MAX))
+            {
+                continue;
             }
 
-            current_key = change_key;
-            current_val = change_val.clone();
+            result += 1;
         }
-
-        // Consider the last part
-        if let Some(val) = current_val {
-            if val.get_lsn_range().end > lsn.start {
-                let kr = Key::from_i128(current_key)..Key::from_i128(end);
-                let lr = lsn.start..val.get_lsn_range().start;
-
-                if !kr.is_empty() {
-                    let new_limit = limit.map(|l| l - 1);
-                    let max_stacked_deltas_underneath = self.count_deltas(&kr, &lr, new_limit)?;
-                    max_stacked_deltas =
-                        std::cmp::max(max_stacked_deltas, 1 + max_stacked_deltas_underneath);
-                }
-            }
-        }
-
-        Ok(max_stacked_deltas)
-    }
-
-    /// Count how many layers we need to visit for given key-lsn pair.
-    ///
-    /// Used as a helper for correctness checks only. Performance not critical.
-    pub fn get_difficulty(&self, lsn: Lsn, key: Key) -> usize {
-        match self.search(key, lsn) {
-            Some(search_result) => {
-                if search_result.layer.is_incremental() {
-                    1 + self.get_difficulty(search_result.lsn_floor, key)
-                } else {
-                    0
-                }
-            }
-            None => 0,
-        }
-    }
-
-    /// Used for correctness checking. Results are expected to be identical to
-    /// self.get_difficulty_map. Assumes self.search is correct.
-    pub fn get_difficulty_map_bruteforce(
-        &self,
-        lsn: Lsn,
-        partitioning: &KeyPartitioning,
-    ) -> Vec<usize> {
-        // Looking at the difficulty as a function of key, it could only increase
-        // when a delta layer starts or an image layer ends. Therefore it's sufficient
-        // to check the difficulties at:
-        // - the key.start for each non-empty part range
-        // - the key.start for each delta
-        // - the key.end for each image
-        let keys_iter: Box<dyn Iterator<Item = Key>> = {
-            let mut keys: Vec<Key> = self
-                .iter_historic_layers()
-                .map(|layer| {
-                    if layer.is_incremental() {
-                        layer.get_key_range().start
-                    } else {
-                        layer.get_key_range().end
-                    }
-                })
-                .collect();
-            keys.sort();
-            Box::new(keys.into_iter())
-        };
-        let mut keys_iter = keys_iter.peekable();
-
-        // Iter the partition and keys together and query all the necessary
-        // keys, computing the max difficulty for each part.
-        partitioning
-            .parts
-            .iter()
-            .map(|part| {
-                let mut difficulty = 0;
-                // Partition ranges are assumed to be sorted and disjoint
-                // TODO assert it
-                for range in &part.ranges {
-                    if !range.is_empty() {
-                        difficulty =
-                            std::cmp::max(difficulty, self.get_difficulty(lsn, range.start));
-                    }
-                    while let Some(key) = keys_iter.peek() {
-                        if key >= &range.end {
-                            break;
-                        }
-                        let key = keys_iter.next().unwrap();
-                        if key < range.start {
-                            continue;
-                        }
-                        difficulty = std::cmp::max(difficulty, self.get_difficulty(lsn, key));
-                    }
-                }
-                difficulty
-            })
-            .collect()
-    }
-
-    /// For each part of a keyspace partitioning, return the maximum number of layers
-    /// that would be needed for page reconstruction in that part at the given LSN.
-    ///
-    /// If `limit` is provided we don't try to count above that number.
-    ///
-    /// This method is used to decide where to create new image layers. Computing the
-    /// result for the entire partitioning at once allows this function to be more
-    /// efficient, and further optimization is possible by using iterators instead,
-    /// to allow early return.
-    ///
-    /// TODO actually use this method instead of count_deltas. Currently we only use
-    ///      it for benchmarks.
-    pub fn get_difficulty_map(
-        &self,
-        lsn: Lsn,
-        partitioning: &KeyPartitioning,
-        limit: Option<usize>,
-    ) -> Vec<usize> {
-        // TODO This is a naive implementation. Perf improvements to do:
-        // 1. Instead of calling self.image_coverage and self.count_deltas,
-        //    iterate the image and delta coverage only once.
-        partitioning
-            .parts
-            .iter()
-            .map(|part| {
-                let mut difficulty = 0;
-                for range in &part.ranges {
-                    if limit == Some(difficulty) {
-                        break;
-                    }
-                    for (img_range, last_img) in self
-                        .image_coverage(range, lsn)
-                        .expect("why would this err?")
-                    {
-                        if limit == Some(difficulty) {
-                            break;
-                        }
-                        let img_lsn = if let Some(last_img) = last_img {
-                            last_img.get_lsn_range().end
-                        } else {
-                            Lsn(0)
-                        };
-
-                        if img_lsn < lsn {
-                            let num_deltas = self
-                                .count_deltas(&img_range, &(img_lsn..lsn), limit)
-                                .expect("why would this err lol?");
-                            difficulty = std::cmp::max(difficulty, num_deltas);
-                        }
-                    }
-                }
-                difficulty
-            })
-            .collect()
+        Ok(result)
     }
 
     /// Return all L0 delta layers
@@ -519,8 +603,8 @@ where
         }
 
         println!("historic_layers:");
-        for layer in self.iter_historic_layers() {
-            layer.dump(verbose)?;
+        for e in self.historic_layers.iter() {
+            e.layer.dump(verbose)?;
         }
         println!("End dump LayerMap");
         Ok(())

pageserver/src/tenant/layer_map/historic_layer_coverage.rs

@@ -1,418 +0,0 @@
-use std::collections::BTreeMap;
-use std::ops::Range;
-
-use super::layer_coverage::LayerCoverageTuple;
-
-/// Efficiently queryable layer coverage for each LSN.
-///
-/// Allows answering layer map queries very efficiently,
-/// but doesn't allow retroactive insertion, which is
-/// sometimes necessary. See BufferedHistoricLayerCoverage.
-pub struct HistoricLayerCoverage<Value> {
-    /// The latest state
-    head: LayerCoverageTuple<Value>,
-
-    /// All previous states
-    historic: BTreeMap<u64, LayerCoverageTuple<Value>>,
-}
-
-impl<T: Clone> Default for HistoricLayerCoverage<T> {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl<Value: Clone> HistoricLayerCoverage<Value> {
-    pub fn new() -> Self {
-        Self {
-            head: LayerCoverageTuple::default(),
-            historic: BTreeMap::default(),
-        }
-    }
-
-    /// Add a layer
-    ///
-    /// Panics if new layer has older lsn.start than an existing layer.
-    /// See BufferedHistoricLayerCoverage for a more general insertion method.
-    pub fn insert(&mut self, key: Range<i128>, lsn: Range<u64>, value: Value, is_image: bool) {
-        // It's only a persistent map, not a retroactive one
-        if let Some(last_entry) = self.historic.iter().rev().next() {
-            let last_lsn = last_entry.0;
-            if lsn.start == *last_lsn {
-                // TODO there are edge cases to take care of
-            }
-            if lsn.start < *last_lsn {
-                panic!("unexpected retroactive insert");
-            }
-        }
-
-        // Insert into data structure
-        if is_image {
-            self.head.image_coverage.insert(key, lsn.clone(), value);
-        } else {
-            self.head.delta_coverage.insert(key, lsn.clone(), value);
-        }
-
-        // Remember history. Clone is O(1)
-        self.historic.insert(lsn.start, self.head.clone());
-    }
-
-    /// Query at a particular LSN, inclusive
-    pub fn get_version(&self, lsn: u64) -> Option<&LayerCoverageTuple<Value>> {
-        match self.historic.range(..=lsn).rev().next() {
-            Some((_, v)) => Some(v),
-            None => None,
-        }
-    }
-
-    /// Remove all entries after a certain LSN
-    pub fn trim(&mut self, begin: &u64) {
-        self.historic.split_off(begin);
-        self.head = self
-            .historic
-            .iter()
-            .rev()
-            .next()
-            .map(|(_, v)| v.clone())
-            .unwrap_or_default();
-    }
-}
-
-/// This is the most basic test that demonstrates intended usage.
-/// All layers in this test have height 1.
-#[test]
-fn test_persistent_simple() {
-    let mut map = HistoricLayerCoverage::<String>::new();
-    map.insert(0..5, 100..101, "Layer 1".to_string(), true);
-    map.insert(3..9, 110..111, "Layer 2".to_string(), true);
-    map.insert(5..6, 120..121, "Layer 3".to_string(), true);
-
-    // After Layer 1 insertion
-    let version = map.get_version(105).unwrap();
-    assert_eq!(version.image_coverage.query(1), Some("Layer 1".to_string()));
-    assert_eq!(version.image_coverage.query(4), Some("Layer 1".to_string()));
-
-    // After Layer 2 insertion
-    let version = map.get_version(115).unwrap();
-    assert_eq!(version.image_coverage.query(4), Some("Layer 2".to_string()));
-    assert_eq!(version.image_coverage.query(8), Some("Layer 2".to_string()));
-    assert_eq!(version.image_coverage.query(11), None);
-
-    // After Layer 3 insertion
-    let version = map.get_version(125).unwrap();
-    assert_eq!(version.image_coverage.query(4), Some("Layer 2".to_string()));
-    assert_eq!(version.image_coverage.query(5), Some("Layer 3".to_string()));
-    assert_eq!(version.image_coverage.query(7), Some("Layer 2".to_string()));
-}
-
-/// Cover simple off-by-one edge cases
-#[test]
-fn test_off_by_one() {
-    let mut map = HistoricLayerCoverage::<String>::new();
-    map.insert(3..5, 100..110, "Layer 1".to_string(), true);
-
-    // Check different LSNs
-    let version = map.get_version(99);
-    assert!(version.is_none());
-    let version = map.get_version(100).unwrap();
-    assert_eq!(version.image_coverage.query(4), Some("Layer 1".to_string()));
-    let version = map.get_version(110).unwrap();
-    assert_eq!(version.image_coverage.query(4), Some("Layer 1".to_string()));
-
-    // Check different keys
-    let version = map.get_version(105).unwrap();
-    assert_eq!(version.image_coverage.query(2), None);
-    assert_eq!(version.image_coverage.query(3), Some("Layer 1".to_string()));
-    assert_eq!(version.image_coverage.query(4), Some("Layer 1".to_string()));
-    assert_eq!(version.image_coverage.query(5), None);
-}
-
-/// Cover edge cases where layers begin or end on the same key
-#[test]
-fn test_key_collision() {
-    let mut map = HistoricLayerCoverage::<String>::new();
-
-    map.insert(3..5, 100..110, "Layer 10".to_string(), true);
-    map.insert(5..8, 100..110, "Layer 11".to_string(), true);
-
-    map.insert(3..4, 200..210, "Layer 20".to_string(), true);
-
-    // Check after layer 11
-    let version = map.get_version(105).unwrap();
-    assert_eq!(version.image_coverage.query(2), None);
-    assert_eq!(
-        version.image_coverage.query(3),
-        Some("Layer 10".to_string())
-    );
-    assert_eq!(
-        version.image_coverage.query(5),
-        Some("Layer 11".to_string())
-    );
-    assert_eq!(
-        version.image_coverage.query(7),
-        Some("Layer 11".to_string())
-    );
-    assert_eq!(version.image_coverage.query(8), None);
-
-    // Check after layer 20
-    let version = map.get_version(205).unwrap();
-    assert_eq!(version.image_coverage.query(2), None);
-    assert_eq!(
-        version.image_coverage.query(3),
-        Some("Layer 20".to_string())
-    );
-    assert_eq!(
-        version.image_coverage.query(5),
-        Some("Layer 11".to_string())
-    );
-    assert_eq!(
-        version.image_coverage.query(7),
-        Some("Layer 11".to_string())
-    );
-    assert_eq!(version.image_coverage.query(8), None);
-}
-
-/// Test when rectangles have nontrivial height and possibly overlap
-#[test]
-fn test_persistent_overlapping() {
-    let mut map = HistoricLayerCoverage::<String>::new();
-
-    // Add 3 key-disjoint layers with varying LSN ranges
-    map.insert(1..2, 100..200, "Layer 1".to_string(), true);
-    map.insert(4..5, 110..200, "Layer 2".to_string(), true);
-    map.insert(7..8, 120..300, "Layer 3".to_string(), true);
-
-    // Add wide and short layer
-    map.insert(0..9, 130..199, "Layer 4".to_string(), true);
-
-    // Add wide layer taller than some
-    map.insert(0..9, 140..201, "Layer 5".to_string(), true);
-
-    // Add wide layer taller than all
-    map.insert(0..9, 150..301, "Layer 6".to_string(), true);
-
-    // After layer 4 insertion
-    let version = map.get_version(135).unwrap();
-    assert_eq!(version.image_coverage.query(0), Some("Layer 4".to_string()));
-    assert_eq!(version.image_coverage.query(1), Some("Layer 1".to_string()));
-    assert_eq!(version.image_coverage.query(2), Some("Layer 4".to_string()));
-    assert_eq!(version.image_coverage.query(4), Some("Layer 2".to_string()));
-    assert_eq!(version.image_coverage.query(5), Some("Layer 4".to_string()));
-    assert_eq!(version.image_coverage.query(7), Some("Layer 3".to_string()));
-    assert_eq!(version.image_coverage.query(8), Some("Layer 4".to_string()));
-
-    // After layer 5 insertion
-    let version = map.get_version(145).unwrap();
-    assert_eq!(version.image_coverage.query(0), Some("Layer 5".to_string()));
-    assert_eq!(version.image_coverage.query(1), Some("Layer 5".to_string()));
-    assert_eq!(version.image_coverage.query(2), Some("Layer 5".to_string()));
-    assert_eq!(version.image_coverage.query(4), Some("Layer 5".to_string()));
-    assert_eq!(version.image_coverage.query(5), Some("Layer 5".to_string()));
-    assert_eq!(version.image_coverage.query(7), Some("Layer 3".to_string()));
-    assert_eq!(version.image_coverage.query(8), Some("Layer 5".to_string()));
-
-    // After layer 6 insertion
-    let version = map.get_version(155).unwrap();
-    assert_eq!(version.image_coverage.query(0), Some("Layer 6".to_string()));
-    assert_eq!(version.image_coverage.query(1), Some("Layer 6".to_string()));
-    assert_eq!(version.image_coverage.query(2), Some("Layer 6".to_string()));
-    assert_eq!(version.image_coverage.query(4), Some("Layer 6".to_string()));
-    assert_eq!(version.image_coverage.query(5), Some("Layer 6".to_string()));
-    assert_eq!(version.image_coverage.query(7), Some("Layer 6".to_string()));
-    assert_eq!(version.image_coverage.query(8), Some("Layer 6".to_string()));
-}
-
-/// Wrapper for HistoricLayerCoverage that allows us to hack around the lack
-/// of support for retroactive insertion by rebuilding the map since the
-/// change.
-///
-/// Why is this needed? We most often insert new layers with newer LSNs,
-/// but during compaction we create layers with non-latest LSN, and during
-/// GC we delete historic layers.
-///
-/// Even though rebuilding is an expensive (N log N) solution to the problem,
-/// it's not critical since we do something equally expensive just to decide
-/// whether or not to create new image layers.
-///
-/// If this becomes an actual bottleneck, one solution would be to build a
-/// segment tree that holds PersistentLayerMaps. Though this would mean that
-/// we take an additional log(N) performance hit for queries, which will probably
-/// still be more critical.
-///
-/// See this for more on persistent and retroactive techniques:
-/// https://www.youtube.com/watch?v=WqCWghETNDc&t=581s
-pub struct BufferedHistoricLayerCoverage<Value> {
-    /// A persistent layer map that we rebuild when we need to retroactively update
-    historic_coverage: HistoricLayerCoverage<Value>,
-
-    /// We buffer insertion into the PersistentLayerMap to decrease the number of rebuilds.
-    ///
-    /// We implicitly assume that layers are identified by their key and lsn range.
-    /// A value of None means we want to delete this item.
-    buffer: BTreeMap<(u64, u64, i128, i128, bool), Option<Value>>,
-
-    /// All current layers. This is not used for search. Only to make rebuilds easier.
-    ///
-    /// We implicitly assume that layers are identified by their key and lsn range.
-    layers: BTreeMap<(u64, u64, i128, i128, bool), Value>,
-}
-
-impl<T: std::fmt::Debug> std::fmt::Debug for BufferedHistoricLayerCoverage<T> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("RetroactiveLayerMap")
-            .field("buffer", &self.buffer)
-            .field("layers", &self.layers)
-            .finish()
-    }
-}
-
-impl<T: Clone> Default for BufferedHistoricLayerCoverage<T> {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl<Value: Clone> BufferedHistoricLayerCoverage<Value> {
-    pub fn new() -> Self {
-        Self {
-            historic_coverage: HistoricLayerCoverage::<Value>::new(),
-            buffer: BTreeMap::new(),
-            layers: BTreeMap::new(),
-        }
-    }
-
-    pub fn insert(&mut self, key: Range<i128>, lsn: Range<u64>, value: Value, is_image: bool) {
-        self.buffer.insert(
-            (lsn.start, lsn.end, key.start, key.end, is_image),
-            Some(value),
-        );
-    }
-
-    pub fn remove(&mut self, key: Range<i128>, lsn: Range<u64>, is_image: bool) {
-        self.buffer
-            .insert((lsn.start, lsn.end, key.start, key.end, is_image), None);
-    }
-
-    pub fn rebuild(&mut self) {
-        // Find the first LSN that needs to be rebuilt
-        let rebuild_since: u64 = match self.buffer.iter().next() {
-            Some(((lsn_start, _, _, _, _), _)) => *lsn_start,
-            None => return, // No need to rebuild if buffer is empty
-        };
-
-        // Apply buffered updates to self.layers
-        self.buffer.retain(|rect, layer| {
-            match layer {
-                Some(l) => {
-                    let existing = self.layers.insert(*rect, l.clone());
-                    if existing.is_some() {
-                        // TODO this happened once. Investigate.
-                        // panic!("can't overwrite layer");
-                    }
-                }
-                None => {
-                    let existing = self.layers.remove(rect);
-                    if existing.is_none() {
-                        panic!("invalid layer deletion");
-                    }
-                }
-            };
-            false
-        });
-
-        // Rebuild
-        self.historic_coverage.trim(&rebuild_since);
-        for ((lsn_start, lsn_end, key_start, key_end, is_image), layer) in
-            self.layers.range((rebuild_since, 0, 0, 0, false)..)
-        {
-            self.historic_coverage.insert(
-                *key_start..*key_end,
-                *lsn_start..*lsn_end,
-                layer.clone(),
-                *is_image,
-            );
-        }
-    }
-
-    /// Iterate all the layers
-    pub fn iter(&self) -> impl '_ + Iterator<Item = Value> {
-        // NOTE we can actually perform this without rebuilding,
-        //      but it's not necessary for now.
-        if !self.buffer.is_empty() {
-            panic!("rebuild pls")
-        }
-
-        self.layers.values().cloned()
-    }
-
-    /// Return a reference to a queryable map, assuming all updates
-    /// have already been processed using self.rebuild()
-    pub fn get(&self) -> anyhow::Result<&HistoricLayerCoverage<Value>> {
-        // NOTE we error here instead of implicitly rebuilding because
-        //      rebuilding is somewhat expensive.
-        // TODO maybe implicitly rebuild and log/sentry an error?
-        if !self.buffer.is_empty() {
-            anyhow::bail!("rebuild required")
-        }
-
-        Ok(&self.historic_coverage)
-    }
-}
-
-#[test]
-fn test_retroactive_regression_1() {
-    let mut map = BufferedHistoricLayerCoverage::new();
-
-    map.insert(
-        0..21267647932558653966460912964485513215,
-        23761336..23761457,
-        "sdfsdfs".to_string(),
-        false,
-    );
-
-    map.rebuild();
-
-    let version = map.get().unwrap().get_version(23761457).unwrap();
-    assert_eq!(
-        version.delta_coverage.query(100),
-        Some("sdfsdfs".to_string())
-    );
-}
-
-#[test]
-fn test_retroactive_simple() {
-    let mut map = BufferedHistoricLayerCoverage::new();
-
-    // Append some images in increasing LSN order
-    map.insert(0..5, 100..101, "Image 1".to_string(), true);
-    map.insert(3..9, 110..111, "Image 2".to_string(), true);
-    map.insert(4..6, 120..121, "Image 3".to_string(), true);
-    map.insert(8..9, 120..121, "Image 4".to_string(), true);
-
-    // Add a delta layer out of order
-    map.insert(2..5, 105..106, "Delta 1".to_string(), true);
-
-    // Rebuild so we can start querying
-    map.rebuild();
-
-    // Query key 4
-    let version = map.get().unwrap().get_version(90);
-    assert!(version.is_none());
-    let version = map.get().unwrap().get_version(102).unwrap();
-    assert_eq!(version.image_coverage.query(4), Some("Image 1".to_string()));
-    let version = map.get().unwrap().get_version(107).unwrap();
-    assert_eq!(version.image_coverage.query(4), Some("Delta 1".to_string()));
-    let version = map.get().unwrap().get_version(115).unwrap();
-    assert_eq!(version.image_coverage.query(4), Some("Image 2".to_string()));
-    let version = map.get().unwrap().get_version(125).unwrap();
-    assert_eq!(version.image_coverage.query(4), Some("Image 3".to_string()));
-
-    // Remove Image 3
-    map.remove(4..6, 120..121, true);
-    map.rebuild();
-
-    // Check deletion worked
-    let version = map.get().unwrap().get_version(125).unwrap();
-    assert_eq!(version.image_coverage.query(4), Some("Image 2".to_string()));
-    assert_eq!(version.image_coverage.query(8), Some("Image 4".to_string()));
-}

pageserver/src/tenant/layer_map/layer_coverage.rs

@@ -1,152 +0,0 @@
-use std::ops::Range;
-
-// There's an alternative non-sync crate im-rc that's supposed to be 20%
-// faster. It's not necessary now, but if extra perf is needed that's an option.
-use im::OrdMap;
-
-/// Data structure that can efficiently:
-/// - find the latest layer by lsn.end at a given key
-/// - iterate the latest layers in a key range
-/// - insert layers in non-decreasing lsn.start order
-///
-/// The struct is parameterized over Value for easier
-/// testing, but in practice it's some sort of layer.
-pub struct LayerCoverage<Value> {
-    /// For every change in coverage (as we sweep the key space)
-    /// we store (lsn.end, value).
-    ///
-    /// We use an immutable/persistent tree so that we can keep historic
-    /// versions of this coverage without cloning the whole thing and
-    /// incurring quadratic memory cost. See HistoricLayerCoverage.
-    nodes: OrdMap<i128, Option<(u64, Value)>>,
-}
-
-impl<T: Clone> Default for LayerCoverage<T> {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl<Value: Clone> LayerCoverage<Value> {
-    pub fn new() -> Self {
-        Self {
-            nodes: OrdMap::default(),
-        }
-    }
-
-    /// Helper function to subdivide the key range without changing any values
-    ///
-    /// Complexity: O(log N)
-    fn add_node(&mut self, key: i128) {
-        let value = match self.nodes.range(..=key).last() {
-            Some((_, Some(v))) => Some(v.clone()),
-            Some((_, None)) => None,
-            None => None,
-        };
-        self.nodes.insert(key, value);
-    }
-
-    /// Insert a layer.
-    ///
-    /// Complexity: worst case O(N), in practice O(log N). See not in implementation.
-    pub fn insert(&mut self, key: Range<i128>, lsn: Range<u64>, value: Value) {
-        // NOTE The order of the following lines is important!!
-
-        // Add nodes at endpoints
-        self.add_node(key.start);
-        self.add_node(key.end);
-
-        // Raise the height where necessary
-        //
-        // NOTE This loop is worst case O(N), but amortized O(log N) in the special
-        // case when rectangles have no height. In practice I don't think we'll see
-        // the kind of layer intersections needed to trigger O(N) behavior. The worst
-        // case is N/2 horizontal layers overlapped with N/2 vertical layers in a
-        // grid pattern.
-        let mut to_update = Vec::new();
-        let mut to_remove = Vec::new();
-        let mut prev_covered = false;
-        for (k, node) in self.nodes.range(key.clone()) {
-            let needs_cover = match node {
-                None => true,
-                Some((h, _)) => h < &lsn.end,
-            };
-            if needs_cover {
-                match prev_covered {
-                    true => to_remove.push(*k),
-                    false => to_update.push(*k),
-                }
-            }
-            prev_covered = needs_cover;
-        }
-        if !prev_covered {
-            to_remove.push(key.end);
-        }
-        for k in to_update {
-            self.nodes.insert(k, Some((lsn.end, value.clone())));
-        }
-        for k in to_remove {
-            self.nodes.remove(&k);
-        }
-    }
-
-    /// Get the latest (by lsn.end) layer at a given key
-    ///
-    /// Complexity: O(log N)
-    pub fn query(&self, key: i128) -> Option<Value> {
-
-        // This is cursed. Commenting out these lines makes the code
-        // correct, even though they don't do anything.
-        let k1 = self.nodes.get_prev(&key)?.0;
-        // let k2 = self.nodes.range(..=key).rev().next()?.0;
-        // assert_eq!(k1, k2);
-
-        self.nodes
-            // .get_prev(&key)?
-            .range(..=key).rev().next()?
-            .1
-            .as_ref()
-            .map(|(_, v)| v.clone())
-    }
-
-    /// Iterate the changes in layer coverage in a given range. You will likely
-    /// want to start with self.query(key.start), and then follow up with self.range
-    ///
-    /// Complexity: O(log N + result_size)
-    pub fn range(&self, key: Range<i128>) -> impl '_ + Iterator<Item = (i128, Option<Value>)> {
-        self.nodes
-            .range(key)
-            .map(|(k, v)| (*k, v.as_ref().map(|x| x.1.clone())))
-    }
-
-    /// O(1) clone
-    pub fn clone(&self) -> Self {
-        Self {
-            nodes: self.nodes.clone(),
-        }
-    }
-}
-
-/// Image and delta coverage at a specific LSN.
-pub struct LayerCoverageTuple<Value> {
-    pub image_coverage: LayerCoverage<Value>,
-    pub delta_coverage: LayerCoverage<Value>,
-}
-
-impl<T: Clone> Default for LayerCoverageTuple<T> {
-    fn default() -> Self {
-        Self {
-            image_coverage: LayerCoverage::default(),
-            delta_coverage: LayerCoverage::default(),
-        }
-    }
-}
-
-impl<Value: Clone> LayerCoverageTuple<Value> {
-    pub fn clone(&self) -> Self {
-        Self {
-            image_coverage: self.image_coverage.clone(),
-            delta_coverage: self.delta_coverage.clone(),
-        }
-    }
-}

pageserver/src/tenant/storage_layer.rs

@@ -196,50 +196,3 @@ pub fn downcast_remote_layer(
         None
     }
 }
-
-impl std::fmt::Debug for dyn Layer {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("Layer")
-            .field("short_id", &self.short_id())
-            .finish()
-    }
-}
-
-/// Holds metadata about a layer without any content. Used mostly for testing.
-pub struct LayerDescriptor {
-    pub key: Range<Key>,
-    pub lsn: Range<Lsn>,
-    pub is_incremental: bool,
-    pub short_id: String,
-}
-
-impl Layer for LayerDescriptor {
-    fn get_key_range(&self) -> Range<Key> {
-        self.key.clone()
-    }
-
-    fn get_lsn_range(&self) -> Range<Lsn> {
-        self.lsn.clone()
-    }
-
-    fn is_incremental(&self) -> bool {
-        self.is_incremental
-    }
-
-    fn get_value_reconstruct_data(
-        &self,
-        _key: Key,
-        _lsn_range: Range<Lsn>,
-        _reconstruct_data: &mut ValueReconstructState,
-    ) -> Result<ValueReconstructResult> {
-        todo!("This method shouldn't be part of the Layer trait")
-    }
-
-    fn short_id(&self) -> String {
-        self.short_id.clone()
-    }
-
-    fn dump(&self, _verbose: bool) -> Result<()> {
-        todo!()
-    }
-}

pageserver/src/tenant/timeline.rs

@@ -1148,7 +1148,6 @@ impl Timeline {
             }
         }
 
-        layers.rebuild_index();
         layers.next_open_layer_at = Some(Lsn(disk_consistent_lsn.0) + 1);
 
         info!(
@@ -1211,11 +1210,7 @@ impl Timeline {
                             anyhow::bail!("could not rename file {local_layer_path:?}: {err:?}");
                         } else {
                             self.metrics.resident_physical_size_gauge.sub(local_size);
-                            {
-                                let mut layers = self.layers.write().unwrap();
-                                layers.remove_historic(local_layer);
-                                layers.rebuild_index();
-                            }
+                            self.layers.write().unwrap().remove_historic(local_layer);
                             // fall-through to adding the remote layer
                         }
                     } else {
@@ -1257,11 +1252,7 @@ impl Timeline {
                     );
                     let remote_layer = Arc::new(remote_layer);
 
-                    {
-                        let mut layers = self.layers.write().unwrap();
-                        layers.insert_historic(remote_layer);
-                        layers.rebuild_index();
-                    }
+                    self.layers.write().unwrap().insert_historic(remote_layer);
                 }
                 LayerFileName::Delta(deltafilename) => {
                     // Create a RemoteLayer for the delta file.
@@ -1284,11 +1275,7 @@ impl Timeline {
                         &remote_layer_metadata,
                     );
                     let remote_layer = Arc::new(remote_layer);
-                    {
-                        let mut layers = self.layers.write().unwrap();
-                        layers.insert_historic(remote_layer);
-                        layers.rebuild_index();
-                    }
+                    self.layers.write().unwrap().insert_historic(remote_layer);
                 }
                 #[cfg(test)]
                 LayerFileName::Test(_) => unreachable!(),
@@ -2198,7 +2185,6 @@ impl Timeline {
         {
             let mut layers = self.layers.write().unwrap();
             layers.insert_historic(Arc::new(new_delta));
-            layers.rebuild_index();
         }
 
         // update the timeline's physical size
@@ -2263,15 +2249,13 @@ impl Timeline {
                 // are some delta layers *later* than current 'lsn', if more WAL was processed and flushed
                 // after we read last_record_lsn, which is passed here in the 'lsn' argument.
                 if img_lsn < lsn {
-                    let threshold = self.get_image_creation_threshold();
-                    let num_deltas =
-                        layers.count_deltas(&img_range, &(img_lsn..lsn), Some(threshold))?;
+                    let num_deltas = layers.count_deltas(&img_range, &(img_lsn..lsn))?;
 
                     debug!(
                         "key range {}-{}, has {} deltas on this timeline in LSN range {}..{}",
                         img_range.start, img_range.end, num_deltas, img_lsn, lsn
                     );
-                    if num_deltas >= threshold {
+                    if num_deltas >= self.get_image_creation_threshold() {
                         return Ok(true);
                     }
                 }
@@ -2376,7 +2360,6 @@ impl Timeline {
                 .add(metadata.len());
             layers.insert_historic(Arc::new(l));
         }
-        layers.rebuild_index();
         drop(layers);
         timer.stop_and_record();
 
@@ -2708,7 +2691,6 @@ impl Timeline {
             l.delete()?;
             layers.remove_historic(l);
         }
-        layers.rebuild_index();
         drop(layers);
 
         // Also schedule the deletions in remote storage
@@ -3022,7 +3004,6 @@ impl Timeline {
                 remote_client.schedule_layer_file_deletion(&layer_names_to_delete)?;
             }
         }
-        layers.rebuild_index();
 
         info!(
             "GC completed removing {} layers, cutoff {}",
@@ -3182,7 +3163,6 @@ impl Timeline {
                         layers.remove_historic(l);
                     }
                     layers.insert_historic(new_layer);
-                    layers.rebuild_index();
                     drop(layers);
 
                     // Now that we've inserted the download into the layer map,

proxy/src/stream.rs

@@ -2,7 +2,7 @@ use crate::error::UserFacingError;
 use anyhow::bail;
 use bytes::BytesMut;
 use pin_project_lite::pin_project;
-use pq_proto::{BeMessage, ConnectionError, FeMessage, FeStartupPacket};
+use pq_proto::{BeMessage, FeMessage, FeStartupPacket, ProtocolError};
 use rustls::ServerConfig;
 use std::pin::Pin;
 use std::sync::Arc;
@@ -53,7 +53,7 @@ impl<S: AsyncRead + Unpin> PqStream<S> {
         // TODO: `FeStartupPacket::read_fut` should return `FeStartupPacket`
         let msg = FeStartupPacket::read_fut(&mut self.stream)
             .await
-            .map_err(ConnectionError::into_io_error)?
+            .map_err(ProtocolError::into_io_error)?
             .ok_or_else(err_connection)?;
 
         match msg {
@@ -75,7 +75,7 @@ impl<S: AsyncRead + Unpin> PqStream<S> {
     async fn read_message(&mut self) -> io::Result<FeMessage> {
         FeMessage::read_fut(&mut self.stream)
             .await
-            .map_err(ConnectionError::into_io_error)?
+            .map_err(ProtocolError::into_io_error)?
             .ok_or_else(err_connection)
     }
 }

rust-toolchain.toml

@@ -4,7 +4,7 @@
 # version, we can consider updating.
 # See https://tracker.debian.org/pkg/rustc for more details on Debian rustc package,
 # we use "unstable" version number as the highest version used in the project by default.
-channel = "1.62.1"
+channel = "1.66.1"
 profile = "default"
 # The default profile includes rustc, rust-std, cargo, rust-docs, rustfmt and clippy.
 # https://rust-lang.github.io/rustup/concepts/profiles.html

safekeeper/Cargo.toml

@@ -14,12 +14,14 @@ clap = { version = "4.0", features = ["derive"] }
 const_format = "0.2.21"
 crc32c = "0.6.0"
 fs2 = "0.4.3"
+futures = "0.3"
 git-version = "0.3.5"
 hex = "0.4.3"
 humantime = "2.1.0"
 hyper = "0.14"
 nix = "0.25"
 once_cell = "1.13.0"
+pin-project-lite = "0.2"
 parking_lot = "0.12.1"
 postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
 postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }

safekeeper/src/bin/safekeeper.rs

@@ -228,20 +228,20 @@ fn start_safekeeper(conf: SafeKeeperConf) -> Result<()> {
 
     let conf_cloned = conf.clone();
     let safekeeper_thread = thread::Builder::new()
-        .name("safekeeper thread".into())
+        .name("WAL service thread".into())
         .spawn(|| wal_service::thread_main(conf_cloned, pg_listener))
         .unwrap();
 
     threads.push(safekeeper_thread);
 
     let conf_ = conf.clone();
-    threads.push(
-        thread::Builder::new()
-            .name("broker thread".into())
-            .spawn(|| {
-                broker::thread_main(conf_);
-            })?,
-    );
+    // threads.push(
+    //     thread::Builder::new()
+    //         .name("broker thread".into())
+    //         .spawn(|| {
+    //             broker::thread_main(conf_);
+    //         })?,
+    // );
 
     let conf_ = conf.clone();
     threads.push(

safekeeper/src/handler.rs

@@ -1,27 +1,23 @@
 //! Part of Safekeeper pretending to be Postgres, i.e. handling Postgres
 //! protocol commands.
 
+use anyhow::{bail, Context};
+use std::str;
+use tracing::{info, info_span, Instrument};
+
 use crate::auth::check_permission;
 use crate::json_ctrl::{handle_json_ctrl, AppendLogicalMessage};
 use crate::receive_wal::ReceiveWalConn;
-
-use crate::send_wal::ReplicationConn;
-
 use crate::{GlobalTimelines, SafeKeeperConf};
-use anyhow::Context;
-
 use postgres_ffi::PG_TLI;
-use regex::Regex;
-
 use pq_proto::{BeMessage, FeStartupPacket, RowDescriptor, INT4_OID, TEXT_OID};
-use std::str;
-use tracing::info;
+use regex::Regex;
 use utils::auth::{Claims, Scope};
 use utils::postgres_backend_async::QueryError;
 use utils::{
     id::{TenantId, TenantTimelineId, TimelineId},
     lsn::Lsn,
-    postgres_backend::{self, PostgresBackend},
+    postgres_backend_async::{self, PostgresBackend},
 };
 
 /// Safekeeper handler of postgres commands
@@ -41,9 +37,11 @@ enum SafekeeperPostgresCommand {
     StartReplication { start_lsn: Lsn },
     IdentifySystem,
     JSONCtrl { cmd: AppendLogicalMessage },
+    Show { guc: String },
 }
 
 fn parse_cmd(cmd: &str) -> anyhow::Result<SafekeeperPostgresCommand> {
+    let cmd_lowercase = cmd.to_ascii_lowercase();
     if cmd.starts_with("START_WAL_PUSH") {
         Ok(SafekeeperPostgresCommand::StartWalPush)
     } else if cmd.starts_with("START_REPLICATION") {
@@ -53,7 +51,7 @@ fn parse_cmd(cmd: &str) -> anyhow::Result<SafekeeperPostgresCommand> {
         let start_lsn = caps
             .next()
             .map(|cap| cap[1].parse::<Lsn>())
-            .context("failed to parse start LSN from START_REPLICATION command")??;
+            .context("parse start LSN from START_REPLICATION command")??;
         Ok(SafekeeperPostgresCommand::StartReplication { start_lsn })
     } else if cmd.starts_with("IDENTIFY_SYSTEM") {
         Ok(SafekeeperPostgresCommand::IdentifySystem)
@@ -62,12 +60,21 @@ fn parse_cmd(cmd: &str) -> anyhow::Result<SafekeeperPostgresCommand> {
         Ok(SafekeeperPostgresCommand::JSONCtrl {
             cmd: serde_json::from_str(cmd)?,
         })
+    } else if cmd_lowercase.starts_with("show") {
+        let re = Regex::new(r"show ((?:[[:alpha:]]|_)+)").unwrap();
+        let mut caps = re.captures_iter(&cmd_lowercase);
+        let guc = caps
+            .next()
+            .map(|cap| cap[1].parse::<String>())
+            .context("parse guc in SHOW command")??;
+        Ok(SafekeeperPostgresCommand::Show { guc })
     } else {
         anyhow::bail!("unsupported command {cmd}");
     }
 }
 
-impl postgres_backend::Handler for SafekeeperPostgresHandler {
+#[async_trait::async_trait]
+impl postgres_backend_async::Handler for SafekeeperPostgresHandler {
     // tenant_id and timeline_id are passed in connection string params
     fn startup(
         &mut self,
@@ -137,7 +144,7 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
         Ok(())
     }
 
-    fn process_query(
+    async fn process_query(
         &mut self,
         pgb: &mut PostgresBackend,
         query_string: &str,
@@ -147,9 +154,11 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
             .starts_with("set datestyle to ")
         {
             // important for debug because psycopg2 executes "SET datestyle TO 'ISO'" on connect
-            pgb.write_message(&BeMessage::CommandComplete(b"SELECT 1"))?;
+            pgb.write_message_flush(&BeMessage::CommandComplete(b"SELECT 1"))
+                .await?;
             return Ok(());
         }
+
         let cmd = parse_cmd(query_string)?;
 
         info!(
@@ -161,14 +170,22 @@ impl postgres_backend::Handler for SafekeeperPostgresHandler {
         let timeline_id = self.timeline_id.context("timelineid is required")?;
         self.check_permission(Some(tenant_id))?;
         self.ttid = TenantTimelineId::new(tenant_id, timeline_id);
+        let span_ttid = self.ttid; // satisfy borrow checker
 
         let res = match cmd {
-            SafekeeperPostgresCommand::StartWalPush => ReceiveWalConn::new(pgb).run(self),
+            // SafekeeperPostgresCommand::StartWalPush => ReceiveWalConn::new(pgb).run(self),
+            SafekeeperPostgresCommand::StartWalPush => Ok(()),
             SafekeeperPostgresCommand::StartReplication { start_lsn } => {
-                ReplicationConn::new(pgb).run(self, pgb, start_lsn)
+                self.handle_start_replication(pgb, start_lsn)
+                    .instrument(info_span!("WAL sender", ttid = %span_ttid))
+                    .await
             }
-            SafekeeperPostgresCommand::IdentifySystem => self.handle_identify_system(pgb),
-            SafekeeperPostgresCommand::JSONCtrl { ref cmd } => handle_json_ctrl(self, pgb, cmd),
+            SafekeeperPostgresCommand::IdentifySystem => self.handle_identify_system(pgb).await,
+            SafekeeperPostgresCommand::Show { guc } => self.handle_show(guc, pgb).await,
+            SafekeeperPostgresCommand::JSONCtrl { ref cmd } => {
+                handle_json_ctrl(self, pgb, cmd).await
+            }
+            _ => unreachable!(),
         };
 
         match res {
@@ -217,7 +234,10 @@ impl SafekeeperPostgresHandler {
     ///
     /// Handle IDENTIFY_SYSTEM replication command
     ///
-    fn handle_identify_system(&mut self, pgb: &mut PostgresBackend) -> Result<(), QueryError> {
+    async fn handle_identify_system(
+        &mut self,
+        pgb: &mut PostgresBackend,
+    ) -> Result<(), QueryError> {
         let tli = GlobalTimelines::get(self.ttid)?;
 
         let lsn = if self.is_walproposer_recovery() {
@@ -235,7 +255,7 @@ impl SafekeeperPostgresHandler {
         let tli_bytes = tli.as_bytes();
         let sysid_bytes = sysid.as_bytes();
 
-        pgb.write_message_noflush(&BeMessage::RowDescription(&[
+        pgb.write_message(&BeMessage::RowDescription(&[
             RowDescriptor {
                 name: b"systemid",
                 typoid: TEXT_OID,
@@ -261,13 +281,48 @@ impl SafekeeperPostgresHandler {
                 ..Default::default()
             },
         ]))?
-        .write_message_noflush(&BeMessage::DataRow(&[
+        .write_message(&BeMessage::DataRow(&[
             Some(sysid_bytes),
             Some(tli_bytes),
             Some(lsn_bytes),
             None,
         ]))?
-        .write_message(&BeMessage::CommandComplete(b"IDENTIFY_SYSTEM"))?;
+        .write_message_flush(&BeMessage::CommandComplete(b"IDENTIFY_SYSTEM"))
+        .await?;
+        Ok(())
+    }
+
+    async fn handle_show(
+        &mut self,
+        guc: String,
+        pgb: &mut PostgresBackend,
+    ) -> Result<(), QueryError> {
+        match guc.as_str() {
+            // pg_receivewal wants it
+            "data_directory_mode" => {
+                pgb.write_message(&BeMessage::RowDescription(&[RowDescriptor::int8_col(
+                    b"data_directory_mode",
+                )]))?
+                // xxx we could return real one, not just 0700
+                .write_message(&BeMessage::DataRow(&[Some(0700.to_string().as_bytes())]))?
+                .write_message(&BeMessage::CommandComplete(b"SELECT 1"))?;
+            }
+            // pg_receivewal wants it
+            "wal_segment_size" => {
+                let tli = GlobalTimelines::get(self.ttid)?;
+                let wal_seg_size = tli.get_state().1.server.wal_seg_size;
+                let wal_seg_size_mb = (wal_seg_size / 1024 / 1024).to_string() + "MB";
+
+                pgb.write_message(&BeMessage::RowDescription(&[RowDescriptor::text_col(
+                    b"wal_segment_size",
+                )]))?
+                .write_message(&BeMessage::DataRow(&[Some(wal_seg_size_mb.as_bytes())]))?
+                .write_message(&BeMessage::CommandComplete(b"SELECT 1"))?;
+            }
+            _ => {
+                return Err(anyhow::anyhow!("SHOW of unknown setting").into());
+            }
+        }
         Ok(())
     }
 

safekeeper/src/http/routes.rs

@@ -8,11 +8,14 @@ use serde::Serialize;
 use serde::Serializer;
 use std::collections::{HashMap, HashSet};
 use std::fmt::Display;
+use std::str::FromStr;
 use std::sync::Arc;
 use storage_broker::proto::SafekeeperTimelineInfo;
 use storage_broker::proto::TenantTimelineId as ProtoTenantTimelineId;
 use tokio::task::JoinError;
 
+use crate::json_ctrl::append_logical_message;
+use crate::json_ctrl::AppendLogicalMessage;
 use crate::safekeeper::ServerInfo;
 use crate::safekeeper::Term;
 
@@ -191,6 +194,50 @@ async fn timeline_create_handler(mut request: Request<Body>) -> Result<Response<
     json_response(StatusCode::OK, ())
 }
 
+// Create fake timeline + insert some valid WAL. Useful to test WAL streaming
+// from safekeeper in isolation, e.g.
+// pg_receivewal -v -d "host=localhost port=5454 options='-c tenant_id=deadbeefdeadbeefdeadbeefdeadbeef timeline_id=deadbeefdeadbeefdeadbeefdeadbeef'" -D ~/tmp/tmp/tmp
+// (hacking pg_receivewal startpos is currently needed though to make pg_receivewal work)
+async fn create_fake_timeline_handler(_request: Request<Body>) -> Result<Response<Body>, ApiError> {
+    let ttid = TenantTimelineId {
+        tenant_id: TenantId::from_str("deadbeefdeadbeefdeadbeefdeadbeef")
+            .expect("timeline_id parsing failed"),
+        timeline_id: TimelineId::from_str("deadbeefdeadbeefdeadbeefdeadbeef")
+            .expect("tenant_id parsing failed"),
+    };
+    let pg_version = 150000;
+    let server_info = ServerInfo {
+        pg_version,
+        system_id: 0,
+        wal_seg_size: WAL_SEGMENT_SIZE as u32,
+    };
+    let init_lsn = Lsn(0x1493AC8);
+    tokio::task::spawn_blocking(move || -> anyhow::Result<()> {
+        let tli = GlobalTimelines::create(ttid, server_info, init_lsn, init_lsn)?;
+        let mut begin_lsn = init_lsn;
+        for _ in 0..16 {
+            let append = AppendLogicalMessage {
+                lm_prefix: "db".to_owned(),
+                lm_message: "hahabubu".to_owned(),
+                set_commit_lsn: true,
+                send_proposer_elected: false, // actually ignored here
+                term: 0,
+                epoch_start_lsn: init_lsn,
+                begin_lsn,
+                truncate_lsn: init_lsn,
+                pg_version,
+            };
+            let inserted = append_logical_message(&tli, &append)?;
+            begin_lsn = inserted.end_lsn;
+        }
+        Ok(())
+    })
+    .await
+    .map_err(|e| ApiError::InternalServerError(e.into()))?
+    .map_err(ApiError::InternalServerError)?;
+    json_response(StatusCode::OK, ())
+}
+
 /// Deactivates the timeline and removes its data directory.
 async fn timeline_delete_force_handler(
     mut request: Request<Body>,
@@ -302,6 +349,7 @@ pub fn make_router(conf: SafeKeeperConf) -> RouterBuilder<hyper::Body, ApiError>
         .get("/v1/status", status_handler)
         // Will be used in the future instead of implicit timeline creation
         .post("/v1/tenant/timeline", timeline_create_handler)
+        .post("/v1/fake_timeline", create_fake_timeline_handler)
         .get(
             "/v1/tenant/:tenant_id/timeline/:timeline_id",
             timeline_status_handler,

safekeeper/src/json_ctrl.rs

@@ -26,26 +26,26 @@ use crate::GlobalTimelines;
 use postgres_ffi::encode_logical_message;
 use postgres_ffi::WAL_SEGMENT_SIZE;
 use pq_proto::{BeMessage, RowDescriptor, TEXT_OID};
-use utils::{lsn::Lsn, postgres_backend::PostgresBackend};
+use utils::{lsn::Lsn, postgres_backend_async::PostgresBackend};
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct AppendLogicalMessage {
     // prefix and message to build LogicalMessage
-    lm_prefix: String,
-    lm_message: String,
+    pub lm_prefix: String,
+    pub lm_message: String,
 
     // if true, commit_lsn will match flush_lsn after append
-    set_commit_lsn: bool,
+    pub set_commit_lsn: bool,
 
     // if true, ProposerElected will be sent before append
-    send_proposer_elected: bool,
+    pub send_proposer_elected: bool,
 
     // fields from AppendRequestHeader
-    term: Term,
-    epoch_start_lsn: Lsn,
-    begin_lsn: Lsn,
-    truncate_lsn: Lsn,
-    pg_version: u32,
+    pub term: Term,
+    pub epoch_start_lsn: Lsn,
+    pub begin_lsn: Lsn,
+    pub truncate_lsn: Lsn,
+    pub pg_version: u32,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -59,7 +59,7 @@ struct AppendResult {
 /// Handles command to craft logical message WAL record with given
 /// content, and then append it with specified term and lsn. This
 /// function is used to test safekeepers in different scenarios.
-pub fn handle_json_ctrl(
+pub async fn handle_json_ctrl(
     spg: &SafekeeperPostgresHandler,
     pgb: &mut PostgresBackend,
     append_request: &AppendLogicalMessage,
@@ -82,14 +82,15 @@ pub fn handle_json_ctrl(
     let response_data = serde_json::to_vec(&response)
         .with_context(|| format!("Response {response:?} is not a json array"))?;
 
-    pgb.write_message_noflush(&BeMessage::RowDescription(&[RowDescriptor {
+    pgb.write_message(&BeMessage::RowDescription(&[RowDescriptor {
         name: b"json",
         typoid: TEXT_OID,
         typlen: -1,
         ..Default::default()
     }]))?
-    .write_message_noflush(&BeMessage::DataRow(&[Some(&response_data)]))?
-    .write_message(&BeMessage::CommandComplete(b"JSON_CTRL"))?;
+    .write_message(&BeMessage::DataRow(&[Some(&response_data)]))?
+    .write_message_flush(&BeMessage::CommandComplete(b"JSON_CTRL"))
+    .await?;
     Ok(())
 }
 
@@ -128,15 +129,15 @@ fn send_proposer_elected(tli: &Arc<Timeline>, term: Term, lsn: Lsn) -> anyhow::R
 }
 
 #[derive(Debug, Serialize, Deserialize)]
-struct InsertedWAL {
+pub struct InsertedWAL {
     begin_lsn: Lsn,
-    end_lsn: Lsn,
+    pub end_lsn: Lsn,
     append_response: AppendResponse,
 }
 
 /// Extend local WAL with new LogicalMessage record. To do that,
 /// create AppendRequest with new WAL and pass it to safekeeper.
-fn append_logical_message(
+pub fn append_logical_message(
     tli: &Arc<Timeline>,
     msg: &AppendLogicalMessage,
 ) -> anyhow::Result<InsertedWAL> {

safekeeper/src/receive_wal.rs

@@ -26,7 +26,7 @@ use crate::safekeeper::ProposerAcceptorMessage;
 
 use crate::handler::SafekeeperPostgresHandler;
 use pq_proto::{BeMessage, FeMessage};
-use utils::{postgres_backend::PostgresBackend, sock_split::ReadStream};
+use utils::{postgres_backend_async::PostgresBackend, sock_split::ReadStream};
 
 pub struct ReceiveWalConn<'pg> {
     /// Postgres connection
@@ -59,82 +59,83 @@ impl<'pg> ReceiveWalConn<'pg> {
         // Notify the libpq client that it's allowed to send `CopyData` messages
         self.pg_backend
             .write_message(&BeMessage::CopyBothResponse)?;
-
-        let r = self
-            .pg_backend
-            .take_stream_in()
-            .ok_or_else(|| anyhow!("failed to take read stream from pgbackend"))?;
-        let mut poll_reader = ProposerPollStream::new(r)?;
+        Ok(())
+        // let r = self
+        //     .pg_backend
+        //     .take_stream_in()
+        //     .ok_or_else(|| anyhow!("failed to take read stream from pgbackend"))?;
+        // let mut poll_reader = ProposerPollStream::new(r)?;
 
         // Receive information about server
-        let next_msg = poll_reader.recv_msg()?;
-        let tli = match next_msg {
-            ProposerAcceptorMessage::Greeting(ref greeting) => {
-                info!(
-                    "start handshake with walproposer {} sysid {} timeline {}",
-                    self.peer_addr, greeting.system_id, greeting.tli,
-                );
-                let server_info = ServerInfo {
-                    pg_version: greeting.pg_version,
-                    system_id: greeting.system_id,
-                    wal_seg_size: greeting.wal_seg_size,
-                };
-                GlobalTimelines::create(spg.ttid, server_info, Lsn::INVALID, Lsn::INVALID)?
-            }
-            _ => {
-                return Err(QueryError::Other(anyhow::anyhow!(
-                    "unexpected message {next_msg:?} instead of greeting"
-                )))
-            }
-        };
+        // let next_msg = poll_reader.recv_msg()?;
+        // let tli = match next_msg {
+        //     ProposerAcceptorMessage::Greeting(ref greeting) => {
+        //         info!(
+        //             "start handshake with walproposer {} sysid {} timeline {}",
+        //             self.peer_addr, greeting.system_id, greeting.tli,
+        //         );
+        //         let server_info = ServerInfo {
+        //             pg_version: greeting.pg_version,
+        //             system_id: greeting.system_id,
+        //             wal_seg_size: greeting.wal_seg_size,
+        //         };
+        //         GlobalTimelines::create(spg.ttid, server_info, Lsn::INVALID, Lsn::INVALID)?
+        //     }
+        //     _ => {
+        //         return Err(QueryError::Other(anyhow::anyhow!(
+        //             "unexpected message {next_msg:?} instead of greeting"
+        //         )))
+        //     }
+        // };
 
-        let mut next_msg = Some(next_msg);
+        // let mut next_msg = None;
 
-        let mut first_time_through = true;
-        let mut _guard: Option<ComputeConnectionGuard> = None;
-        loop {
-            if matches!(next_msg, Some(ProposerAcceptorMessage::AppendRequest(_))) {
-                // poll AppendRequest's without blocking and write WAL to disk without flushing,
-                // while it's readily available
-                while let Some(ProposerAcceptorMessage::AppendRequest(append_request)) = next_msg {
-                    let msg = ProposerAcceptorMessage::NoFlushAppendRequest(append_request);
+        // let mut first_time_through = true;
+        // let mut _guard: Option<ComputeConnectionGuard> = None;
+        // loop {
+        //     if matches!(next_msg, Some(ProposerAcceptorMessage::AppendRequest(_))) {
+        //         // poll AppendRequest's without blocking and write WAL to disk without flushing,
+        //         // while it's readily available
+        //         while let Some(ProposerAcceptorMessage::AppendRequest(append_request)) = next_msg {
+        //             let msg = ProposerAcceptorMessage::NoFlushAppendRequest(append_request);
 
-                    let reply = tli.process_msg(&msg)?;
-                    if let Some(reply) = reply {
-                        self.write_msg(&reply)?;
-                    }
+        //             let reply = tli.process_msg(&msg)?;
+        //             if let Some(reply) = reply {
+        //                 self.write_msg(&reply)?;
+        //             }
 
-                    next_msg = poll_reader.poll_msg();
-                }
+        //             // next_msg = poll_reader.poll_msg();
+        //             next_msg = poll_reader.poll_msg();
+        //         }
 
-                // flush all written WAL to the disk
-                let reply = tli.process_msg(&ProposerAcceptorMessage::FlushWAL)?;
-                if let Some(reply) = reply {
-                    self.write_msg(&reply)?;
-                }
-            } else if let Some(msg) = next_msg.take() {
-                // process other message
-                let reply = tli.process_msg(&msg)?;
-                if let Some(reply) = reply {
-                    self.write_msg(&reply)?;
-                }
-            }
-            if first_time_through {
-                // Register the connection and defer unregister. Do that only
-                // after processing first message, as it sets wal_seg_size,
-                // wanted by many.
-                tli.on_compute_connect()?;
-                _guard = Some(ComputeConnectionGuard {
-                    timeline: Arc::clone(&tli),
-                });
-                first_time_through = false;
-            }
+        //         // flush all written WAL to the disk
+        //         let reply = tli.process_msg(&ProposerAcceptorMessage::FlushWAL)?;
+        //         if let Some(reply) = reply {
+        //             self.write_msg(&reply)?;
+        //         }
+        //     } else if let Some(msg) = next_msg.take() {
+        //         // process other message
+        //         let reply = tli.process_msg(&msg)?;
+        //         if let Some(reply) = reply {
+        //             self.write_msg(&reply)?;
+        //         }
+        //     }
+        //     if first_time_through {
+        //         // Register the connection and defer unregister. Do that only
+        //         // after processing first message, as it sets wal_seg_size,
+        //         // wanted by many.
+        //         tli.on_compute_connect()?;
+        //         _guard = Some(ComputeConnectionGuard {
+        //             timeline: Arc::clone(&tli),
+        //         });
+        //         first_time_through = false;
+        //     }
 
-            // blocking wait for the next message
-            if next_msg.is_none() {
-                next_msg = Some(poll_reader.recv_msg()?);
-            }
-        }
+        //     // blocking wait for the next message
+        //     if next_msg.is_none() {
+        //         next_msg = Some(poll_reader.recv_msg()?);
+        //     }
+        // }
     }
 }
 
@@ -144,37 +145,37 @@ struct ProposerPollStream {
 }
 
 impl ProposerPollStream {
-    fn new(mut r: ReadStream) -> anyhow::Result<Self> {
-        let (msg_tx, msg_rx) = channel();
+    // fn new(mut r: ReadStream) -> anyhow::Result<Self> {
+    //     let (msg_tx, msg_rx) = channel();
 
-        let read_thread = thread::Builder::new()
-            .name("Read WAL thread".into())
-            .spawn(move || -> Result<(), QueryError> {
-                loop {
-                    let copy_data = match FeMessage::read(&mut r)? {
-                        Some(FeMessage::CopyData(bytes)) => Ok(bytes),
-                        Some(msg) => Err(QueryError::Other(anyhow::anyhow!(
-                            "expected `CopyData` message, found {msg:?}"
-                        ))),
-                        None => Err(QueryError::from(std::io::Error::new(
-                            std::io::ErrorKind::ConnectionAborted,
-                            "walproposer closed the connection",
-                        ))),
-                    }?;
+    //     let read_thread = thread::Builder::new()
+    //         .name("Read WAL thread".into())
+    //         .spawn(move || -> Result<(), QueryError> {
+    //             loop {
+    //                 let copy_data = match FeMessage::read(&mut r)? {
+    //                     Some(FeMessage::CopyData(bytes)) => Ok(bytes),
+    //                     Some(msg) => Err(QueryError::Other(anyhow::anyhow!(
+    //                         "expected `CopyData` message, found {msg:?}"
+    //                     ))),
+    //                     None => Err(QueryError::from(std::io::Error::new(
+    //                         std::io::ErrorKind::ConnectionAborted,
+    //                         "walproposer closed the connection",
+    //                     ))),
+    //                 }?;
 
-                    let msg = ProposerAcceptorMessage::parse(copy_data)?;
-                    msg_tx
-                        .send(msg)
-                        .context("Failed to send the proposer message")?;
-                }
-                // msg_tx will be dropped here, this will also close msg_rx
-            })?;
+    //                 let msg = ProposerAcceptorMessage::parse(copy_data)?;
+    //                 msg_tx
+    //                     .send(msg)
+    //                     .context("Failed to send the proposer message")?;
+    //             }
+    //             // msg_tx will be dropped here, this will also close msg_rx
+    //         })?;
 
-        Ok(Self {
-            msg_rx,
-            read_thread: Some(read_thread),
-        })
-    }
+    //     Ok(Self {
+    //         msg_rx,
+    //         read_thread: Some(read_thread),
+    //     })
+    // }
 
     fn recv_msg(&mut self) -> Result<ProposerAcceptorMessage, QueryError> {
         self.msg_rx.recv().map_err(|_| {

safekeeper/src/send_wal.rs

@@ -1,28 +1,36 @@
 //! This module implements the streaming side of replication protocol, starting
 //! with the "START_REPLICATION" message.
 
+use anyhow::Context as AnyhowContext;
+use bytes::Bytes;
+use futures::future::BoxFuture;
+use pin_project_lite::pin_project;
+use postgres_ffi::get_current_timestamp;
+use postgres_ffi::{TimestampTz, MAX_SEND_SIZE};
+use serde::{Deserialize, Serialize};
+use std::cell::RefCell;
+use std::cmp::min;
+use std::future::Future;
+use std::pin::Pin;
+use std::rc::Rc;
+use std::sync::Arc;
+use std::task::{ready, Context, Poll};
+use std::time::Duration;
+use std::{io, str, thread};
+use tokio::sync::watch::Receiver;
+use tokio::time::timeout;
+use tracing::*;
+use utils::postgres_backend_async::QueryError;
+use utils::send_rc::RefCellSend;
+use utils::send_rc::SendRc;
+
+use pq_proto::{BeMessage, FeMessage, ReplicationFeedback, WalSndKeepAlive, XLogDataBody};
+use utils::{bin_ser::BeSer, lsn::Lsn, postgres_backend_async::PostgresBackend};
+
 use crate::handler::SafekeeperPostgresHandler;
 use crate::timeline::{ReplicaState, Timeline};
 use crate::wal_storage::WalReader;
 use crate::GlobalTimelines;
-use anyhow::Context;
-
-use bytes::Bytes;
-use postgres_ffi::get_current_timestamp;
-use postgres_ffi::{TimestampTz, MAX_SEND_SIZE};
-use serde::{Deserialize, Serialize};
-use std::cmp::min;
-use std::net::Shutdown;
-use std::sync::Arc;
-use std::time::Duration;
-use std::{io, str, thread};
-use utils::postgres_backend_async::QueryError;
-
-use pq_proto::{BeMessage, FeMessage, ReplicationFeedback, WalSndKeepAlive, XLogDataBody};
-use tokio::sync::watch::Receiver;
-use tokio::time::timeout;
-use tracing::*;
-use utils::{bin_ser::BeSer, lsn::Lsn, postgres_backend::PostgresBackend, sock_split::ReadStream};
 
 // See: https://www.postgresql.org/docs/13/protocol-replication.html
 const HOT_STANDBY_FEEDBACK_TAG_BYTE: u8 = b'h';
@@ -60,13 +68,6 @@ pub struct StandbyReply {
     pub reply_requested: bool,
 }
 
-/// A network connection that's speaking the replication protocol.
-pub struct ReplicationConn {
-    /// This is an `Option` because we will spawn a background thread that will
-    /// `take` it from us.
-    stream_in: Option<ReadStream>,
-}
-
 /// Scope guard to unregister replication connection from timeline
 struct ReplicationConnGuard {
     replica: usize, // replica internal ID assigned by timeline
@@ -79,230 +80,418 @@ impl Drop for ReplicationConnGuard {
     }
 }
 
-impl ReplicationConn {
-    /// Create a new `ReplicationConn`
-    pub fn new(pgb: &mut PostgresBackend) -> Self {
-        Self {
-            stream_in: pgb.take_stream_in(),
-        }
-    }
-
-    /// Handle incoming messages from the network.
-    /// This is spawned into the background by `handle_start_replication`.
-    fn background_thread(
-        mut stream_in: ReadStream,
-        replica_guard: Arc<ReplicationConnGuard>,
-    ) -> anyhow::Result<()> {
-        let replica_id = replica_guard.replica;
-        let timeline = &replica_guard.timeline;
-
-        let mut state = ReplicaState::new();
-        // Wait for replica's feedback.
-        while let Some(msg) = FeMessage::read(&mut stream_in)? {
-            match &msg {
-                FeMessage::CopyData(m) => {
-                    // There's three possible data messages that the client is supposed to send here:
-                    // `HotStandbyFeedback` and `StandbyStatusUpdate` and `NeonStandbyFeedback`.
-
-                    match m.first().cloned() {
-                        Some(HOT_STANDBY_FEEDBACK_TAG_BYTE) => {
-                            // Note: deserializing is on m[1..] because we skip the tag byte.
-                            state.hs_feedback = HotStandbyFeedback::des(&m[1..])
-                                .context("failed to deserialize HotStandbyFeedback")?;
-                            timeline.update_replica_state(replica_id, state);
-                        }
-                        Some(STANDBY_STATUS_UPDATE_TAG_BYTE) => {
-                            let _reply = StandbyReply::des(&m[1..])
-                                .context("failed to deserialize StandbyReply")?;
-                            // This must be a regular postgres replica,
-                            // because pageserver doesn't send this type of messages to safekeeper.
-                            // Currently this is not implemented, so this message is ignored.
-
-                            warn!("unexpected StandbyReply. Read-only postgres replicas are not supported in safekeepers yet.");
-                            // timeline.update_replica_state(replica_id, Some(state));
-                        }
-                        Some(NEON_STATUS_UPDATE_TAG_BYTE) => {
-                            // Note: deserializing is on m[9..] because we skip the tag byte and len bytes.
-                            let buf = Bytes::copy_from_slice(&m[9..]);
-                            let reply = ReplicationFeedback::parse(buf);
-
-                            trace!("ReplicationFeedback is {:?}", reply);
-                            // Only pageserver sends ReplicationFeedback, so set the flag.
-                            // This replica is the source of information to resend to compute.
-                            state.pageserver_feedback = Some(reply);
-
-                            timeline.update_replica_state(replica_id, state);
-                        }
-                        _ => warn!("unexpected message {:?}", msg),
-                    }
-                }
-                FeMessage::Sync => {}
-                FeMessage::CopyFail => {
-                    // Shutdown the connection, because rust-postgres client cannot be dropped
-                    // when connection is alive.
-                    let _ = stream_in.shutdown(Shutdown::Both);
-                    anyhow::bail!("Copy failed");
-                }
-                _ => {
-                    // We only handle `CopyData`, 'Sync', 'CopyFail' messages. Anything else is ignored.
-                    info!("unexpected message {:?}", msg);
-                }
-            }
-        }
-
-        Ok(())
-    }
-
-    ///
-    /// Handle START_REPLICATION replication command
-    ///
-    pub fn run(
+impl SafekeeperPostgresHandler {
+    pub async fn handle_start_replication(
         &mut self,
-        spg: &mut SafekeeperPostgresHandler,
         pgb: &mut PostgresBackend,
-        mut start_pos: Lsn,
+        start_pos: Lsn,
     ) -> Result<(), QueryError> {
-        let _enter = info_span!("WAL sender", ttid = %spg.ttid).entered();
-
-        let tli = GlobalTimelines::get(spg.ttid)?;
-
-        // spawn the background thread which receives HotStandbyFeedback messages.
-        let bg_timeline = Arc::clone(&tli);
-        let bg_stream_in = self.stream_in.take().unwrap();
-        let bg_timeline_id = spg.timeline_id.unwrap();
+        let appname = self.appname.clone();
+        let tli = GlobalTimelines::get(self.ttid)?;
 
         let state = ReplicaState::new();
         // This replica_id is used below to check if it's time to stop replication.
-        let replica_id = bg_timeline.add_replica(state);
+        let replica_id = tli.add_replica(state);
 
         // Use a guard object to remove our entry from the timeline, when the background
         // thread and us have both finished using it.
-        let replica_guard = Arc::new(ReplicationConnGuard {
+        let _guard = Arc::new(ReplicationConnGuard {
             replica: replica_id,
-            timeline: bg_timeline,
+            timeline: tli.clone(),
         });
-        let bg_replica_guard = Arc::clone(&replica_guard);
 
-        // TODO: here we got two threads, one for writing WAL and one for receiving
-        // feedback. If one of them fails, we should shutdown the other one too.
-        let _ = thread::Builder::new()
-            .name("HotStandbyFeedback thread".into())
-            .spawn(move || {
-                let _enter =
-                    info_span!("HotStandbyFeedback thread", timeline = %bg_timeline_id).entered();
-                if let Err(err) = Self::background_thread(bg_stream_in, bg_replica_guard) {
-                    error!("Replication background thread failed: {}", err);
+        // Walproposer gets special handling: safekeeper must give proposer all
+        // local WAL till the end, whether committed or not (walproposer will
+        // hang otherwise). That's because walproposer runs the consensus and
+        // synchronizes safekeepers on the most advanced one.
+        //
+        // There is a small risk of this WAL getting concurrently garbaged if
+        // another compute rises which collects majority and starts fixing log
+        // on this safekeeper itself. That's ok as (old) proposer will never be
+        // able to commit such WAL.
+        let stop_pos: Option<Lsn> = if self.is_walproposer_recovery() {
+            let wal_end = tli.get_flush_lsn();
+            Some(wal_end)
+        } else {
+            None
+        };
+        let end_pos = stop_pos.unwrap_or(Lsn::INVALID);
+
+        info!(
+            "starting streaming from {:?} till {:?}",
+            start_pos, stop_pos
+        );
+
+        // switch to copy
+        pgb.write_message(&BeMessage::CopyBothResponse)?;
+
+        let (_, persisted_state) = tli.get_state();
+        let wal_reader = WalReader::new(
+            self.conf.workdir.clone(),
+            self.conf.timeline_dir(&tli.ttid),
+            &persisted_state,
+            start_pos,
+            self.conf.wal_backup_enabled,
+        )?;
+        let write_ctx = SendRc::new(WriteContext {
+            wal_reader: RefCell::new(wal_reader),
+            send_buf: RefCell::new([0; MAX_SEND_SIZE]),
+        });
+
+        let mut c = ReplicationContext {
+            tli,
+            replica_id,
+            appname,
+            pgb,
+            start_pos,
+            end_pos,
+            stop_pos,
+            write_ctx,
+            feedback: ReplicaState::new(),
+        };
+
+        let _phantom_wf = c.wait_wal_fut();
+        let real_end_pos = c.end_pos;
+        c.end_pos = c.start_pos + 1; // to well form read_wal future
+        let _phantom_rf = c.read_wal_fut();
+        c.end_pos = real_end_pos;
+
+        ReplicationHandler {
+            c,
+            write_state: WriteState::FlushWal,
+            _phantom_wf,
+            _phantom_rf,
+        }
+        .await
+    }
+}
+
+pin_project! {
+    /// START_REPLICATION stream driver: sends WAL and receives feedback.
+    struct ReplicationHandler<'a, WF, RF>
+    where
+        WF: Future<Output = anyhow::Result<Option<Lsn>>>,
+        RF: Future<Output = anyhow::Result<usize>>,
+    {
+        c: ReplicationContext<'a>,
+        #[pin]
+        write_state: WriteState<WF, RF>,
+        // To deduce anonymous types.
+        _phantom_wf: WF,
+        _phantom_rf: RF,
+    }
+}
+
+/// Data ReplicationHandler maintains. Separated so we could generate WriteState
+/// futures during init, deducing their type.
+struct ReplicationContext<'a> {
+    tli: Arc<Timeline>,
+    appname: Option<String>,
+    replica_id: usize,
+    pgb: &'a mut PostgresBackend,
+    // Position since which we are sending next chunk.
+    start_pos: Lsn,
+    // WAL up to this position is known to be locally available.
+    end_pos: Lsn,
+    // If present, terminate after reaching this position; used by walproposer
+    // in recovery.
+    stop_pos: Option<Lsn>,
+    // This data is needed to create Future sending WAL, so we need to both have
+    // it here (to create new future) and borrow it to the future itself.
+    // Essentially this is a self referential struct. To satisfy borrow checker,
+    // use Rc<RefCell>. To make ReplicationHandler itself Send'able future, wrap
+    // it into SendRc; this is safe as ReplicationHandler is passed between
+    // threads only as a whole (during rescheduling).
+    //
+    // Right now we're in CurrentThread runtime, so Send is somewhat redundant;
+    // however, otherwise we'd need to inconveniently have separate !Send
+    // version of pg backend Handler trait (and work with LocalSet).
+    write_ctx: SendRc<WriteContext>,
+    feedback: ReplicaState,
+}
+
+// State which ReplicationHandler needs to create futures sending data.
+struct WriteContext {
+    wal_reader: RefCell<WalReader>,
+    // buffer for readling WAL into to send it
+    send_buf: RefCell<[u8; MAX_SEND_SIZE]>,
+}
+
+// Yield points of WAL sending machinery.
+pin_project! {
+    #[project = WriteStateProj]
+    enum WriteState<WF, RF>
+    where
+        WF: Future<Output = anyhow::Result<Option<Lsn>>>,
+        RF: Future<Output = anyhow::Result<usize>>,
+    {
+        WaitWal{ #[pin] fut: WF},
+        ReadWal{ #[pin] fut: RF},
+        FlushWal,
+    }
+}
+
+impl<WF, RF> Future for ReplicationHandler<'_, WF, RF>
+where
+    WF: Future<Output = anyhow::Result<Option<Lsn>>>,
+    RF: Future<Output = anyhow::Result<usize>>,
+{
+    type Output = Result<(), QueryError>;
+
+    // We need to read feedback from the socket and write data there at the same
+    // time. To avoid having to split socket, which creates messy split-join
+    // APIs, is problematic with TLS [1] and needs to manage two tasks, just run
+    // single task and use poll interfaces, basically manual state machine,
+    // which is simple here.
+    //
+    // [1] https://github.com/tokio-rs/tls/issues/40
+    //
+    // Completes only when the stream is over, technically on error currently.
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        if let Poll::Ready(r) = self.as_mut().poll_read(cx) {
+            return Poll::Ready(r);
+        }
+        self.as_mut().poll_write(cx)
+    }
+}
+
+impl<WF, RF> ReplicationHandler<'_, WF, RF>
+where
+    WF: Future<Output = anyhow::Result<Option<Lsn>>>,
+    RF: Future<Output = anyhow::Result<usize>>,
+{
+    // Poll reading, i.e. getting feedback and processing it. Completes only on error/end of stream.
+    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), QueryError>> {
+        loop {
+            match ready!(self.as_mut().project().c.pgb.poll_read_message(cx)) {
+                Ok(Some(msg)) => self.as_mut().handle_feedback(&msg)?,
+                Ok(None) => {
+                    return Poll::Ready(Err(QueryError::Other(anyhow::anyhow!(
+                        "EOF on replication stream"
+                    ))))
                 }
-            })?;
-
-        let runtime = tokio::runtime::Builder::new_current_thread()
-            .enable_all()
-            .build()?;
-
-        runtime.block_on(async move {
-            let (inmem_state, persisted_state) = tli.get_state();
-            // add persisted_state.timeline_start_lsn == Lsn(0) check
-
-            // Walproposer gets special handling: safekeeper must give proposer all
-            // local WAL till the end, whether committed or not (walproposer will
-            // hang otherwise). That's because walproposer runs the consensus and
-            // synchronizes safekeepers on the most advanced one.
-            //
-            // There is a small risk of this WAL getting concurrently garbaged if
-            // another compute rises which collects majority and starts fixing log
-            // on this safekeeper itself. That's ok as (old) proposer will never be
-            // able to commit such WAL.
-            let stop_pos: Option<Lsn> = if spg.is_walproposer_recovery() {
-                let wal_end = tli.get_flush_lsn();
-                Some(wal_end)
-            } else {
-                None
+                Err(err) => return Poll::Ready(Err(err.into())),
             };
+        }
+    }
 
-            info!("Start replication from {:?} till {:?}", start_pos, stop_pos);
-
-            // switch to copy
-            pgb.write_message(&BeMessage::CopyBothResponse)?;
-
-            let mut end_pos = stop_pos.unwrap_or(inmem_state.commit_lsn);
-
-            let mut wal_reader = WalReader::new(
-                spg.conf.workdir.clone(),
-                spg.conf.timeline_dir(&tli.ttid),
-                &persisted_state,
-                start_pos,
-                spg.conf.wal_backup_enabled,
-            )?;
-
-            // buffer for wal sending, limited by MAX_SEND_SIZE
-            let mut send_buf = vec![0u8; MAX_SEND_SIZE];
-
-            // watcher for commit_lsn updates
-            let mut commit_lsn_watch_rx = tli.get_commit_lsn_watch_rx();
-
-            loop {
-                if let Some(stop_pos) = stop_pos {
-                    if start_pos >= stop_pos {
-                        break; /* recovery finished */
+    fn handle_feedback(self: Pin<&mut Self>, msg: &FeMessage) -> Result<(), QueryError> {
+        let this = self.project();
+        match &msg {
+            FeMessage::CopyData(m) => {
+                // There's three possible data messages that the client is supposed to send here:
+                // `HotStandbyFeedback` and `StandbyStatusUpdate` and `NeonStandbyFeedback`.
+                match m.first().cloned() {
+                    Some(HOT_STANDBY_FEEDBACK_TAG_BYTE) => {
+                        // Note: deserializing is on m[1..] because we skip the tag byte.
+                        this.c.feedback.hs_feedback = HotStandbyFeedback::des(&m[1..])
+                            .context("failed to deserialize HotStandbyFeedback")?;
+                        this.c
+                            .tli
+                            .update_replica_state(this.c.replica_id, this.c.feedback);
                     }
-                    end_pos = stop_pos;
-                } else {
-                    /* Wait until we have some data to stream */
-                    let lsn = wait_for_lsn(&mut commit_lsn_watch_rx, start_pos).await?;
+                    Some(STANDBY_STATUS_UPDATE_TAG_BYTE) => {
+                        let _reply = StandbyReply::des(&m[1..])
+                            .context("failed to deserialize StandbyReply")?;
+                        // This must be a regular postgres replica,
+                        // because pageserver doesn't send this type of messages to safekeeper.
+                        // Currently we just ignore this, tracking progress for them is not supported.
+                    }
+                    Some(NEON_STATUS_UPDATE_TAG_BYTE) => {
+                        // Note: deserializing is on m[9..] because we skip the tag byte and len bytes.
+                        let buf = Bytes::copy_from_slice(&m[9..]);
+                        let reply = ReplicationFeedback::parse(buf);
 
-                    if let Some(lsn) = lsn {
-                        end_pos = lsn;
-                    } else {
-                        // TODO: also check once in a while whether we are walsender
-                        // to right pageserver.
-                        if tli.should_walsender_stop(replica_id) {
-                            // Shut down, timeline is suspended.
-                            return Err(QueryError::from(io::Error::new(
-                                io::ErrorKind::ConnectionAborted,
-                                format!("end streaming to {:?}", spg.appname),
-                            )));
-                        }
+                        trace!("ReplicationFeedback is {:?}", reply);
+                        // Only pageserver sends ReplicationFeedback, so set the flag.
+                        // This replica is the source of information to resend to compute.
+                        this.c.feedback.pageserver_feedback = Some(reply);
 
-                        // timeout expired: request pageserver status
-                        pgb.write_message(&BeMessage::KeepAlive(WalSndKeepAlive {
-                            sent_ptr: end_pos.0,
-                            timestamp: get_current_timestamp(),
-                            request_reply: true,
-                        }))?;
+                        this.c
+                            .tli
+                            .update_replica_state(this.c.replica_id, this.c.feedback);
+                    }
+                    _ => warn!("unexpected message {:?}", msg),
+                }
+            }
+            FeMessage::CopyFail => {
+                // XXX we should probably (tell pgb to) close the socket, as
+                // CopyFail in duplex copy is somewhat unexpected (at least to
+                // PG walsender; evidently client should finish it with
+                // CopyDone). Note that sync rust-postgres client (which we
+                // don't use anymore) hangs otherwise.
+                // https://github.com/sfackler/rust-postgres/issues/755
+                // https://github.com/neondatabase/neon/issues/935
+                //
+                return Err(anyhow::anyhow!("unexpected CopyFail").into());
+            }
+            _ => {
+                return Err(
+                    anyhow::anyhow!("unexpected message {:?} in replication stream", msg).into(),
+                );
+            }
+        };
+        Ok(())
+    }
+
+    // Poll writing, i.e. sending more WAL. Completes only on error or when we
+    // decide to shutdown connection -- receiver is caughtup and there is no
+    // active computes; this is still handled as Err though.
+    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), QueryError>> {
+        // send while we don't block or error out
+        loop {
+            match &mut self.as_mut().project().write_state.project() {
+                WriteStateProj::WaitWal { fut } => match ready!(fut.as_mut().poll(cx))? {
+                    Some(lsn) => {
+                        self.as_mut().project().c.end_pos = lsn;
+                        self.as_mut().start_read_wal();
                         continue;
                     }
+                    // Timed out waiting for WAL, send keepalive and possibly terminate.
+                    None => {
+                        let mut this = self.as_mut().project();
+                        if this.c.tli.should_walsender_stop(this.c.replica_id) {
+                            // Terminate if there is nothing more to send.
+                            // TODO close the stream properly
+                            return Poll::Ready(Err(anyhow::anyhow!(format!(
+                                "ending streaming to {:?} at {}, receiver is caughtup and there is no computes",
+                                self.c.appname, self.c.start_pos,
+                            )).into()));
+                        }
+                        let end_pos = this.c.end_pos.0;
+                        this.c
+                            .pgb
+                            .write_message(&BeMessage::KeepAlive(WalSndKeepAlive {
+                                sent_ptr: end_pos,
+                                timestamp: get_current_timestamp(),
+                                request_reply: true,
+                            }))?;
+                        /* flush KA */
+                        this.write_state.set(WriteState::FlushWal);
+                    }
+                },
+                WriteStateProj::ReadWal { fut } => {
+                    let read_len = ready!(fut.as_mut().poll(cx))?;
+                    assert!(read_len > 0, "read_len={}", read_len);
+
+                    let mut this = self.as_mut().project();
+                    let write_ctx_clone = this.c.write_ctx.clone();
+                    let send_buf = &write_ctx_clone.send_buf.borrow()[..read_len];
+                    let chunk_end = this.c.start_pos + read_len as u64;
+                    // write data to the output buffer
+                    this.c
+                        .pgb
+                        .write_message(&BeMessage::XLogData(XLogDataBody {
+                            wal_start: this.c.start_pos.0,
+                            wal_end: chunk_end.0,
+                            timestamp: get_current_timestamp(),
+                            data: send_buf,
+                        }))
+                        .context("Failed to write XLogData")?;
+                    trace!("wrote a chunk of wal {}-{}", this.c.start_pos, chunk_end);
+                    this.c.start_pos = chunk_end;
+                    // and flush it
+                    this.write_state.set(WriteState::FlushWal);
                 }
+                WriteStateProj::FlushWal => {
+                    let this = self.as_mut().project();
 
-                let send_size = end_pos.checked_sub(start_pos).unwrap().0 as usize;
-                let send_size = min(send_size, send_buf.len());
-
-                let send_buf = &mut send_buf[..send_size];
-
-                // read wal into buffer
-                let send_size = wal_reader.read(send_buf).await?;
-                let send_buf = &send_buf[..send_size];
-
-                // Write some data to the network socket.
-                pgb.write_message(&BeMessage::XLogData(XLogDataBody {
-                    wal_start: start_pos.0,
-                    wal_end: end_pos.0,
-                    timestamp: get_current_timestamp(),
-                    data: send_buf,
-                }))
-                .context("Failed to send XLogData")?;
-
-                start_pos += send_size as u64;
-                trace!("sent WAL up to {}", start_pos);
+                    ready!(this.c.pgb.poll_flush(cx))?;
+                    // If we are streaming to walproposer, check it is time to stop.
+                    if let Some(stop_pos) = this.c.stop_pos {
+                        if this.c.start_pos >= stop_pos {
+                            // recovery finished
+                            // TODO close the stream properly
+                            return Poll::Ready(Err(anyhow::anyhow!(format!(
+                                "ending streaming to walproposer at {}, receiver is caughtup and there is no computes",
+                                this.c.start_pos)).into()));
+                        }
+                        self.as_mut().start_read_wal();
+                        continue;
+                    } else {
+                        // if we don't know next portion is already available, wait
+                        // for it; otherwise proceed to sending
+                        if self.c.end_pos <= self.c.start_pos {
+                            self.as_mut().start_wait_wal();
+                        } else {
+                            self.as_mut().start_read_wal();
+                        }
+                    }
+                }
             }
+        }
+    }
 
-            Ok(())
-        })
+    // Start waiting for WAL, creating future doing that.
+    fn start_wait_wal(self: Pin<&mut Self>) {
+        let fut = self.c.wait_wal_fut();
+        self.project().write_state.set(WriteState::WaitWal {
+            fut: {
+                // SAFETY: this function is the only way to assign WaitWal to
+                // write_state. We just workaround impossibility of specifying
+                // async fn type, which is anonymous.
+                // transmute_copy is used as transmute refuses generic param:
+                // https://users.rust-lang.org/t/transmute-doesnt-work-on-generic-types/87272
+                assert_eq!(std::mem::size_of::<WF>(), std::mem::size_of_val(&fut));
+                let t = unsafe { std::mem::transmute_copy(&fut) };
+                std::mem::forget(fut);
+                t
+            },
+        });
+    }
+
+    // Switch into reading WAL state, creating Future doing that.
+    fn start_read_wal(self: Pin<&mut Self>) {
+        let fut = self.c.read_wal_fut();
+        self.project().write_state.set(WriteState::ReadWal {
+            fut: {
+                // SAFETY: this function is the only way to assign ReadWal to
+                // write_state. We just workaround impossibility of specifying
+                // async fn type, which is anonymous.
+                // transmute_copy is used as transmute refuses generic param:
+                // https://users.rust-lang.org/t/transmute-doesnt-work-on-generic-types/87272
+                assert_eq!(std::mem::size_of::<RF>(), std::mem::size_of_val(&fut));
+                let t = unsafe { std::mem::transmute_copy(&fut) };
+                std::mem::forget(fut);
+                t
+            },
+        });
+    }
+}
+
+impl ReplicationContext<'_> {
+    // Create future waiting for WAL.
+    fn wait_wal_fut(&self) -> impl Future<Output = anyhow::Result<Option<Lsn>>> {
+        let mut commit_lsn_watch_rx = self.tli.get_commit_lsn_watch_rx();
+        let start_pos = self.start_pos;
+        async move { wait_for_lsn(&mut commit_lsn_watch_rx, start_pos).await }
+    }
+
+    // Create future reading WAL.
+    fn read_wal_fut(&self) -> impl Future<Output = anyhow::Result<usize>> {
+        let mut send_size = self
+            .end_pos
+            .checked_sub(self.start_pos)
+            .expect("reading wal without waiting for it first")
+            .0 as usize;
+        send_size = min(send_size, self.write_ctx.send_buf.borrow().len());
+        let write_ctx_fut = self.write_ctx.clone();
+        async move {
+            let mut wal_reader_ref = write_ctx_fut.wal_reader.borrow_mut_send();
+            let mut send_buf_ref = write_ctx_fut.send_buf.borrow_mut_send();
+
+            let send_buf = &mut send_buf_ref[..send_size];
+            wal_reader_ref.read(send_buf).await
+        }
     }
 }
 
 const POLL_STATE_TIMEOUT: Duration = Duration::from_secs(1);
 
-// Wait until we have commit_lsn > lsn or timeout expires. Returns latest commit_lsn.
+// Wait until we have commit_lsn > lsn or timeout expires. Returns
+// - Ok(Some(commit_lsn)) if needed lsn is successfully observed;
+// - Ok(None) if timeout expired;
+// - Err in case of error (if watch channel is in trouble, shouldn't happen).
 async fn wait_for_lsn(rx: &mut Receiver<Lsn>, lsn: Lsn) -> anyhow::Result<Option<Lsn>> {
     let commit_lsn: Lsn = *rx.borrow();
     if commit_lsn > lsn {

safekeeper/src/wal_backup.rs

@@ -346,7 +346,9 @@ impl WalBackupTask {
                         backup_lsn, commit_lsn, e
                     );
 
-                    retry_attempt = retry_attempt.saturating_add(1);
+                    if retry_attempt < u32::MAX {
+                        retry_attempt += 1;
+                    }
                 }
             }
         }
@@ -385,7 +387,7 @@ async fn backup_single_segment(
 ) -> Result<()> {
     let segment_file_path = seg.file_path(timeline_dir)?;
     let remote_segment_path = segment_file_path
-        .strip_prefix(workspace_dir)
+        .strip_prefix(&workspace_dir)
         .context("Failed to strip workspace dir prefix")
         .and_then(RemotePath::new)
         .with_context(|| {
@@ -467,7 +469,7 @@ async fn backup_object(source_file: &Path, target_file: &RemotePath, size: usize
 pub async fn read_object(
     file_path: &RemotePath,
     offset: u64,
-) -> anyhow::Result<Pin<Box<dyn tokio::io::AsyncRead>>> {
+) -> anyhow::Result<Pin<Box<dyn tokio::io::AsyncRead + Send + Sync>>> {
     let storage = REMOTE_STORAGE
         .get()
         .context("Failed to get remote storage")?

safekeeper/src/wal_service.rs

@@ -2,36 +2,54 @@
 //!   WAL service listens for client connections and
 //!   receive WAL from wal_proposer and send it to WAL receivers
 //!
+use anyhow::{Context, Result};
 use regex::Regex;
-use std::net::{TcpListener, TcpStream};
-use std::thread;
+use std::{future, thread};
+use tokio::net::TcpStream;
 use tracing::*;
 use utils::postgres_backend_async::QueryError;
 
 use crate::handler::SafekeeperPostgresHandler;
 use crate::SafeKeeperConf;
-use utils::postgres_backend::{AuthType, PostgresBackend};
+use utils::postgres_backend_async::{AuthType, PostgresBackend};
 
 /// Accept incoming TCP connections and spawn them into a background thread.
-pub fn thread_main(conf: SafeKeeperConf, listener: TcpListener) -> ! {
-    loop {
-        match listener.accept() {
-            Ok((socket, peer_addr)) => {
-                debug!("accepted connection from {}", peer_addr);
-                let conf = conf.clone();
+pub fn thread_main(conf: SafeKeeperConf, pg_listener: std::net::TcpListener) {
+    let runtime = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .context("create runtime")
+        // todo catch error in main thread
+        .expect("failed to create runtime");
 
-                let _ = thread::Builder::new()
-                    .name("WAL service thread".into())
-                    .spawn(move || {
-                        if let Err(err) = handle_socket(socket, conf) {
-                            error!("connection handler exited: {}", err);
-                        }
-                    })
-                    .unwrap();
+    runtime
+        .block_on(async move {
+            // Tokio's from_std won't do this for us, per its comment.
+            pg_listener.set_nonblocking(true)?;
+            let listener = tokio::net::TcpListener::from_std(pg_listener)?;
+
+            loop {
+                match listener.accept().await {
+                    Ok((socket, peer_addr)) => {
+                        debug!("accepted connection from {}", peer_addr);
+                        let conf = conf.clone();
+
+                        let _ = thread::Builder::new()
+                            .name("WAL service thread".into())
+                            .spawn(move || {
+                                if let Err(err) = handle_socket(socket, conf) {
+                                    error!("connection handler exited: {}", err);
+                                }
+                            })
+                            .unwrap();
+                    }
+                    Err(e) => error!("Failed to accept connection: {}", e),
+                }
             }
-            Err(e) => error!("Failed to accept connection: {}", e),
-        }
-    }
+            #[allow(unreachable_code)] // hint compiler the closure return type
+            Ok::<(), anyhow::Error>(())
+        })
+        .expect("listener failed")
 }
 
 // Get unique thread id (Rust internal), with ThreadId removed for shorter printing
@@ -44,9 +62,14 @@ fn get_tid() -> u64 {
 
 /// This is run by `thread_main` above, inside a background thread.
 ///
-fn handle_socket(socket: TcpStream, conf: SafeKeeperConf) -> Result<(), QueryError> {
+fn handle_socket(mut socket: TcpStream, conf: SafeKeeperConf) -> Result<(), QueryError> {
     let _enter = info_span!("", tid = ?get_tid()).entered();
 
+    let runtime = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()?;
+    let local = tokio::task::LocalSet::new();
+
     socket.set_nodelay(true)?;
 
     let auth_type = match conf.auth {
@@ -54,9 +77,13 @@ fn handle_socket(socket: TcpStream, conf: SafeKeeperConf) -> Result<(), QueryErr
         Some(_) => AuthType::NeonJWT,
     };
     let mut conn_handler = SafekeeperPostgresHandler::new(conf);
-    let pgbackend = PostgresBackend::new(socket, auth_type, None, false)?;
-    // libpq replication protocol between safekeeper and replicas/pagers
-    pgbackend.run(&mut conn_handler)?;
+    let pgbackend = PostgresBackend::new(socket, auth_type, None)?;
+    // libpq protocol between safekeeper and walproposer / pageserver
+    // We don't use shutdown.
+    local.block_on(
+        &runtime,
+        pgbackend.run(&mut conn_handler, || future::pending::<()>()),
+    )?;
 
     Ok(())
 }

safekeeper/src/wal_storage.rs

@@ -450,7 +450,7 @@ pub struct WalReader {
     timeline_dir: PathBuf,
     wal_seg_size: usize,
     pos: Lsn,
-    wal_segment: Option<Pin<Box<dyn AsyncRead>>>,
+    wal_segment: Option<Pin<Box<dyn AsyncRead + Send + Sync>>>,
 
     // S3 will be used to read WAL if LSN is not available locally
     enable_remote_read: bool,
@@ -491,6 +491,11 @@ impl WalReader {
         })
     }
 
+    pub async fn fake_read(&mut self) -> Result<usize> {
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+        Ok(self.pos.0 as usize)
+    }
+
     pub async fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
         let mut wal_segment = match self.wal_segment.take() {
             Some(reader) => reader,
@@ -517,7 +522,7 @@ impl WalReader {
     }
 
     /// Open WAL segment at the current position of the reader.
-    async fn open_segment(&self) -> Result<Pin<Box<dyn AsyncRead>>> {
+    async fn open_segment(&self) -> Result<Pin<Box<dyn AsyncRead + Send + Sync>>> {
         let xlogoff = self.pos.segment_offset(self.wal_seg_size);
         let segno = self.pos.segment_number(self.wal_seg_size);
         let wal_file_name = XLogFileName(PG_TLI, segno, self.wal_seg_size);

workspace_hack/Cargo.toml

@@ -25,7 +25,6 @@ futures-task = { version = "0.3", default-features = false, features = ["alloc",
 futures-util = { version = "0.3", features = ["alloc", "async-await", "async-await-macro", "channel", "futures-channel", "futures-io", "futures-macro", "futures-sink", "io", "memchr", "sink", "slab", "std"] }
 hashbrown = { version = "0.12", features = ["ahash", "inline-more", "raw"] }
 indexmap = { version = "1", default-features = false, features = ["std"] }
-itertools = { version = "0.10", features = ["use_alloc", "use_std"] }
 libc = { version = "0.2", features = ["extra_traits", "std"] }
 log = { version = "0.4", default-features = false, features = ["serde", "std"] }
 memchr = { version = "2", features = ["std"] }
@@ -54,7 +53,6 @@ bytes = { version = "1", features = ["serde", "std"] }
 either = { version = "1", features = ["use_std"] }
 hashbrown = { version = "0.12", features = ["ahash", "inline-more", "raw"] }
 indexmap = { version = "1", default-features = false, features = ["std"] }
-itertools = { version = "0.10", features = ["use_alloc", "use_std"] }
 libc = { version = "0.2", features = ["extra_traits", "std"] }
 log = { version = "0.4", default-features = false, features = ["serde", "std"] }
 memchr = { version = "2", features = ["std"] }