run benchmarks

This code was mostly copied from walsender.c and the idea was to keep it
similar to walsender.c, so that we can easily copy-paste future upstream
changes to walsender.c to waproposer_utils.c, too. But right now I see that
deleting it doesn't break anything, so it's better to remove unused parts.

.github/workflows/build_and_test.yml

@@ -814,7 +814,7 @@ jobs:
       run:
         shell: sh -eu {0}
     env:
-      VM_BUILDER_VERSION: v0.15.0-alpha1
+      VM_BUILDER_VERSION: v0.15.4
 
     steps:
       - name: Checkout

compute_tools/src/bin/compute_ctl.rs

@@ -58,12 +58,16 @@ use compute_tools::monitor::launch_monitor;
 use compute_tools::params::*;
 use compute_tools::spec::*;
 
+// this is an arbitrary build tag. Fine as a default / for testing purposes
+// in-case of not-set environment var
+const BUILD_TAG_DEFAULT: &str = "5670669815";
+
 fn main() -> Result<()> {
     init_tracing_and_logging(DEFAULT_LOG_LEVEL)?;
 
-    // if there is no build tag environment variable, use a "default value"
-    // the use-case for this "default value" is regression tests that don't require BUILD_TAG
-    let build_tag = option_env!("BUILD_TAG").unwrap_or("5670669815").to_string();
+    let build_tag = option_env!("BUILD_TAG")
+        .unwrap_or(BUILD_TAG_DEFAULT)
+        .to_string();
     info!("build_tag: {build_tag}");
 
     let matches = cli().get_matches();

pageserver/src/page_cache.rs

@@ -47,13 +47,11 @@ use std::{
 
 use anyhow::Context;
 use once_cell::sync::OnceCell;
-use tracing::error;
 use utils::{
     id::{TenantId, TimelineId},
     lsn::Lsn,
 };
 
-use crate::tenant::writeback_ephemeral_file;
 use crate::{metrics::PageCacheSizeMetrics, repository::Key};
 
 static PAGE_CACHE: OnceCell<PageCache> = OnceCell::new();
@@ -97,10 +95,6 @@ enum CacheKey {
         hash_key: MaterializedPageHashKey,
         lsn: Lsn,
     },
-    EphemeralPage {
-        file_id: u64,
-        blkno: u32,
-    },
     ImmutableFilePage {
         file_id: u64,
         blkno: u32,
@@ -128,7 +122,6 @@ struct Slot {
 struct SlotInner {
     key: Option<CacheKey>,
     buf: &'static mut [u8; PAGE_SZ],
-    dirty: bool,
 }
 
 impl Slot {
@@ -177,8 +170,6 @@ pub struct PageCache {
     /// can have a separate mapping map, next to this field.
     materialized_page_map: RwLock<HashMap<MaterializedPageHashKey, Vec<Version>>>,
 
-    ephemeral_page_map: RwLock<HashMap<(u64, u32), usize>>,
-
     immutable_page_map: RwLock<HashMap<(u64, u32), usize>>,
 
     /// The actual buffers with their metadata.
@@ -258,14 +249,6 @@ impl PageWriteGuard<'_> {
         );
         self.valid = true;
     }
-    pub fn mark_dirty(&mut self) {
-        // only ephemeral pages can be dirty ATM.
-        assert!(matches!(
-            self.inner.key,
-            Some(CacheKey::EphemeralPage { .. })
-        ));
-        self.inner.dirty = true;
-    }
 }
 
 impl Drop for PageWriteGuard<'_> {
@@ -280,7 +263,6 @@ impl Drop for PageWriteGuard<'_> {
             let self_key = self.inner.key.as_ref().unwrap();
             PAGE_CACHE.get().unwrap().remove_mapping(self_key);
             self.inner.key = None;
-            self.inner.dirty = false;
         }
     }
 }
@@ -388,41 +370,7 @@ impl PageCache {
         Ok(())
     }
 
-    // Section 1.2: Public interface functions for working with Ephemeral pages.
-
-    pub fn read_ephemeral_buf(&self, file_id: u64, blkno: u32) -> anyhow::Result<ReadBufResult> {
-        let mut cache_key = CacheKey::EphemeralPage { file_id, blkno };
-
-        self.lock_for_read(&mut cache_key)
-    }
-
-    pub fn write_ephemeral_buf(&self, file_id: u64, blkno: u32) -> anyhow::Result<WriteBufResult> {
-        let cache_key = CacheKey::EphemeralPage { file_id, blkno };
-
-        self.lock_for_write(&cache_key)
-    }
-
-    /// Immediately drop all buffers belonging to given file, without writeback
-    pub fn drop_buffers_for_ephemeral(&self, drop_file_id: u64) {
-        for slot_idx in 0..self.slots.len() {
-            let slot = &self.slots[slot_idx];
-
-            let mut inner = slot.inner.write().unwrap();
-            if let Some(key) = &inner.key {
-                match key {
-                    CacheKey::EphemeralPage { file_id, blkno: _ } if *file_id == drop_file_id => {
-                        // remove mapping for old buffer
-                        self.remove_mapping(key);
-                        inner.key = None;
-                        inner.dirty = false;
-                    }
-                    _ => {}
-                }
-            }
-        }
-    }
-
-    // Section 1.3: Public interface functions for working with immutable file pages.
+    // Section 1.2: Public interface functions for working with immutable file pages.
 
     pub fn read_immutable_buf(&self, file_id: u64, blkno: u32) -> anyhow::Result<ReadBufResult> {
         let mut cache_key = CacheKey::ImmutableFilePage { file_id, blkno };
@@ -444,7 +392,6 @@ impl PageCache {
                         // remove mapping for old buffer
                         self.remove_mapping(key);
                         inner.key = None;
-                        inner.dirty = false;
                     }
                     _ => {}
                 }
@@ -522,10 +469,6 @@ impl PageCache {
             CacheKey::MaterializedPage { .. } => {
                 unreachable!("Materialized pages use lookup_materialized_page")
             }
-            CacheKey::EphemeralPage { .. } => (
-                &crate::metrics::PAGE_CACHE.read_accesses_ephemeral,
-                &crate::metrics::PAGE_CACHE.read_hits_ephemeral,
-            ),
             CacheKey::ImmutableFilePage { .. } => (
                 &crate::metrics::PAGE_CACHE.read_accesses_immutable,
                 &crate::metrics::PAGE_CACHE.read_hits_immutable,
@@ -566,7 +509,6 @@ impl PageCache {
             // Make the slot ready
             let slot = &self.slots[slot_idx];
             inner.key = Some(cache_key.clone());
-            inner.dirty = false;
             slot.usage_count.store(1, Ordering::Relaxed);
 
             return Ok(ReadBufResult::NotFound(PageWriteGuard {
@@ -628,7 +570,6 @@ impl PageCache {
             // Make the slot ready
             let slot = &self.slots[slot_idx];
             inner.key = Some(cache_key.clone());
-            inner.dirty = false;
             slot.usage_count.store(1, Ordering::Relaxed);
 
             return Ok(WriteBufResult::NotFound(PageWriteGuard {
@@ -667,10 +608,6 @@ impl PageCache {
                 *lsn = version.lsn;
                 Some(version.slot_idx)
             }
-            CacheKey::EphemeralPage { file_id, blkno } => {
-                let map = self.ephemeral_page_map.read().unwrap();
-                Some(*map.get(&(*file_id, *blkno))?)
-            }
             CacheKey::ImmutableFilePage { file_id, blkno } => {
                 let map = self.immutable_page_map.read().unwrap();
                 Some(*map.get(&(*file_id, *blkno))?)
@@ -694,10 +631,6 @@ impl PageCache {
                     None
                 }
             }
-            CacheKey::EphemeralPage { file_id, blkno } => {
-                let map = self.ephemeral_page_map.read().unwrap();
-                Some(*map.get(&(*file_id, *blkno))?)
-            }
             CacheKey::ImmutableFilePage { file_id, blkno } => {
                 let map = self.immutable_page_map.read().unwrap();
                 Some(*map.get(&(*file_id, *blkno))?)
@@ -731,12 +664,6 @@ impl PageCache {
                     panic!("could not find old key in mapping")
                 }
             }
-            CacheKey::EphemeralPage { file_id, blkno } => {
-                let mut map = self.ephemeral_page_map.write().unwrap();
-                map.remove(&(*file_id, *blkno))
-                    .expect("could not find old key in mapping");
-                self.size_metrics.current_bytes_ephemeral.sub_page_sz(1);
-            }
             CacheKey::ImmutableFilePage { file_id, blkno } => {
                 let mut map = self.immutable_page_map.write().unwrap();
                 map.remove(&(*file_id, *blkno))
@@ -776,17 +703,7 @@ impl PageCache {
                     }
                 }
             }
-            CacheKey::EphemeralPage { file_id, blkno } => {
-                let mut map = self.ephemeral_page_map.write().unwrap();
-                match map.entry((*file_id, *blkno)) {
-                    Entry::Occupied(entry) => Some(*entry.get()),
-                    Entry::Vacant(entry) => {
-                        entry.insert(slot_idx);
-                        self.size_metrics.current_bytes_ephemeral.add_page_sz(1);
-                        None
-                    }
-                }
-            }
+
             CacheKey::ImmutableFilePage { file_id, blkno } => {
                 let mut map = self.immutable_page_map.write().unwrap();
                 match map.entry((*file_id, *blkno)) {
@@ -837,25 +754,8 @@ impl PageCache {
                     }
                 };
                 if let Some(old_key) = &inner.key {
-                    if inner.dirty {
-                        if let Err(err) = Self::writeback(old_key, inner.buf) {
-                            // Writing the page to disk failed.
-                            //
-                            // FIXME: What to do here, when? We could propagate the error to the
-                            // caller, but victim buffer is generally unrelated to the original
-                            // call. It can even belong to a different tenant. Currently, we
-                            // report the error to the log and continue the clock sweep to find
-                            // a different victim. But if the problem persists, the page cache
-                            // could fill up with dirty pages that we cannot evict, and we will
-                            // loop retrying the writebacks indefinitely.
-                            error!("writeback of buffer {:?} failed: {}", old_key, err);
-                            continue;
-                        }
-                    }
-
                     // remove mapping for old buffer
                     self.remove_mapping(old_key);
-                    inner.dirty = false;
                     inner.key = None;
                 }
                 return Ok((slot_idx, inner));
@@ -863,28 +763,6 @@ impl PageCache {
         }
     }
 
-    fn writeback(cache_key: &CacheKey, buf: &[u8]) -> Result<(), std::io::Error> {
-        match cache_key {
-            CacheKey::MaterializedPage {
-                hash_key: _,
-                lsn: _,
-            } => Err(std::io::Error::new(
-                std::io::ErrorKind::Other,
-                "unexpected dirty materialized page",
-            )),
-            CacheKey::EphemeralPage { file_id, blkno } => {
-                writeback_ephemeral_file(*file_id, *blkno, buf)
-            }
-            CacheKey::ImmutableFilePage {
-                file_id: _,
-                blkno: _,
-            } => Err(std::io::Error::new(
-                std::io::ErrorKind::Other,
-                "unexpected dirty immutable page",
-            )),
-        }
-    }
-
     /// Initialize a new page cache
     ///
     /// This should be called only once at page server startup.
@@ -895,7 +773,6 @@ impl PageCache {
 
         let size_metrics = &crate::metrics::PAGE_CACHE_SIZE;
         size_metrics.max_bytes.set_page_sz(num_pages);
-        size_metrics.current_bytes_ephemeral.set_page_sz(0);
         size_metrics.current_bytes_immutable.set_page_sz(0);
         size_metrics.current_bytes_materialized_page.set_page_sz(0);
 
@@ -905,11 +782,7 @@ impl PageCache {
                 let buf: &mut [u8; PAGE_SZ] = chunk.try_into().unwrap();
 
                 Slot {
-                    inner: RwLock::new(SlotInner {
-                        key: None,
-                        buf,
-                        dirty: false,
-                    }),
+                    inner: RwLock::new(SlotInner { key: None, buf }),
                     usage_count: AtomicU8::new(0),
                 }
             })
@@ -917,7 +790,6 @@ impl PageCache {
 
         Self {
             materialized_page_map: Default::default(),
-            ephemeral_page_map: Default::default(),
             immutable_page_map: Default::default(),
             slots,
             next_evict_slot: AtomicUsize::new(0),

pageserver/src/tenant.rs

@@ -136,9 +136,6 @@ pub use timeline::{
     LocalLayerInfoForDiskUsageEviction, LogicalSizeCalculationCause, PageReconstructError, Timeline,
 };
 
-// re-export this function so that page_cache.rs can use it.
-pub use crate::tenant::ephemeral_file::writeback as writeback_ephemeral_file;
-
 // re-export for use in remote_timeline_client.rs
 pub use crate::tenant::metadata::save_metadata;
 

pageserver/src/tenant/ephemeral_file.rs

@@ -2,7 +2,7 @@
 //! used to keep in-memory layers spilled on disk.
 
 use crate::config::PageServerConf;
-use crate::page_cache::{self, ReadBufResult, WriteBufResult, PAGE_SZ};
+use crate::page_cache::PAGE_SZ;
 use crate::tenant::blob_io::BlobWriter;
 use crate::tenant::block_io::BlockReader;
 use crate::virtual_file::VirtualFile;
@@ -10,7 +10,7 @@ use once_cell::sync::Lazy;
 use std::cmp::min;
 use std::collections::HashMap;
 use std::fs::OpenOptions;
-use std::io::{self, ErrorKind};
+use std::io::{self};
 use std::ops::DerefMut;
 use std::path::PathBuf;
 use std::sync::{Arc, RwLock};
@@ -19,6 +19,9 @@ use utils::id::{TenantId, TimelineId};
 
 use std::os::unix::fs::FileExt;
 
+mod buffer_pool;
+mod dirty_buffer;
+
 ///
 /// This is the global cache of file descriptors (File objects).
 ///
@@ -94,27 +97,13 @@ impl EphemeralFile {
         Ok(())
     }
 
-    fn get_buf_for_write(&self, blkno: u32) -> Result<page_cache::PageWriteGuard, io::Error> {
-        // Look up the right page
-        let cache = page_cache::get();
-        let mut write_guard = match cache
-            .write_ephemeral_buf(self.file_id, blkno)
-            .map_err(|e| to_io_error(e, "Failed to write ephemeral buf"))?
-        {
-            WriteBufResult::Found(guard) => guard,
-            WriteBufResult::NotFound(mut guard) => {
-                // Read the page from disk into the buffer
-                // TODO: if we're overwriting the whole page, no need to read it in first
-                self.fill_buffer(guard.deref_mut(), blkno)?;
-                guard.mark_valid();
-
-                // And then fall through to modify it.
-                guard
-            }
-        };
-        write_guard.mark_dirty();
-
-        Ok(write_guard)
+    fn get_buf_for_write(&self, blkno: u32) -> Result<dirty_buffer::Buffer, io::Error> {
+        let pool = buffer_pool::get();
+        let mut buf = pool.get_buffer();
+        // Read the page from disk into the buffer
+        // TODO: if we're overwriting the whole page, no need to read it in first
+        self.fill_buffer(buf.deref_mut(), blkno)?;
+        Ok(dirty_buffer::Buffer::new(self, buf, blkno))
     }
 }
 
@@ -127,75 +116,6 @@ pub fn is_ephemeral_file(filename: &str) -> bool {
     }
 }
 
-impl FileExt for EphemeralFile {
-    fn read_at(&self, dstbuf: &mut [u8], offset: u64) -> Result<usize, io::Error> {
-        // Look up the right page
-        let blkno = (offset / PAGE_SZ as u64) as u32;
-        let off = offset as usize % PAGE_SZ;
-        let len = min(PAGE_SZ - off, dstbuf.len());
-
-        let read_guard;
-        let mut write_guard;
-
-        let cache = page_cache::get();
-        let buf = match cache
-            .read_ephemeral_buf(self.file_id, blkno)
-            .map_err(|e| to_io_error(e, "Failed to read ephemeral buf"))?
-        {
-            ReadBufResult::Found(guard) => {
-                read_guard = guard;
-                read_guard.as_ref()
-            }
-            ReadBufResult::NotFound(guard) => {
-                // Read the page from disk into the buffer
-                write_guard = guard;
-                self.fill_buffer(write_guard.deref_mut(), blkno)?;
-                write_guard.mark_valid();
-
-                // And then fall through to read the requested slice from the
-                // buffer.
-                write_guard.as_ref()
-            }
-        };
-
-        dstbuf[0..len].copy_from_slice(&buf[off..(off + len)]);
-        Ok(len)
-    }
-
-    fn write_at(&self, srcbuf: &[u8], offset: u64) -> Result<usize, io::Error> {
-        // Look up the right page
-        let blkno = (offset / PAGE_SZ as u64) as u32;
-        let off = offset as usize % PAGE_SZ;
-        let len = min(PAGE_SZ - off, srcbuf.len());
-
-        let mut write_guard;
-        let cache = page_cache::get();
-        let buf = match cache
-            .write_ephemeral_buf(self.file_id, blkno)
-            .map_err(|e| to_io_error(e, "Failed to write ephemeral buf"))?
-        {
-            WriteBufResult::Found(guard) => {
-                write_guard = guard;
-                write_guard.deref_mut()
-            }
-            WriteBufResult::NotFound(guard) => {
-                // Read the page from disk into the buffer
-                // TODO: if we're overwriting the whole page, no need to read it in first
-                write_guard = guard;
-                self.fill_buffer(write_guard.deref_mut(), blkno)?;
-                write_guard.mark_valid();
-
-                // And then fall through to modify it.
-                write_guard.deref_mut()
-            }
-        };
-
-        buf[off..(off + len)].copy_from_slice(&srcbuf[0..len]);
-        write_guard.mark_dirty();
-        Ok(len)
-    }
-}
-
 impl BlobWriter for EphemeralFile {
     fn write_blob(&mut self, srcbuf: &[u8]) -> Result<u64, io::Error> {
         let pos = self.size;
@@ -217,6 +137,7 @@ impl BlobWriter for EphemeralFile {
                 // it needs to be split across pages
                 buf[off..(off + thislen)].copy_from_slice(&len_buf[..thislen]);
                 blknum += 1;
+                buf.writeback()?;
                 buf = self.get_buf_for_write(blknum)?;
                 buf[0..4 - thislen].copy_from_slice(&len_buf[thislen..]);
                 off = 4 - thislen;
@@ -232,6 +153,7 @@ impl BlobWriter for EphemeralFile {
             let mut page_remain = PAGE_SZ - off;
             if page_remain == 0 {
                 blknum += 1;
+                buf.writeback()?;
                 buf = self.get_buf_for_write(blknum)?;
                 off = 0;
                 page_remain = PAGE_SZ;
@@ -241,7 +163,8 @@ impl BlobWriter for EphemeralFile {
             off += this_blk_len;
             buf_remain = &buf_remain[this_blk_len..];
         }
-        drop(buf);
+
+        buf.writeback()?;
 
         if srcbuf.len() < 0x80 {
             self.size += 1;
@@ -256,10 +179,6 @@ impl BlobWriter for EphemeralFile {
 
 impl Drop for EphemeralFile {
     fn drop(&mut self) {
-        // drop all pages from page cache
-        let cache = page_cache::get();
-        cache.drop_buffers_for_ephemeral(self.file_id);
-
         // remove entry from the hash map
         EPHEMERAL_FILES.write().unwrap().files.remove(&self.file_id);
 
@@ -281,62 +200,24 @@ impl Drop for EphemeralFile {
     }
 }
 
-pub fn writeback(file_id: u64, blkno: u32, buf: &[u8]) -> Result<(), io::Error> {
-    if let Some(file) = EPHEMERAL_FILES.read().unwrap().files.get(&file_id) {
-        match file.write_all_at(buf, blkno as u64 * PAGE_SZ as u64) {
-            Ok(_) => Ok(()),
-            Err(e) => Err(io::Error::new(
-                ErrorKind::Other,
-                format!(
-                    "failed to write back to ephemeral file at {} error: {}",
-                    file.path.display(),
-                    e
-                ),
-            )),
-        }
-    } else {
-        Err(io::Error::new(
-            ErrorKind::Other,
-            "could not write back page, not found in ephemeral files hash",
-        ))
-    }
-}
-
 impl BlockReader for EphemeralFile {
-    type BlockLease = page_cache::PageReadGuard<'static>;
+    type BlockLease = buffer_pool::Handle;
 
     fn read_blk(&self, blknum: u32) -> Result<Self::BlockLease, io::Error> {
-        // Look up the right page
-        let cache = page_cache::get();
-        loop {
-            match cache
-                .read_ephemeral_buf(self.file_id, blknum)
-                .map_err(|e| to_io_error(e, "Failed to read ephemeral buf"))?
-            {
-                ReadBufResult::Found(guard) => return Ok(guard),
-                ReadBufResult::NotFound(mut write_guard) => {
-                    // Read the page from disk into the buffer
-                    self.fill_buffer(write_guard.deref_mut(), blknum)?;
-                    write_guard.mark_valid();
-
-                    // Swap for read lock
-                    continue;
-                }
-            };
-        }
+        // Read the page from disk into the buffer
+        let pool = buffer_pool::get();
+        let mut buf = pool.get_buffer();
+        self.fill_buffer(buf.deref_mut(), blknum)?;
+        Ok(buf)
     }
 }
 
-fn to_io_error(e: anyhow::Error, context: &str) -> io::Error {
-    io::Error::new(ErrorKind::Other, format!("{context}: {e:#}"))
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::tenant::blob_io::BlobWriter;
     use crate::tenant::block_io::BlockCursor;
-    use rand::{seq::SliceRandom, thread_rng, RngCore};
+    use rand::{thread_rng, RngCore};
     use std::fs;
     use std::str::FromStr;
 
@@ -357,50 +238,6 @@ mod tests {
         Ok((conf, tenant_id, timeline_id))
     }
 
-    // Helper function to slurp contents of a file, starting at the current position,
-    // into a string
-    fn read_string(efile: &EphemeralFile, offset: u64, len: usize) -> Result<String, io::Error> {
-        let mut buf = Vec::new();
-        buf.resize(len, 0u8);
-
-        efile.read_exact_at(&mut buf, offset)?;
-
-        Ok(String::from_utf8_lossy(&buf)
-            .trim_end_matches('\0')
-            .to_string())
-    }
-
-    #[test]
-    fn test_ephemeral_files() -> Result<(), io::Error> {
-        let (conf, tenant_id, timeline_id) = harness("ephemeral_files")?;
-
-        let file_a = EphemeralFile::create(conf, tenant_id, timeline_id)?;
-
-        file_a.write_all_at(b"foo", 0)?;
-        assert_eq!("foo", read_string(&file_a, 0, 20)?);
-
-        file_a.write_all_at(b"bar", 3)?;
-        assert_eq!("foobar", read_string(&file_a, 0, 20)?);
-
-        // Open a lot of files, enough to cause some page evictions.
-        let mut efiles = Vec::new();
-        for fileno in 0..100 {
-            let efile = EphemeralFile::create(conf, tenant_id, timeline_id)?;
-            efile.write_all_at(format!("file {}", fileno).as_bytes(), 0)?;
-            assert_eq!(format!("file {}", fileno), read_string(&efile, 0, 10)?);
-            efiles.push((fileno, efile));
-        }
-
-        // Check that all the files can still be read from. Use them in random order for
-        // good measure.
-        efiles.as_mut_slice().shuffle(&mut thread_rng());
-        for (fileno, efile) in efiles.iter_mut() {
-            assert_eq!(format!("file {}", fileno), read_string(efile, 0, 10)?);
-        }
-
-        Ok(())
-    }
-
     #[test]
     fn test_ephemeral_blobs() -> Result<(), io::Error> {
         let (conf, tenant_id, timeline_id) = harness("ephemeral_blobs")?;

pageserver/src/tenant/ephemeral_file/buffer_pool.rs

@@ -0,0 +1,66 @@
+//! Buffer pool for ephemeral file buffers.
+//!
+//! Currently this is a very simple implementation that just uses `malloc`.
+//! But the interface is such that we can switch to a more sophisticated
+//! implementation later, e.g., one that caps that amount of memory used.
+
+use std::ops::{Deref, DerefMut};
+
+use crate::page_cache::PAGE_SZ;
+
+pub struct BufferPool;
+
+const POOL: BufferPool = BufferPool;
+
+pub(super) fn get() -> &'static BufferPool {
+    &POOL
+}
+
+impl BufferPool {
+    /// Get a [`Handle`] to a buffer in the pool.
+    ///
+    /// The buffer is guaranteed to be zeroed out.
+    ///
+    /// The implementation may block to wait for buffers to become available,
+    /// and a future async version of this method may `.await` internally to
+    /// wait for buffers to become available.
+    ///
+    /// To avoid deadlocks, a thread/task must get all the buffers it needs
+    /// with a single call to `get_buffer`. Without this rule, a deadlock
+    /// can happen. Take for example a buffer pool with 2 buffers X, Y
+    /// and a program with two threads A and B, each requiring 2 buffers.
+    /// If A gets X and B gets Y, then both threads will block forever trying
+    /// to get their second buffer.
+    pub fn get_buffer(&self) -> Handle {
+        Handle {
+            data: vec![0; PAGE_SZ],
+        }
+    }
+}
+
+pub struct Handle {
+    data: Vec<u8>,
+}
+
+impl std::fmt::Debug for Handle {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Handle")
+            .field("data", &self.data.as_ptr())
+            .finish()
+    }
+}
+
+impl Deref for Handle {
+    type Target = [u8; PAGE_SZ];
+    fn deref(&self) -> &Self::Target {
+        let slice: &[u8] = &self.data[..];
+        slice.try_into().unwrap()
+    }
+}
+
+impl DerefMut for Handle {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        let slice: &mut [u8] = &mut self.data[..];
+        slice.try_into().unwrap()
+    }
+}

pageserver/src/tenant/ephemeral_file/dirty_buffer.rs

@@ -0,0 +1,111 @@
+//! Newtypes to ensure that dirty buffers are written back to the filesystem before they are dropped.
+
+use std::io::ErrorKind;
+use std::ops::Deref;
+use std::ops::DerefMut;
+use std::os::unix::prelude::FileExt;
+
+use crate::page_cache::PAGE_SZ;
+
+use super::buffer_pool;
+use super::EphemeralFile;
+
+pub(super) struct Buffer<'f> {
+    inner: Inner<'f>,
+}
+
+enum Inner<'f> {
+    Dirty {
+        ephemeral_file: &'f EphemeralFile,
+        buf: buffer_pool::Handle,
+        blkno: u32,
+    },
+    WritebackOngoing,
+    WrittenBack,
+    WritebackError,
+    Dropped,
+}
+
+impl<'f> Buffer<'f> {
+    pub(super) fn new(
+        ephemeral_file: &'f EphemeralFile,
+        buf: buffer_pool::Handle,
+        blkno: u32,
+    ) -> Self {
+        Self {
+            inner: Inner::Dirty {
+                ephemeral_file,
+                buf,
+                blkno,
+            },
+        }
+    }
+    pub(super) fn writeback(mut self) -> Result<(), std::io::Error> {
+        let Inner::Dirty {
+        ephemeral_file,
+        buf,
+        blkno,
+    } = std::mem::replace(&mut self.inner, Inner::WritebackOngoing) else {
+        unreachable!("writeback consumes");
+    };
+        match ephemeral_file
+            .file
+            .write_all_at(buf.deref(), blkno as u64 * PAGE_SZ as u64)
+        {
+            Ok(_) => {
+                self.inner = Inner::WrittenBack;
+                Ok(())
+            }
+            Err(e) => {
+                self.inner = Inner::WritebackError;
+                Err(std::io::Error::new(
+                    ErrorKind::Other,
+                    format!(
+                        "failed to write back to ephemeral file at {} error: {}",
+                        ephemeral_file.file.path.display(),
+                        e
+                    ),
+                ))
+            }
+        }
+    }
+}
+
+impl<'f> Deref for Buffer<'f> {
+    type Target = [u8];
+
+    fn deref(&self) -> &[u8] {
+        match &self.inner {
+            Inner::Dirty { buf, .. } => &**buf,
+            Inner::WritebackOngoing => unreachable!("writeback consumes"),
+            Inner::WrittenBack => unreachable!("writeback consumes"),
+            Inner::WritebackError => unreachable!("writeback consumes"),
+            Inner::Dropped => unreachable!(),
+        }
+    }
+}
+
+impl<'f> DerefMut for Buffer<'f> {
+    fn deref_mut(&mut self) -> &mut [u8] {
+        match &mut self.inner {
+            Inner::Dirty { buf, .. } => &mut **buf,
+            Inner::WritebackOngoing => unreachable!("writeback consumes"),
+            Inner::WrittenBack => unreachable!("writeback consumes"),
+            Inner::WritebackError => unreachable!("writeback consumes"),
+            Inner::Dropped => unreachable!(),
+        }
+    }
+}
+
+impl Drop for Buffer<'_> {
+    fn drop(&mut self) {
+        let prev = std::mem::replace(&mut self.inner, Inner::Dropped);
+        match prev {
+            // TODO: check this at compile time
+            Inner::Dirty { .. } => panic!("dropped dirty buffer, need to writeback() first"),
+            Inner::WritebackOngoing => unreachable!("transitory state"),
+            Inner::WrittenBack | Inner::WritebackError => {}
+            Inner::Dropped => unreachable!("drop only happens once"),
+        }
+    }
+}

pgxn/neon/walproposer_utils.c

@@ -37,68 +37,14 @@ static XLogSegNo walpropSegNo = 0;
 
 /* START cloned file-local variables and functions from walsender.c */
 
-/*
- * xlogreader used for replication.  Note that a WAL sender doing physical
- * replication does not need xlogreader to read WAL, but it needs one to
- * keep a state of its work.
- */
-static XLogReaderState *xlogreader = NULL;
-
-/*
- * These variables keep track of the state of the timeline we're currently
- * sending. sendTimeLine identifies the timeline. If sendTimeLineIsHistoric,
- * the timeline is not the latest timeline on this server, and the server's
- * history forked off from that timeline at sendTimeLineValidUpto.
- */
-static TimeLineID sendTimeLine = 0;
-static TimeLineID sendTimeLineNextTLI = 0;
-static bool sendTimeLineIsHistoric = false;
-static XLogRecPtr sendTimeLineValidUpto = InvalidXLogRecPtr;
-
-/*
- * Timestamp of last ProcessRepliesIfAny() that saw a reply from the
- * standby. Set to 0 if wal_sender_timeout doesn't need to be active.
- */
-static TimestampTz last_reply_timestamp = 0;
-
-/* Have we sent a heartbeat message asking for reply, since last reply? */
-static bool waiting_for_ping_response = false;
-
-static bool streamingDoneSending;
-static bool streamingDoneReceiving;
-
-/* Are we there yet? */
-static bool WalSndCaughtUp = false;
-
-/* Flags set by signal handlers for later service in main loop */
-static volatile sig_atomic_t got_STOPPING = false;
-
 /*
  * How far have we sent WAL already? This is also advertised in
  * MyWalSnd->sentPtr.  (Actually, this is the next WAL location to send.)
  */
 static XLogRecPtr sentPtr = InvalidXLogRecPtr;
 
-/*
- * This is set while we are streaming. When not set
- * PROCSIG_WALSND_INIT_STOPPING signal will be handled like SIGTERM. When set,
- * the main loop is responsible for checking got_STOPPING and terminating when
- * it's set (after streaming any remaining WAL).
- */
-static volatile sig_atomic_t replication_active = false;
-
-typedef void (*WalSndSendDataCallback) (void);
-static void WalSndLoop(WalSndSendDataCallback send_data);
-static void XLogSendPhysical(void);
-#if PG_VERSION_NUM >= 150000
-static XLogRecPtr GetStandbyFlushRecPtr(TimeLineID *tli);
-#else
-static XLogRecPtr GetStandbyFlushRecPtr(void);
-#endif
-
-static void WalSndSegmentOpen(XLogReaderState *state, XLogSegNo nextSegNo,
-							  TimeLineID *tli_p);
-
+static void WalSndLoop(void);
+static void XLogBroadcastWalProposer(void);
 /* END cloned file-level variables and functions from walsender.c */
 
 int
@@ -506,7 +452,7 @@ XLogWalPropClose(XLogRecPtr recptr)
 /* START of cloned functions from walsender.c */
 
 /*
- * Handle START_REPLICATION command.
+ * Subscribe for new WAL and stream it in the loop to safekeepers.
  *
  * At the moment, this never returns, but an ereport(ERROR) will take us back
  * to the main loop.
@@ -524,18 +470,6 @@ StartProposerReplication(StartReplicationCmd *cmd)
 				 errmsg("IDENTIFY_SYSTEM has not been run before START_REPLICATION")));
 #endif
 
-	/* create xlogreader for physical replication */
-	xlogreader =
-		XLogReaderAllocate(wal_segment_size, NULL,
-						   XL_ROUTINE(.segment_open = WalSndSegmentOpen,
-									  .segment_close = wal_segment_close),
-						   NULL);
-
-	if (!xlogreader)
-		ereport(ERROR,
-				(errcode(ERRCODE_OUT_OF_MEMORY),
-				 errmsg("out of memory")));
-
 	/*
 	 * We assume here that we're logging enough information in the WAL for
 	 * log-shipping, since this is checked in PostmasterMain().
@@ -569,341 +503,61 @@ StartProposerReplication(StartReplicationCmd *cmd)
 	 * we keep this code around to lighten the load for when we need it.
 	 */
 #if PG_VERSION_NUM >= 150000
-	if (am_cascading_walsender)
-	{
-		/* this also updates ThisTimeLineID */
-		FlushPtr = GetStandbyFlushRecPtr(&currTLI);
-	}
-	else
-		FlushPtr = GetFlushRecPtr(&currTLI);
+	FlushPtr = GetFlushRecPtr(&currTLI);
 #else
-	if (am_cascading_walsender)
-	{
-		/* this also updates ThisTimeLineID */
-		FlushPtr = GetStandbyFlushRecPtr();
-	}
-	else
-		FlushPtr = GetFlushRecPtr();
-
+	FlushPtr = GetFlushRecPtr();
 	currTLI = ThisTimeLineID;
 #endif
 
+	/*
+	 * When we first start replication the standby will be behind the
+	 * primary. For some applications, for example synchronous
+	 * replication, it is important to have a clear state for this initial
+	 * catchup mode, so we can trigger actions when we change streaming
+	 * state later. We may stay in this state for a long time, which is
+	 * exactly why we want to be able to monitor whether or not we are
+	 * still here.
+	 */
+	WalSndSetState(WALSNDSTATE_CATCHUP);
 
-	if (cmd->timeline != 0)
+	/*
+	 * Don't allow a request to stream from a future point in WAL that
+	 * hasn't been flushed to disk in this server yet.
+	 */
+	if (FlushPtr < cmd->startpoint)
 	{
-		XLogRecPtr	switchpoint;
-
-		sendTimeLine = cmd->timeline;
-		if (sendTimeLine == currTLI)
-		{
-			sendTimeLineIsHistoric = false;
-			sendTimeLineValidUpto = InvalidXLogRecPtr;
-		}
-		else
-		{
-			List	   *timeLineHistory;
-
-			sendTimeLineIsHistoric = true;
-
-			/*
-			 * Check that the timeline the client requested exists, and the
-			 * requested start location is on that timeline.
-			 */
-			timeLineHistory = readTimeLineHistory(currTLI);
-			switchpoint = tliSwitchPoint(cmd->timeline, timeLineHistory,
-										 &sendTimeLineNextTLI);
-			list_free_deep(timeLineHistory);
-
-			/*
-			 * Found the requested timeline in the history. Check that
-			 * requested startpoint is on that timeline in our history.
-			 *
-			 * This is quite loose on purpose. We only check that we didn't
-			 * fork off the requested timeline before the switchpoint. We
-			 * don't check that we switched *to* it before the requested
-			 * starting point. This is because the client can legitimately
-			 * request to start replication from the beginning of the WAL
-			 * segment that contains switchpoint, but on the new timeline, so
-			 * that it doesn't end up with a partial segment. If you ask for
-			 * too old a starting point, you'll get an error later when we
-			 * fail to find the requested WAL segment in pg_wal.
-			 *
-			 * XXX: we could be more strict here and only allow a startpoint
-			 * that's older than the switchpoint, if it's still in the same
-			 * WAL segment.
-			 */
-			if (!XLogRecPtrIsInvalid(switchpoint) &&
-				switchpoint < cmd->startpoint)
-			{
-				ereport(ERROR,
-						(errmsg("requested starting point %X/%X on timeline %u is not in this server's history",
-								LSN_FORMAT_ARGS(cmd->startpoint),
-								cmd->timeline),
-						 errdetail("This server's history forked from timeline %u at %X/%X.",
-								   cmd->timeline,
-								   LSN_FORMAT_ARGS(switchpoint))));
-			}
-			sendTimeLineValidUpto = switchpoint;
-		}
-	}
-	else
-	{
-		sendTimeLine = currTLI;
-		sendTimeLineValidUpto = InvalidXLogRecPtr;
-		sendTimeLineIsHistoric = false;
+		ereport(ERROR,
+				(errmsg("requested starting point %X/%X is ahead of the WAL flush position of this server %X/%X",
+						LSN_FORMAT_ARGS(cmd->startpoint),
+						LSN_FORMAT_ARGS(FlushPtr))));
 	}
 
-	streamingDoneSending = streamingDoneReceiving = false;
+	/* Start streaming from the requested point */
+	sentPtr = cmd->startpoint;
 
-	/* If there is nothing to stream, don't even enter COPY mode */
-	if (!sendTimeLineIsHistoric || cmd->startpoint < sendTimeLineValidUpto)
-	{
-		/*
-		 * When we first start replication the standby will be behind the
-		 * primary. For some applications, for example synchronous
-		 * replication, it is important to have a clear state for this initial
-		 * catchup mode, so we can trigger actions when we change streaming
-		 * state later. We may stay in this state for a long time, which is
-		 * exactly why we want to be able to monitor whether or not we are
-		 * still here.
-		 */
-		WalSndSetState(WALSNDSTATE_CATCHUP);
+	/* Initialize shared memory status, too */
+	SpinLockAcquire(&MyWalSnd->mutex);
+	MyWalSnd->sentPtr = sentPtr;
+	SpinLockRelease(&MyWalSnd->mutex);
 
-		/*
-		 * Don't allow a request to stream from a future point in WAL that
-		 * hasn't been flushed to disk in this server yet.
-		 */
-		if (FlushPtr < cmd->startpoint)
-		{
-			ereport(ERROR,
-					(errmsg("requested starting point %X/%X is ahead of the WAL flush position of this server %X/%X",
-							LSN_FORMAT_ARGS(cmd->startpoint),
-							LSN_FORMAT_ARGS(FlushPtr))));
-		}
+	SyncRepInitConfig();
 
-		/* Start streaming from the requested point */
-		sentPtr = cmd->startpoint;
+	/* Infinite send loop, never returns */
+	WalSndLoop();
 
-		/* Initialize shared memory status, too */
-		SpinLockAcquire(&MyWalSnd->mutex);
-		MyWalSnd->sentPtr = sentPtr;
-		SpinLockRelease(&MyWalSnd->mutex);
-
-		SyncRepInitConfig();
-
-		/* Main loop of walsender */
-		replication_active = true;
-
-		WalSndLoop(XLogSendPhysical);
-
-		replication_active = false;
-		if (got_STOPPING)
-			proc_exit(0);
-		WalSndSetState(WALSNDSTATE_STARTUP);
-
-		Assert(streamingDoneSending && streamingDoneReceiving);
-	}
+	WalSndSetState(WALSNDSTATE_STARTUP);
 
 	if (cmd->slotname)
 		ReplicationSlotRelease();
-
-	/*
-	 * Copy is finished now. Send a single-row result set indicating the next
-	 * timeline.
-	 */
-	if (sendTimeLineIsHistoric)
-	{
-		char		startpos_str[8 + 1 + 8 + 1];
-		DestReceiver *dest;
-		TupOutputState *tstate;
-		TupleDesc	tupdesc;
-		Datum		values[2];
-		bool		nulls[2];
-
-		snprintf(startpos_str, sizeof(startpos_str), "%X/%X",
-				 LSN_FORMAT_ARGS(sendTimeLineValidUpto));
-
-		dest = CreateDestReceiver(DestRemoteSimple);
-		MemSet(nulls, false, sizeof(nulls));
-
-		/*
-		 * Need a tuple descriptor representing two columns. int8 may seem
-		 * like a surprising data type for this, but in theory int4 would not
-		 * be wide enough for this, as TimeLineID is unsigned.
-		 */
-		tupdesc = CreateTemplateTupleDesc(2);
-		TupleDescInitBuiltinEntry(tupdesc, (AttrNumber) 1, "next_tli",
-								  INT8OID, -1, 0);
-		TupleDescInitBuiltinEntry(tupdesc, (AttrNumber) 2, "next_tli_startpos",
-								  TEXTOID, -1, 0);
-
-		/* prepare for projection of tuple */
-		tstate = begin_tup_output_tupdesc(dest, tupdesc, &TTSOpsVirtual);
-
-		values[0] = Int64GetDatum((int64) sendTimeLineNextTLI);
-		values[1] = CStringGetTextDatum(startpos_str);
-
-		/* send it to dest */
-		do_tup_output(tstate, values, nulls);
-
-		end_tup_output(tstate);
-	}
-
-	/* Send CommandComplete message */
-	EndReplicationCommand("START_STREAMING");
 }
 
-#if PG_VERSION_NUM >= 150000
-static XLogRecPtr
-GetStandbyFlushRecPtr(TimeLineID *tli)
-{
-	XLogRecPtr	replayPtr;
-	TimeLineID	replayTLI;
-	XLogRecPtr	receivePtr;
-	TimeLineID	receiveTLI;
-	XLogRecPtr	result;
-
-	/*
-	 * We can safely send what's already been replayed. Also, if walreceiver
-	 * is streaming WAL from the same timeline, we can send anything that it
-	 * has streamed, but hasn't been replayed yet.
-	 */
-
-	receivePtr = GetWalRcvFlushRecPtr(NULL, &receiveTLI);
-	replayPtr = GetXLogReplayRecPtr(&replayTLI);
-
-	*tli = replayTLI;
-
-	result = replayPtr;
-	if (receiveTLI == replayTLI && receivePtr > replayPtr)
-		result = receivePtr;
-
-	return result;
-}
-#else
 /*
- * Returns the latest point in WAL that has been safely flushed to disk, and
- * can be sent to the standby. This should only be called when in recovery,
- * ie. we're streaming to a cascaded standby.
- *
- * As a side-effect, ThisTimeLineID is updated to the TLI of the last
- * replayed WAL record.
+ * Main loop that waits for LSN updates and calls the walproposer.
+ * Synchronous replication sets latch in WalSndWakeup at walsender.c
  */
-static XLogRecPtr
-GetStandbyFlushRecPtr(void)
-{
-	XLogRecPtr	replayPtr;
-	TimeLineID	replayTLI;
-	XLogRecPtr	receivePtr;
-	TimeLineID	receiveTLI;
-	XLogRecPtr	result;
-
-	/*
-	 * We can safely send what's already been replayed. Also, if walreceiver
-	 * is streaming WAL from the same timeline, we can send anything that it
-	 * has streamed, but hasn't been replayed yet.
-	 */
-
-	receivePtr = GetWalRcvFlushRecPtr(NULL, &receiveTLI);
-	replayPtr = GetXLogReplayRecPtr(&replayTLI);
-
-	ThisTimeLineID = replayTLI;
-
-	result = replayPtr;
-	if (receiveTLI == ThisTimeLineID && receivePtr > replayPtr)
-		result = receivePtr;
-
-	return result;
-}
-#endif
-
-
-
-/* XLogReaderRoutine->segment_open callback */
 static void
-WalSndSegmentOpen(XLogReaderState *state, XLogSegNo nextSegNo,
-				  TimeLineID *tli_p)
+WalSndLoop(void)
 {
-	char		path[MAXPGPATH];
-
-	/*-------
-	 * When reading from a historic timeline, and there is a timeline switch
-	 * within this segment, read from the WAL segment belonging to the new
-	 * timeline.
-	 *
-	 * For example, imagine that this server is currently on timeline 5, and
-	 * we're streaming timeline 4. The switch from timeline 4 to 5 happened at
-	 * 0/13002088. In pg_wal, we have these files:
-	 *
-	 * ...
-	 * 000000040000000000000012
-	 * 000000040000000000000013
-	 * 000000050000000000000013
-	 * 000000050000000000000014
-	 * ...
-	 *
-	 * In this situation, when requested to send the WAL from segment 0x13, on
-	 * timeline 4, we read the WAL from file 000000050000000000000013. Archive
-	 * recovery prefers files from newer timelines, so if the segment was
-	 * restored from the archive on this server, the file belonging to the old
-	 * timeline, 000000040000000000000013, might not exist. Their contents are
-	 * equal up to the switchpoint, because at a timeline switch, the used
-	 * portion of the old segment is copied to the new file.  -------
-	 */
-	*tli_p = sendTimeLine;
-	if (sendTimeLineIsHistoric)
-	{
-		XLogSegNo	endSegNo;
-
-		XLByteToSeg(sendTimeLineValidUpto, endSegNo, state->segcxt.ws_segsize);
-		if (nextSegNo == endSegNo)
-			*tli_p = sendTimeLineNextTLI;
-	}
-
-	XLogFilePath(path, *tli_p, nextSegNo, state->segcxt.ws_segsize);
-	state->seg.ws_file = BasicOpenFile(path, O_RDONLY | PG_BINARY);
-	if (state->seg.ws_file >= 0)
-		return;
-
-	/*
-	 * If the file is not found, assume it's because the standby asked for a
-	 * too old WAL segment that has already been removed or recycled.
-	 */
-	if (errno == ENOENT)
-	{
-		char		xlogfname[MAXFNAMELEN];
-		int			save_errno = errno;
-
-		XLogFileName(xlogfname, *tli_p, nextSegNo, wal_segment_size);
-		errno = save_errno;
-		ereport(ERROR,
-				(errcode_for_file_access(),
-				 errmsg("requested WAL segment %s has already been removed",
-						xlogfname)));
-	}
-	else
-		ereport(ERROR,
-				(errcode_for_file_access(),
-				 errmsg("could not open file \"%s\": %m",
-						path)));
-}
-
-
-/* Main loop of walsender process that streams the WAL over Copy messages. */
-static void
-WalSndLoop(WalSndSendDataCallback send_data)
-{
-	/*
-	 * Initialize the last reply timestamp. That enables timeout processing
-	 * from hereon.
-	 */
-	last_reply_timestamp = GetCurrentTimestamp();
-	waiting_for_ping_response = false;
-
-	/*
-	 * Loop until we reach the end of this timeline or the client requests to
-	 * stop streaming.
-	 */
 	for (;;)
 	{
 		/* Clear any already-pending wakeups */
@@ -911,153 +565,41 @@ WalSndLoop(WalSndSendDataCallback send_data)
 
 		CHECK_FOR_INTERRUPTS();
 
-		/* Process any requests or signals received recently */
-		if (ConfigReloadPending)
-		{
-			ConfigReloadPending = false;
-			ProcessConfigFile(PGC_SIGHUP);
-			SyncRepInitConfig();
-		}
+		XLogBroadcastWalProposer();
 
-		/* always true */
-		if (am_wal_proposer)
-		{
-			send_data();
-			if (WalSndCaughtUp)
-			{
-				if (MyWalSnd->state == WALSNDSTATE_CATCHUP)
-					WalSndSetState(WALSNDSTATE_STREAMING);
-				WalProposerPoll();
-				WalSndCaughtUp = false;
-			}
-			continue;
-		}
+		if (MyWalSnd->state == WALSNDSTATE_CATCHUP)
+			WalSndSetState(WALSNDSTATE_STREAMING);
+		WalProposerPoll();
 	}
 }
 
 /*
- * Send out the WAL in its normal physical/stored form.
- *
- * Read up to MAX_SEND_SIZE bytes of WAL that's been flushed to disk,
- * but not yet sent to the client, and buffer it in the libpq output
- * buffer.
- *
- * If there is no unsent WAL remaining, WalSndCaughtUp is set to true,
- * otherwise WalSndCaughtUp is set to false.
+ * Notify walproposer about the new WAL position.
  */
 static void
-XLogSendPhysical(void)
+XLogBroadcastWalProposer(void)
 {
-	XLogRecPtr	SendRqstPtr;
 	XLogRecPtr	startptr;
 	XLogRecPtr	endptr;
-	Size		nbytes PG_USED_FOR_ASSERTS_ONLY;
-	TimeLineID	currTLI;
 
-	/* If requested switch the WAL sender to the stopping state. */
-	if (got_STOPPING)
-		WalSndSetState(WALSNDSTATE_STOPPING);
+	/* Start from the last sent position */
+	startptr = sentPtr;
 
-	if (streamingDoneSending)
-	{
-		WalSndCaughtUp = true;
-		return;
-	}
-
-	/* Figure out how far we can safely send the WAL. */
-	if (sendTimeLineIsHistoric)
-	{
-		/*
-		 * Streaming an old timeline that's in this server's history, but is
-		 * not the one we're currently inserting or replaying. It can be
-		 * streamed up to the point where we switched off that timeline.
-		 */
-		SendRqstPtr = sendTimeLineValidUpto;
-	}
-	else if (am_cascading_walsender)
-	{
-		/*
-		 * Streaming the latest timeline on a standby.
-		 *
-		 * Attempt to send all WAL that has already been replayed, so that we
-		 * know it's valid. If we're receiving WAL through streaming
-		 * replication, it's also OK to send any WAL that has been received
-		 * but not replayed.
-		 *
-		 * The timeline we're recovering from can change, or we can be
-		 * promoted. In either case, the current timeline becomes historic. We
-		 * need to detect that so that we don't try to stream past the point
-		 * where we switched to another timeline. We check for promotion or
-		 * timeline switch after calculating FlushPtr, to avoid a race
-		 * condition: if the timeline becomes historic just after we checked
-		 * that it was still current, it's still be OK to stream it up to the
-		 * FlushPtr that was calculated before it became historic.
-		 */
-		bool		becameHistoric = false;
+	/*
+	 * Streaming the current timeline on a primary.
+	 *
+	 * Attempt to send all data that's already been written out and
+	 * fsync'd to disk.  We cannot go further than what's been written out
+	 * given the current implementation of WALRead().  And in any case
+	 * it's unsafe to send WAL that is not securely down to disk on the
+	 * primary: if the primary subsequently crashes and restarts, standbys
+	 * must not have applied any WAL that got lost on the primary.
+	 */
 #if PG_VERSION_NUM >= 150000
-		SendRqstPtr = GetStandbyFlushRecPtr(&currTLI);
+	endptr = GetFlushRecPtr(NULL);
 #else
-		SendRqstPtr = GetStandbyFlushRecPtr();
-		currTLI = ThisTimeLineID;
+	endptr = GetFlushRecPtr();
 #endif
-		if (!RecoveryInProgress())
-		{
-			/*
-			 * We have been promoted. RecoveryInProgress() updated
-			 * ThisTimeLineID to the new current timeline.
-			 */
-			am_cascading_walsender = false;
-			becameHistoric = true;
-		}
-		else
-		{
-			/*
-			 * Still a cascading standby. But is the timeline we're sending
-			 * still the one recovery is recovering from? currTLI was updated
-			 * by the GetStandbyFlushRecPtr() call above.
-			 */
-			if (sendTimeLine != currTLI)
-				becameHistoric = true;
-		}
-
-		if (becameHistoric)
-		{
-			/*
-			 * The timeline we were sending has become historic. Read the
-			 * timeline history file of the new timeline to see where exactly
-			 * we forked off from the timeline we were sending.
-			 */
-			List	   *history;
-
-			history = readTimeLineHistory(currTLI);
-			sendTimeLineValidUpto = tliSwitchPoint(sendTimeLine, history, &sendTimeLineNextTLI);
-
-			Assert(sendTimeLine < sendTimeLineNextTLI);
-			list_free_deep(history);
-
-			sendTimeLineIsHistoric = true;
-
-			SendRqstPtr = sendTimeLineValidUpto;
-		}
-	}
-	else
-	{
-		/*
-		 * Streaming the current timeline on a primary.
-		 *
-		 * Attempt to send all data that's already been written out and
-		 * fsync'd to disk.  We cannot go further than what's been written out
-		 * given the current implementation of WALRead().  And in any case
-		 * it's unsafe to send WAL that is not securely down to disk on the
-		 * primary: if the primary subsequently crashes and restarts, standbys
-		 * must not have applied any WAL that got lost on the primary.
-		 */
-#if PG_VERSION_NUM >= 150000
-		SendRqstPtr = GetFlushRecPtr(NULL);
-#else
-		SendRqstPtr = GetFlushRecPtr();
-#endif
-	}
 
 	/*
 	 * Record the current system time as an approximation of the time at which
@@ -1083,91 +625,14 @@ XLogSendPhysical(void)
 	 * that arbitrary LSN is eventually reported as written, flushed and
 	 * applied, so that it can measure the elapsed time.
 	 */
-	LagTrackerWrite(SendRqstPtr, GetCurrentTimestamp());
-
-	/*
-	 * If this is a historic timeline and we've reached the point where we
-	 * forked to the next timeline, stop streaming.
-	 *
-	 * Note: We might already have sent WAL > sendTimeLineValidUpto. The
-	 * startup process will normally replay all WAL that has been received
-	 * from the primary, before promoting, but if the WAL streaming is
-	 * terminated at a WAL page boundary, the valid portion of the timeline
-	 * might end in the middle of a WAL record. We might've already sent the
-	 * first half of that partial WAL record to the cascading standby, so that
-	 * sentPtr > sendTimeLineValidUpto. That's OK; the cascading standby can't
-	 * replay the partial WAL record either, so it can still follow our
-	 * timeline switch.
-	 */
-	if (sendTimeLineIsHistoric && sendTimeLineValidUpto <= sentPtr)
-	{
-		/* close the current file. */
-		if (xlogreader->seg.ws_file >= 0)
-			wal_segment_close(xlogreader);
-
-		/* Send CopyDone */
-		pq_putmessage_noblock('c', NULL, 0);
-		streamingDoneSending = true;
-
-		WalSndCaughtUp = true;
-
-		elog(DEBUG1, "walsender reached end of timeline at %X/%X (sent up to %X/%X)",
-			 LSN_FORMAT_ARGS(sendTimeLineValidUpto),
-			 LSN_FORMAT_ARGS(sentPtr));
-		return;
-	}
+	LagTrackerWrite(endptr, GetCurrentTimestamp());
 
 	/* Do we have any work to do? */
-	Assert(sentPtr <= SendRqstPtr);
-	if (SendRqstPtr <= sentPtr)
-	{
-		WalSndCaughtUp = true;
+	Assert(startptr <= endptr);
+	if (endptr <= startptr)
 		return;
-	}
 
-	/*
-	 * Figure out how much to send in one message. If there's no more than
-	 * MAX_SEND_SIZE bytes to send, send everything. Otherwise send
-	 * MAX_SEND_SIZE bytes, but round back to logfile or page boundary.
-	 *
-	 * The rounding is not only for performance reasons. Walreceiver relies on
-	 * the fact that we never split a WAL record across two messages. Since a
-	 * long WAL record is split at page boundary into continuation records,
-	 * page boundary is always a safe cut-off point. We also assume that
-	 * SendRqstPtr never points to the middle of a WAL record.
-	 */
-	startptr = sentPtr;
-	endptr = startptr;
-	endptr += MAX_SEND_SIZE;
-
-	/* if we went beyond SendRqstPtr, back off */
-	if (SendRqstPtr <= endptr)
-	{
-		endptr = SendRqstPtr;
-		if (sendTimeLineIsHistoric)
-			WalSndCaughtUp = false;
-		else
-			WalSndCaughtUp = true;
-	}
-	else
-	{
-		/* round down to page boundary. */
-		endptr -= (endptr % XLOG_BLCKSZ);
-		WalSndCaughtUp = false;
-	}
-
-	nbytes = endptr - startptr;
-	Assert(nbytes <= MAX_SEND_SIZE);
-
-	/* always true */
-	if (am_wal_proposer)
-	{
-		WalProposerBroadcast(startptr, endptr);
-	}
-	else
-	{
-		/* code removed for brevity */
-	}
+	WalProposerBroadcast(startptr, endptr);
 	sentPtr = endptr;
 
 	/* Update shared memory status */

safekeeper/src/handler.rs

@@ -2,8 +2,8 @@
 //! protocol commands.
 
 use anyhow::Context;
-use std::str;
 use std::str::FromStr;
+use std::str::{self};
 use tokio::io::{AsyncRead, AsyncWrite};
 use tracing::{info, info_span, Instrument};
 
@@ -11,6 +11,7 @@ use crate::auth::check_permission;
 use crate::json_ctrl::{handle_json_ctrl, AppendLogicalMessage};
 
 use crate::metrics::{TrafficMetrics, PG_QUERIES_FINISHED, PG_QUERIES_RECEIVED};
+use crate::safekeeper::Term;
 use crate::timeline::TimelineError;
 use crate::wal_service::ConnectionId;
 use crate::{GlobalTimelines, SafeKeeperConf};
@@ -44,7 +45,7 @@ pub struct SafekeeperPostgresHandler {
 /// Parsed Postgres command.
 enum SafekeeperPostgresCommand {
     StartWalPush,
-    StartReplication { start_lsn: Lsn },
+    StartReplication { start_lsn: Lsn, term: Option<Term> },
     IdentifySystem,
     TimelineStatus,
     JSONCtrl { cmd: AppendLogicalMessage },
@@ -55,15 +56,21 @@ fn parse_cmd(cmd: &str) -> anyhow::Result<SafekeeperPostgresCommand> {
         Ok(SafekeeperPostgresCommand::StartWalPush)
     } else if cmd.starts_with("START_REPLICATION") {
         let re = Regex::new(
-            r"START_REPLICATION(?: SLOT [^ ]+)?(?: PHYSICAL)? ([[:xdigit:]]+/[[:xdigit:]]+)",
+            // We follow postgres START_REPLICATION LOGICAL options to pass term.
+            r"START_REPLICATION(?: SLOT [^ ]+)?(?: PHYSICAL)? ([[:xdigit:]]+/[[:xdigit:]]+)(?: \(term='(\d+)'\))?",
         )
         .unwrap();
-        let mut caps = re.captures_iter(cmd);
-        let start_lsn = caps
-            .next()
-            .map(|cap| Lsn::from_str(&cap[1]))
-            .context("parse start LSN from START_REPLICATION command")??;
-        Ok(SafekeeperPostgresCommand::StartReplication { start_lsn })
+        let caps = re
+            .captures(cmd)
+            .context(format!("failed to parse START_REPLICATION command {}", cmd))?;
+        let start_lsn =
+            Lsn::from_str(&caps[1]).context("parse start LSN from START_REPLICATION command")?;
+        let term = if let Some(m) = caps.get(2) {
+            Some(m.as_str().parse::<u64>().context("invalid term")?)
+        } else {
+            None
+        };
+        Ok(SafekeeperPostgresCommand::StartReplication { start_lsn, term })
     } else if cmd.starts_with("IDENTIFY_SYSTEM") {
         Ok(SafekeeperPostgresCommand::IdentifySystem)
     } else if cmd.starts_with("TIMELINE_STATUS") {
@@ -218,8 +225,8 @@ impl<IO: AsyncRead + AsyncWrite + Unpin + Send> postgres_backend::Handler<IO>
                     .instrument(info_span!("WAL receiver", ttid = %span_ttid))
                     .await
             }
-            SafekeeperPostgresCommand::StartReplication { start_lsn } => {
-                self.handle_start_replication(pgb, start_lsn)
+            SafekeeperPostgresCommand::StartReplication { start_lsn, term } => {
+                self.handle_start_replication(pgb, start_lsn, term)
                     .instrument(info_span!("WAL sender", ttid = %span_ttid))
                     .await
             }

safekeeper/src/send_wal.rs

@@ -2,6 +2,7 @@
 //! with the "START_REPLICATION" message, and registry of walsenders.
 
 use crate::handler::SafekeeperPostgresHandler;
+use crate::safekeeper::Term;
 use crate::timeline::Timeline;
 use crate::wal_service::ConnectionId;
 use crate::wal_storage::WalReader;
@@ -359,8 +360,12 @@ impl SafekeeperPostgresHandler {
         &mut self,
         pgb: &mut PostgresBackend<IO>,
         start_pos: Lsn,
+        term: Option<Term>,
     ) -> Result<(), QueryError> {
-        if let Err(end) = self.handle_start_replication_guts(pgb, start_pos).await {
+        if let Err(end) = self
+            .handle_start_replication_guts(pgb, start_pos, term)
+            .await
+        {
             // Log the result and probably send it to the client, closing the stream.
             pgb.handle_copy_stream_end(end).await;
         }
@@ -371,6 +376,7 @@ impl SafekeeperPostgresHandler {
         &mut self,
         pgb: &mut PostgresBackend<IO>,
         start_pos: Lsn,
+        term: Option<Term>,
     ) -> Result<(), CopyStreamHandlerEnd> {
         let appname = self.appname.clone();
         let tli =
@@ -440,6 +446,7 @@ impl SafekeeperPostgresHandler {
             start_pos,
             end_pos,
             stop_pos,
+            term,
             commit_lsn_watch_rx,
             ws_guard: ws_guard.clone(),
             wal_reader,
@@ -476,6 +483,10 @@ struct WalSender<'a, IO> {
     // If present, terminate after reaching this position; used by walproposer
     // in recovery.
     stop_pos: Option<Lsn>,
+    /// When streaming uncommitted part, the term the client acts as the leader
+    /// in. Streaming is stopped if local term changes to a different (higher)
+    /// value.
+    term: Option<Term>,
     commit_lsn_watch_rx: Receiver<Lsn>,
     ws_guard: Arc<WalSenderGuard>,
     wal_reader: WalReader,
@@ -518,8 +529,18 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> WalSender<'_, IO> {
                 .0 as usize;
             send_size = min(send_size, self.send_buf.len());
             let send_buf = &mut self.send_buf[..send_size];
-            // read wal into buffer
-            send_size = self.wal_reader.read(send_buf).await?;
+            let send_size: usize;
+            {
+                // If uncommitted part is being pulled, check that the term is
+                // still the expected one.
+                let _term_guard = if let Some(t) = self.term {
+                    Some(self.tli.acquire_term(t).await?)
+                } else {
+                    None
+                };
+                // read wal into buffer
+                send_size = self.wal_reader.read(send_buf).await?
+            };
             let send_buf = &send_buf[..send_size];
 
             // and send it

safekeeper/src/timeline.rs

@@ -499,6 +499,19 @@ impl Timeline {
         false
     }
 
+    /// Ensure taht current term is t, erroring otherwise, and lock the state.
+    pub async fn acquire_term(&self, t: Term) -> Result<MutexGuard<SharedState>> {
+        let ss = self.write_shared_state().await;
+        if ss.sk.state.acceptor_state.term != t {
+            bail!(
+                "failed to acquire term {}, current term {}",
+                t,
+                ss.sk.state.acceptor_state.term
+            );
+        }
+        Ok(ss)
+    }
+
     /// Returns whether s3 offloading is required and sets current status as
     /// matching it.
     pub async fn wal_backup_attend(&self) -> bool {

test_runner/regress/data/extension_test/5670669815/v14/ext_index.json

@@ -1,12 +1,24 @@
 {
-    "public_extensions": [ "weighted_mean" ],
-    "library_index": { "weighted_mean": "weighted_mean" },
+    "public_extensions": [
+        "anon",
+        "pg_buffercache"
+    ],
+    "library_index": {
+        "anon": "anon",
+        "pg_buffercache": "pg_buffercache"
+    },
     "extension_data": {
-        "weighted_mean": {
+        "pg_buffercache": {
             "control_data": {
-                "weighted_mean.control": "# weighted_mean extension\ncomment = 'A weighted mean aggregate function extension'\ndefault_version = '1.0.1'\nmodule_pathname = '$libdir/weighted_mean'\nrelocatable = true\ntrusted = true\n"
+                "pg_buffercache.control": "# pg_buffercache extension \ncomment = 'examine the shared buffer cache' \ndefault_version = '1.3' \nmodule_pathname = '$libdir/pg_buffercache' \nrelocatable = true \ntrusted=true"
             },
-            "archive_path": "5670669815/v14/extensions/weighted_mean.tar.zst"
+            "archive_path": "5670669815/v14/extensions/pg_buffercache.tar.zst"
+        },
+        "anon": {
+            "control_data": {
+                "anon.control": "# PostgreSQL Anonymizer (anon) extension \ncomment = 'Data anonymization tools' \ndefault_version = '1.1.0' \ndirectory='extension/anon' \nrelocatable = false \nrequires = 'pgcrypto' \nsuperuser = false \nmodule_pathname = '$libdir/anon' \ntrusted = true \n"
+            },
+            "archive_path": "5670669815/v14/extensions/anon.tar.zst"
         }
     }
-}
+}

test_runner/regress/data/extension_test/5670669815/v15/ext_index.json

@@ -1,16 +1,17 @@
 {
     "public_extensions": [
-        "weighted_mean"
+        "anon"
     ],
     "library_index": {
-        "weighted_mean": "weighted_mean"
+        "anon": "anon"
     },
     "extension_data": {
-        "weighted_mean": {
+        "anon": {
             "control_data": {
-                "weighted_mean.control": "# weighted_mean extension\ncomment = 'A weighted mean aggregate function extension'\ndefault_version = '1.0.1'\nmodule_pathname = '$libdir/weighted_mean'\nrelocatable = true\ntrusted = true\n"
+                "anon.control": "# PostgreSQL Anonymizer (anon) extension \ncomment = 'Data anonymization tools' \ndefault_version = '1.1.0' \ndirectory='extension/anon' \nrelocatable = false \nrequires = 'pgcrypto' \nsuperuser = false \nmodule_pathname = '$libdir/anon' \ntrusted = true \n"
             },
-            "archive_path": "5670669815/v15/extensions/weighted_mean.tar.zst"
+            "archive_path": "5670669815/v15/extensions/anon.tar.zst"
         }
     }
 }
+

test_runner/regress/data/extension_test/README.md

@@ -1,51 +0,0 @@
-
-This is some data for testing remote extensions
-It was generated by this docker code
-Basically just make installing the extension.
-
-```
-#########################################################################################
-#
-# Layer "weighted-mean-pg-build"
-# compile weighted-mean extension
-# Note: Please do not include this in the compute image,
-# it is used for testing remote extension functionality
-# so the extension *must* not be part of the compute image already
-#
-#########################################################################################
-FROM build-deps AS weighted-mean-pg-build
-COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
-ENV PATH "/usr/local/pgsql/bin/:$PATH"
-RUN wget https://github.com/Kozea/weighted_mean/archive/master.tar.gz -O weighted_mean.tar.gz &&\
-    mkdir weighted_mean-src && cd weighted_mean-src && tar xvzf ../weighted_mean.tar.gz --strip-components=1 -C . && \
-    find /usr/local/pgsql -type f | sed 's|^/usr/local/pgsql/||' > /before.txt &&\
-    make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/weighted_mean.control && \
-    find /usr/local/pgsql -type f | sed 's|^/usr/local/pgsql/||' > /after.txt &&\
-    mkdir -p /extensions/weighted_mean && cp /usr/local/pgsql/share/extension/weighted_mean.control /extensions/weighted_mean && \
-    sort -o /before.txt /before.txt && sort -o /after.txt /after.txt && \
-    comm -13 /before.txt /after.txt | tar --directory=/usr/local/pgsql --zstd -cf /extensions/weighted_mean.tar.zst -T -
-```
-
-weighted_mean
-
-Copyright (c) Ronan Dunklau 2012-2015
-
-Permission to use, copy, modify, and distribute this software and its
-documentation for any purpose, without fee, and without a written agreement
-is hereby granted, provided that the above copyright notice and this
-paragraph and the following two paragraphs appear in all copies.
-
-IN NO EVENT SHALL RONAN DUNKLAU BE LIABLE TO ANY PARTY FOR
-DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, INCLUDING
-LOST PROFITS, ARISING OUT OF THE USE OF THIS SOFTWARE AND ITS
-DOCUMENTATION, EVEN IF RONAN DUNKLAU HAS BEEN ADVISED OF THE
-POSSIBILITY OF SUCH DAMAGE.
-
-RONAN DUNKLAU SPECIFICALLY DISCLAIMS ANY WARRANTIES,
-INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
-AND FITNESS FOR A PARTICULAR PURPOSE.  THE SOFTWARE PROVIDED HEREUNDER IS
-ON AN "AS IS" BASIS, AND RONAN DUNKLAU HAS NO OBLIGATIONS TO
-PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
-
-

test_runner/regress/test_download_extensions.py

@@ -1,23 +1,64 @@
 import os
 import shutil
-import uuid
+import threading
 from contextlib import closing
 from pathlib import Path
-from typing import Optional
 
 import pytest
-from fixtures.broker import NeonBroker
 from fixtures.log_helper import log
 from fixtures.neon_fixtures import (
     NeonEnvBuilder,
 )
 from fixtures.pg_version import PgVersion
-from fixtures.port_distributor import PortDistributor
-from fixtures.remote_storage import (
-    MockS3Server,
-    RemoteStorageKind,
-    available_s3_storages,
-)
+from fixtures.remote_storage import RemoteStorageKind, available_s3_storages
+
+
+# Cleaning up downloaded files is important for local tests
+# or else one test could reuse the files from another test or another test run
+def cleanup(pg_version):
+    PGDIR = Path(f"pg_install/v{pg_version}")
+
+    LIB_DIR = PGDIR / Path("lib/postgresql")
+    cleanup_lib_globs = ["anon*", "postgis*", "pg_buffercache*"]
+    cleanup_lib_glob_paths = [LIB_DIR.glob(x) for x in cleanup_lib_globs]
+
+    SHARE_DIR = PGDIR / Path("share/postgresql/extension")
+    cleanup_ext_globs = [
+        "anon*",
+        "address_standardizer*",
+        "postgis*",
+        "pageinspect*",
+        "pg_buffercache*",
+        "pgrouting*",
+    ]
+    cleanup_ext_glob_paths = [SHARE_DIR.glob(x) for x in cleanup_ext_globs]
+
+    all_glob_paths = cleanup_lib_glob_paths + cleanup_ext_glob_paths
+    all_cleanup_files = []
+    for file_glob in all_glob_paths:
+        for file in file_glob:
+            all_cleanup_files.append(file)
+
+    for file in all_cleanup_files:
+        try:
+            os.remove(file)
+            log.info(f"removed file {file}")
+        except Exception as err:
+            log.info(
+                f"skipping remove of file {file} because it doesn't exist.\
+                      this may be expected or unexpected depending on the test {err}"
+            )
+
+    cleanup_folders = [SHARE_DIR / Path("anon"), PGDIR / Path("download_extensions")]
+    for folder in cleanup_folders:
+        try:
+            shutil.rmtree(folder)
+            log.info(f"removed folder {folder}")
+        except Exception as err:
+            log.info(
+                f"skipping remove of folder {folder} because it doesn't exist.\
+                      this may be expected or unexpected depending on the test {err}"
+            )
 
 
 def upload_files(env):
@@ -39,194 +80,248 @@ def upload_files(env):
     os.chdir("../../../..")
 
 
-# creates the weighted_mean extension and runs test queries to verify it works
-def weighted_mean_test(endpoint, message: Optional[str] = None, create: bool = True):
-    if message is not None:
-        log.info(message)
-    with closing(endpoint.connect()) as conn:
-        with conn.cursor() as cur:
-            # Check that appropriate control files were downloaded
-            cur.execute("SELECT * FROM pg_available_extensions")
-            all_extensions = [x[0] for x in cur.fetchall()]
-            log.info(all_extensions)
-            assert "weighted_mean" in all_extensions
-
-            # test is from: https://github.com/Kozea/weighted_mean/tree/master/test
-            if create:
-                cur.execute("CREATE extension weighted_mean")
-            cur.execute(
-                "create temp table test as (\
-            select a::numeric, b::numeric\
-            from\
-                generate_series(1, 100) as a(a),\
-                generate_series(1, 100) as b(b));"
-            )
-            cur.execute("select weighted_mean(a,b) from test;")
-            x = cur.fetchone()
-            log.info(x)
-            assert str(x) == "(Decimal('50.5000000000000000'),)"
-            cur.execute("update test set b = 0;")
-            cur.execute("select weighted_mean(a,b) from test;")
-            x = cur.fetchone()
-            log.info(x)
-            assert str(x) == "(Decimal('0'),)"
-
-
 # Test downloading remote extension.
+@pytest.mark.parametrize("remote_storage_kind", available_s3_storages())
+@pytest.mark.skip(reason="https://github.com/neondatabase/neon/issues/4949")
 def test_remote_extensions(
+    neon_env_builder: NeonEnvBuilder,
+    remote_storage_kind: RemoteStorageKind,
     pg_version: PgVersion,
-    test_output_dir: Path,
-    port_distributor: PortDistributor,
-    mock_s3_server: MockS3Server,
-    neon_binpath: Path,
-    pg_distrib_dir: Path,
-    default_broker: NeonBroker,
-    run_id: uuid.UUID,
 ):
-    alt_pgdir = test_output_dir / "pg_install"
-    log.info(f"Copying {pg_distrib_dir} (which is big) to {alt_pgdir}")
-    shutil.copytree(pg_distrib_dir / pg_version.v_prefixed, alt_pgdir / pg_version.v_prefixed)
+    neon_env_builder.enable_remote_storage(
+        remote_storage_kind=remote_storage_kind,
+        test_name="test_remote_extensions",
+        enable_remote_extensions=True,
+    )
+    env = neon_env_builder.init_start()
+    tenant_id, _ = env.neon_cli.create_tenant()
+    env.neon_cli.create_timeline("test_remote_extensions", tenant_id=tenant_id)
 
-    with NeonEnvBuilder(
-        repo_dir=test_output_dir / "repo",
-        port_distributor=port_distributor,
-        mock_s3_server=mock_s3_server,
-        neon_binpath=neon_binpath,
-        pg_distrib_dir=alt_pgdir,
-        pg_version=pg_version,
-        broker=default_broker,
-        run_id=run_id,
-        preserve_database_files=False,
-    ) as neon_env_builder:
-        log.info(port_distributor)
-        log.info(mock_s3_server)
-        log.info(neon_binpath)
-        log.info(pg_distrib_dir)
-        log.info(pg_version)
-        log.info(default_broker)
-        log.info(run_id)
+    assert env.ext_remote_storage is not None  # satisfy mypy
+    assert env.remote_storage_client is not None  # satisfy mypy
 
-        neon_env_builder.enable_remote_storage(
-            remote_storage_kind=RemoteStorageKind.MOCK_S3,
-            test_name="test_remote_extensions",
-            enable_remote_extensions=True,
-        )
-        env = neon_env_builder.init_start()
-        tenant_id, _ = env.neon_cli.create_tenant()
-        env.neon_cli.create_timeline("test_remote_extensions", tenant_id=tenant_id)
-
-        assert env.ext_remote_storage is not None  # satisfy mypy
+    # For MOCK_S3 we upload test files.
+    # For REAL_S3 we use the files already in the bucket
+    if remote_storage_kind == RemoteStorageKind.MOCK_S3:
         upload_files(env)
 
-        # Start a compute node and check that it can download the extensions
-        # and use them to CREATE EXTENSION and LOAD
-        endpoint = env.endpoints.create_start(
-            "test_remote_extensions",
-            tenant_id=tenant_id,
-            remote_ext_config=env.ext_remote_storage.to_string(),
-            # config_lines=["log_min_messages=debug3"],
-        )
-        weighted_mean_test(endpoint)
+    # Start a compute node and check that it can download the extensions
+    # and use them to CREATE EXTENSION and LOAD
+    endpoint = env.endpoints.create_start(
+        "test_remote_extensions",
+        tenant_id=tenant_id,
+        remote_ext_config=env.ext_remote_storage.to_string(),
+        # config_lines=["log_min_messages=debug3"],
+    )
+    try:
+        with closing(endpoint.connect()) as conn:
+            with conn.cursor() as cur:
+                # Check that appropriate control files were downloaded
+                cur.execute("SELECT * FROM pg_available_extensions")
+                all_extensions = [x[0] for x in cur.fetchall()]
+                log.info(all_extensions)
+                assert "anon" in all_extensions
 
-        # # Test that extension is downloaded after endpoint restart,
-        # # when the library is used in the query.
-        # #
-        # # Run the test with mutliple simultaneous connections to an endpoint.
-        # # to ensure that the extension is downloaded only once.
+                # postgis is on real s3 but not mock s3.
+                # it's kind of a big file, would rather not upload to github
+                if remote_storage_kind == RemoteStorageKind.REAL_S3:
+                    assert "postgis" in all_extensions
+                    # this may fail locally if dependency is missing
+                    # we don't really care about the error,
+                    # we just want to make sure it downloaded
+                    try:
+                        cur.execute("CREATE EXTENSION postgis")
+                    except Exception as err:
+                        log.info(f"(expected) error creating postgis extension: {err}")
+                        # we do not check the error, so this is basically a NO-OP
+                        # however checking the log you can make sure that it worked
+                        # and also get valuable information about how long loading the extension took
 
-        # import subprocess
-        # import threading
-        # # shutdown compute node
-        # endpoint.stop()
-        # # remove extension files locally
-        # SHAREDIR = subprocess.check_output(
-        #     [alt_pgdir / f"{pg_version.v_prefixed}/bin/pg_config", "--sharedir"]
-        # )
-        # SHAREDIRPATH = Path(SHAREDIR.decode("utf-8").strip())
-        # log.info("SHAREDIRPATH: %s", SHAREDIRPATH)
-        # (SHAREDIRPATH / "extension/weighted_mean--1.0.1.sql").unlink()
-        # (SHAREDIRPATH / "extension/weighted_mean.control").unlink()
+                # this is expected to fail on my computer because I don't have the pgcrypto extension
+                try:
+                    cur.execute("CREATE EXTENSION anon")
+                except Exception as err:
+                    log.info("error creating anon extension")
+                    assert "pgcrypto" in str(err), "unexpected error creating anon extension"
+    finally:
+        cleanup(pg_version)
 
-        # # spin up compute node again (there are no extension files available, because compute is stateless)
-        # endpoint = env.endpoints.create_start(
-        #     "test_remote_extensions",
-        #     tenant_id=tenant_id,
-        #     remote_ext_config=env.ext_remote_storage.to_string(),
-        #     # config_lines=["log_min_messages=debug3"],
-        # )
 
-        # # connect to compute node and run the query
-        # # that will trigger the download of the extension
-        # threads = [
-        #     threading.Thread(
-        #         target=weighted_mean_test, args=(endpoint, f"this is thread {i}", False)
-        #     )
-        #     for i in range(2)
-        # ]
-        # for thread in threads:
-        #     thread.start()
-        # for thread in threads:
-        #     thread.join()
+# Test downloading remote library.
+@pytest.mark.parametrize("remote_storage_kind", available_s3_storages())
+@pytest.mark.skip(reason="https://github.com/neondatabase/neon/issues/4949")
+def test_remote_library(
+    neon_env_builder: NeonEnvBuilder,
+    remote_storage_kind: RemoteStorageKind,
+    pg_version: PgVersion,
+):
+    neon_env_builder.enable_remote_storage(
+        remote_storage_kind=remote_storage_kind,
+        test_name="test_remote_library",
+        enable_remote_extensions=True,
+    )
+    env = neon_env_builder.init_start()
+    tenant_id, _ = env.neon_cli.create_tenant()
+    env.neon_cli.create_timeline("test_remote_library", tenant_id=tenant_id)
+
+    assert env.ext_remote_storage is not None  # satisfy mypy
+    assert env.remote_storage_client is not None  # satisfy mypy
+
+    # For MOCK_S3 we upload test files.
+    # For REAL_S3 we use the files already in the bucket
+    if remote_storage_kind == RemoteStorageKind.MOCK_S3:
+        upload_files(env)
+
+    # and use them to run LOAD library
+    endpoint = env.endpoints.create_start(
+        "test_remote_library",
+        tenant_id=tenant_id,
+        remote_ext_config=env.ext_remote_storage.to_string(),
+        # config_lines=["log_min_messages=debug3"],
+    )
+    try:
+        with closing(endpoint.connect()) as conn:
+            with conn.cursor() as cur:
+                # try to load library
+                try:
+                    cur.execute("LOAD 'anon'")
+                except Exception as err:
+                    log.info(f"error loading anon library: {err}")
+                    raise AssertionError("unexpected error loading anon library") from err
+
+                # test library which name is different from extension name
+                # this may fail locally if dependency is missing
+                # however, it does successfully download the postgis archive
+                if remote_storage_kind == RemoteStorageKind.REAL_S3:
+                    try:
+                        cur.execute("LOAD 'postgis_topology-3'")
+                    except Exception as err:
+                        log.info("error loading postgis_topology-3")
+                        assert "No such file or directory" in str(
+                            err
+                        ), "unexpected error loading postgis_topology-3"
+    finally:
+        cleanup(pg_version)
 
 
 # Here we test a complex extension
 # which has multiple extensions in one archive
 # using postgis as an example
-@pytest.mark.skipif(
-    RemoteStorageKind.REAL_S3 not in available_s3_storages(),
-    reason="skipping test because real s3 not enabled",
-)
+# @pytest.mark.skipif(
+#    RemoteStorageKind.REAL_S3 not in available_s3_storages(),
+#    reason="skipping test because real s3 not enabled",
+# )
+@pytest.mark.skip(reason="https://github.com/neondatabase/neon/issues/4949")
 def test_multiple_extensions_one_archive(
+    neon_env_builder: NeonEnvBuilder,
     pg_version: PgVersion,
-    test_output_dir: Path,
-    port_distributor: PortDistributor,
-    mock_s3_server: MockS3Server,
-    neon_binpath: Path,
-    pg_distrib_dir: Path,
-    default_broker: NeonBroker,
-    run_id: uuid.UUID,
 ):
-    alt_pgdir = test_output_dir / "pg_install"
-    log.info(f"Copying {pg_distrib_dir} (which is big) to {alt_pgdir}.")
-    shutil.copytree(pg_distrib_dir / pg_version.v_prefixed, alt_pgdir / pg_version.v_prefixed)
+    neon_env_builder.enable_remote_storage(
+        remote_storage_kind=RemoteStorageKind.REAL_S3,
+        test_name="test_multiple_extensions_one_archive",
+        enable_remote_extensions=True,
+    )
+    env = neon_env_builder.init_start()
+    tenant_id, _ = env.neon_cli.create_tenant()
+    env.neon_cli.create_timeline("test_multiple_extensions_one_archive", tenant_id=tenant_id)
 
-    with NeonEnvBuilder(
-        repo_dir=test_output_dir / "repo",
-        port_distributor=port_distributor,
-        mock_s3_server=mock_s3_server,
-        neon_binpath=neon_binpath,
-        pg_distrib_dir=alt_pgdir,
-        pg_version=pg_version,
-        broker=default_broker,
-        run_id=run_id,
-        preserve_database_files=False,
-    ) as neon_env_builder:
-        neon_env_builder.enable_remote_storage(
-            remote_storage_kind=RemoteStorageKind.REAL_S3,
-            test_name="test_multiple_extensions_one_archive",
-            enable_remote_extensions=True,
-        )
-        env = neon_env_builder.init_start()
-        tenant_id, _ = env.neon_cli.create_tenant()
-        env.neon_cli.create_timeline("test_multiple_extensions_one_archive", tenant_id=tenant_id)
+    assert env.ext_remote_storage is not None  # satisfy mypy
+    assert env.remote_storage_client is not None  # satisfy mypy
 
-        assert env.ext_remote_storage is not None  # satisfy mypy
+    endpoint = env.endpoints.create_start(
+        "test_multiple_extensions_one_archive",
+        tenant_id=tenant_id,
+        remote_ext_config=env.ext_remote_storage.to_string(),
+    )
+    with closing(endpoint.connect()) as conn:
+        with conn.cursor() as cur:
+            cur.execute("CREATE EXTENSION address_standardizer;")
+            cur.execute("CREATE EXTENSION address_standardizer_data_us;")
+            # execute query to ensure that it works
+            cur.execute(
+                "SELECT house_num, name, suftype, city, country, state, unit \
+                        FROM standardize_address('us_lex', 'us_gaz', 'us_rules', \
+                        'One Rust Place, Boston, MA 02109');"
+            )
+            res = cur.fetchall()
+            log.info(res)
+            assert len(res) > 0
 
-        endpoint = env.endpoints.create_start(
-            "test_multiple_extensions_one_archive",
-            tenant_id=tenant_id,
-            remote_ext_config=env.ext_remote_storage.to_string(),
-        )
+    cleanup(pg_version)
+
+
+# Test that extension is downloaded after endpoint restart,
+# when the library is used in the query.
+#
+# Run the test with mutliple simultaneous connections to an endpoint.
+# to ensure that the extension is downloaded only once.
+#
+@pytest.mark.skip(reason="https://github.com/neondatabase/neon/issues/4949")
+def test_extension_download_after_restart(
+    neon_env_builder: NeonEnvBuilder,
+    pg_version: PgVersion,
+):
+    if "15" in pg_version:  # SKIP v15 for now because test set only has extension built for v14
+        return None
+
+    neon_env_builder.enable_remote_storage(
+        remote_storage_kind=RemoteStorageKind.MOCK_S3,
+        test_name="test_extension_download_after_restart",
+        enable_remote_extensions=True,
+    )
+    env = neon_env_builder.init_start()
+    tenant_id, _ = env.neon_cli.create_tenant()
+    env.neon_cli.create_timeline("test_extension_download_after_restart", tenant_id=tenant_id)
+
+    assert env.ext_remote_storage is not None  # satisfy mypy
+    assert env.remote_storage_client is not None  # satisfy mypy
+
+    # For MOCK_S3 we upload test files.
+    upload_files(env)
+
+    endpoint = env.endpoints.create_start(
+        "test_extension_download_after_restart",
+        tenant_id=tenant_id,
+        remote_ext_config=env.ext_remote_storage.to_string(),
+        config_lines=["log_min_messages=debug3"],
+    )
+    with closing(endpoint.connect()) as conn:
+        with conn.cursor() as cur:
+            cur.execute("CREATE extension pg_buffercache;")
+            cur.execute("SELECT * from pg_buffercache;")
+            res = cur.fetchall()
+            assert len(res) > 0
+            log.info(res)
+
+    # shutdown compute node
+    endpoint.stop()
+    # remove extension files locally
+    cleanup(pg_version)
+
+    # spin up compute node again (there are no extension files available, because compute is stateless)
+    endpoint = env.endpoints.create_start(
+        "test_extension_download_after_restart",
+        tenant_id=tenant_id,
+        remote_ext_config=env.ext_remote_storage.to_string(),
+        config_lines=["log_min_messages=debug3"],
+    )
+
+    # connect to compute node and run the query
+    # that will trigger the download of the extension
+    def run_query(endpoint, thread_id: int):
+        log.info("thread_id {%d} starting", thread_id)
         with closing(endpoint.connect()) as conn:
             with conn.cursor() as cur:
-                cur.execute("CREATE EXTENSION address_standardizer;")
-                cur.execute("CREATE EXTENSION address_standardizer_data_us;")
-                # execute query to ensure that it works
-                cur.execute(
-                    "SELECT house_num, name, suftype, city, country, state, unit \
-                            FROM standardize_address('us_lex', 'us_gaz', 'us_rules', \
-                            'One Rust Place, Boston, MA 02109');"
-                )
+                cur.execute("SELECT * from pg_buffercache;")
                 res = cur.fetchall()
-                log.info(res)
                 assert len(res) > 0
+                log.info("thread_id {%d}, res = %s", thread_id, res)
+
+    threads = [threading.Thread(target=run_query, args=(endpoint, i)) for i in range(2)]
+
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+
+    cleanup(pg_version)

test_runner/regress/test_wal_acceptor.py

@@ -869,6 +869,49 @@ def test_timeline_status(neon_env_builder: NeonEnvBuilder, auth_enabled: bool):
     assert debug_dump_1["config"]["id"] == env.safekeepers[0].id
 
 
+class DummyConsumer(object):
+    def __call__(self, msg):
+        pass
+
+
+def test_start_replication_term(neon_env_builder: NeonEnvBuilder):
+    """
+    Test START_REPLICATION of uncommitted part specifying leader term. It must
+    error if safekeeper switched to different term.
+    """
+
+    env = neon_env_builder.init_start()
+
+    env.neon_cli.create_branch("test_start_replication_term")
+    endpoint = env.endpoints.create_start("test_start_replication_term")
+
+    endpoint.safe_psql("CREATE TABLE t(key int primary key, value text)")
+
+    # learn neon timeline from compute
+    tenant_id = TenantId(endpoint.safe_psql("show neon.tenant_id")[0][0])
+    timeline_id = TimelineId(endpoint.safe_psql("show neon.timeline_id")[0][0])
+
+    sk = env.safekeepers[0]
+    sk_http_cli = sk.http_client()
+    tli_status = sk_http_cli.timeline_status(tenant_id, timeline_id)
+    timeline_start_lsn = tli_status.timeline_start_lsn
+
+    conn_opts = {
+        "host": "127.0.0.1",
+        "options": f"-c timeline_id={timeline_id} tenant_id={tenant_id}",
+        "port": sk.port.pg,
+        "connection_factory": psycopg2.extras.PhysicalReplicationConnection,
+    }
+    sk_pg_conn = psycopg2.connect(**conn_opts)  # type: ignore
+    with sk_pg_conn.cursor() as cur:
+        # should fail, as first start has term 2
+        cur.start_replication_expert(f"START_REPLICATION {timeline_start_lsn} (term='3')")
+        dummy_consumer = DummyConsumer()
+        with pytest.raises(psycopg2.errors.InternalError_) as excinfo:
+            cur.consume_stream(dummy_consumer)
+        assert "failed to acquire term 3" in str(excinfo.value)
+
+
 # Test auth on WAL service (postgres protocol) ports.
 def test_sk_auth(neon_env_builder: NeonEnvBuilder):
     neon_env_builder.auth_enabled = True