Fix

Cargo.toml

@@ -1,3 +1,14 @@
+# 'named-profiles' feature was stabilized in cargo 1.57. This line makes the
+# build work with older cargo versions.
+#
+# We have this because as of this writing, the latest cargo Debian package
+# that's available is 1.56. (Confusingly, the Debian package version number
+# is 0.57, whereas 'cargo --version' says 1.56.)
+#
+# See https://tracker.debian.org/pkg/cargo for the current status of the
+# package. When that gets updated, we can remove this.
+cargo-features = ["named-profiles"]
+
 [workspace]
 members = [
     "compute_tools",

Dockerfile.compute-node-v14

@@ -29,12 +29,7 @@ RUN cd postgres && \
     make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C contrib/ install && \
     # Install headers
     make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/include install && \
-    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/interfaces/libpq install && \
-    # Enable some of contrib extensions
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/bloom.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgrowlocks.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/intagg.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgstattuple.control
+    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/interfaces/libpq install
 
 #########################################################################################
 #
@@ -60,10 +55,7 @@ RUN wget https://download.osgeo.org/postgis/source/postgis-3.3.1.tar.gz && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_raster.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_tiger_geocoder.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_topology.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer_data_us.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_sfcgal.control
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_topology.control
 
 #########################################################################################
 #

Dockerfile.compute-node-v15

@@ -29,12 +29,7 @@ RUN cd postgres && \
     make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C contrib/ install && \
     # Install headers
     make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/include install && \
-    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/interfaces/libpq install && \
-    # Enable some of contrib extensions
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/bloom.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgrowlocks.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/intagg.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgstattuple.control
+    make MAKELEVEL=0 -j $(getconf _NPROCESSORS_ONLN) -s -C src/interfaces/libpq install
 
 #########################################################################################
 #
@@ -60,10 +55,7 @@ RUN wget https://download.osgeo.org/postgis/source/postgis-3.3.1.tar.gz && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_raster.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_tiger_geocoder.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_topology.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer_data_us.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_sfcgal.control
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_topology.control
 
 #########################################################################################
 #

pageserver/Cargo.toml

@@ -67,7 +67,9 @@ 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" }
+rpds = "0.12.0"
 reqwest = { version = "0.11", default-features = false, features = ["rustls-tls"] }
+im = "15.1.0"
 
 [dev-dependencies]
 criterion = "0.4"

pageserver/benches/bench_layer_map.rs

@@ -1,13 +1,12 @@
-use anyhow::Result;
+use pageserver::keyspace::{KeyPartitioning, KeySpace};
 use pageserver::repository::Key;
 use pageserver::tenant::layer_map::LayerMap;
-use pageserver::tenant::storage_layer::{DeltaFileName, ImageFileName, ValueReconstructState};
-use pageserver::tenant::storage_layer::{Layer, ValueReconstructResult};
+use pageserver::tenant::storage_layer::Layer;
+use pageserver::tenant::storage_layer::{DeltaFileName, ImageFileName, LayerDescriptor};
 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;
@@ -17,80 +16,8 @@ use utils::lsn::Lsn;
 
 use criterion::{criterion_group, criterion_main, Criterion};
 
-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();
+fn build_layer_map(filename_dump: PathBuf) -> LayerMap<LayerDescriptor> {
+    let mut layer_map = LayerMap::<LayerDescriptor>::default();
 
     let mut min_lsn = Lsn(u64::MAX);
     let mut max_lsn = Lsn(0);
@@ -100,17 +27,21 @@ fn build_layer_map(filename_dump: PathBuf) -> LayerMap<dyn Layer> {
     for fname in filenames {
         let fname = &fname.unwrap();
         if let Some(imgfilename) = ImageFileName::parse_str(fname) {
-            let layer = DummyImage {
-                key_range: imgfilename.key_range,
-                lsn: imgfilename.lsn,
+            let layer = LayerDescriptor {
+                key: imgfilename.key_range,
+                lsn: imgfilename.lsn..(imgfilename.lsn + 1),
+                is_incremental: false,
+                short_id: fname.to_string(),
             };
             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 = DummyDelta {
-                key_range: deltafilename.key_range,
-                lsn_range: deltafilename.lsn_range.clone(),
+            let layer = LayerDescriptor {
+                key: deltafilename.key_range.clone(),
+                lsn: deltafilename.lsn_range.clone(),
+                is_incremental: true,
+                short_id: fname.to_string(),
             };
             layer_map.insert_historic(Arc::new(layer));
             min_lsn = min(min_lsn, deltafilename.lsn_range.start);
@@ -119,6 +50,7 @@ fn build_layer_map(filename_dump: PathBuf) -> LayerMap<dyn Layer> {
             panic!("unexpected filename {fname}");
         }
     }
+    layer_map.rebuild_index();
 
     println!("min: {min_lsn}, max: {max_lsn}");
 
@@ -126,7 +58,7 @@ fn build_layer_map(filename_dump: PathBuf) -> LayerMap<dyn Layer> {
 }
 
 /// Construct a layer map query pattern for benchmarks
-fn uniform_query_pattern(layer_map: &LayerMap<dyn Layer>) -> Vec<(Key, Lsn)> {
+fn uniform_query_pattern(layer_map: &LayerMap<LayerDescriptor>) -> 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
@@ -150,6 +82,41 @@ fn uniform_query_pattern(layer_map: &LayerMap<dyn Layer>) -> 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.
@@ -183,24 +150,71 @@ 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) {
-    // TODO consider compressing this file
+    // Init layer map
+    let now = Instant::now();
     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);
 
-    // Test with uniform query pattern
-    c.bench_function("real_map_uniform_queries", |b| {
+    // 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| {
         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
@@ -208,39 +222,38 @@ 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 = DummyImage {
-            key_range: zero.add(10 * i32)..zero.add(10 * i32 + 1),
-            lsn: Lsn(10 * i),
+        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),
         };
         layer_map.insert_historic(Arc::new(layer));
     }
-
-    // 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());
+    layer_map.rebuild_index();
+    println!("Finished layer map 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, 1)
+        .choose_multiple(rng, 100)
         .copied()
         .collect();
 
     // Define and name the benchmark function
-    c.bench_function("sequential_uniform_queries", |b| {
-        // Run the search queries
+    let mut group = c.benchmark_group("sequential");
+    group.bench_function("uniform_queries", |b| {
         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/basebackup.rs

@@ -10,7 +10,7 @@
 //! This module is responsible for creation of such tarball
 //! from data stored in object storage.
 //!
-use anyhow::{anyhow, ensure, Context, Result};
+use anyhow::{anyhow, bail, ensure, Context};
 use bytes::{BufMut, BytesMut};
 use fail::fail_point;
 use std::fmt::Write as FmtWrite;
@@ -27,8 +27,7 @@ use tracing::*;
 ///
 use tokio_tar::{Builder, EntryType, Header};
 
-use crate::tenant::TimelineRequestContext;
-use crate::tenant::{PageReconstructError, Timeline};
+use crate::tenant::{with_ondemand_download, Timeline};
 use pageserver_api::reltag::{RelTag, SlruKind};
 
 use postgres_ffi::pg_constants::{DEFAULTTABLESPACE_OID, GLOBALTABLESPACE_OID};
@@ -53,8 +52,7 @@ pub async fn send_basebackup_tarball<'a, W>(
     req_lsn: Option<Lsn>,
     prev_lsn: Option<Lsn>,
     full_backup: bool,
-    ctx: &'a TimelineRequestContext,
-) -> Result<(), PageReconstructError>
+) -> anyhow::Result<()>
 where
     W: AsyncWrite + Send + Sync + Unpin,
 {
@@ -93,10 +91,8 @@ where
 
     // Consolidate the derived and the provided prev_lsn values
     let prev_lsn = if let Some(provided_prev_lsn) = prev_lsn {
-        if backup_prev != Lsn(0) && backup_prev != provided_prev_lsn {
-            return Err(PageReconstructError::Other(anyhow!(
-                "prev LSN doesn't match"
-            )));
+        if backup_prev != Lsn(0) {
+            ensure!(backup_prev == provided_prev_lsn);
         }
         provided_prev_lsn
     } else {
@@ -114,7 +110,6 @@ where
         lsn: backup_lsn,
         prev_record_lsn: prev_lsn,
         full_backup,
-        ctx,
     };
     basebackup
         .send_tarball()
@@ -134,14 +129,13 @@ where
     lsn: Lsn,
     prev_record_lsn: Lsn,
     full_backup: bool,
-    ctx: &'a TimelineRequestContext,
 }
 
 impl<'a, W> Basebackup<'a, W>
 where
     W: AsyncWrite + Send + Sync + Unpin,
 {
-    async fn send_tarball(mut self) -> Result<(), PageReconstructError> {
+    async fn send_tarball(mut self) -> anyhow::Result<()> {
         // TODO include checksum
 
         // Create pgdata subdirs structure
@@ -177,10 +171,8 @@ where
             SlruKind::MultiXactOffsets,
             SlruKind::MultiXactMembers,
         ] {
-            for segno in self
-                .timeline
-                .list_slru_segments(kind, self.lsn, self.ctx)
-                .await?
+            for segno in
+                with_ondemand_download(|| self.timeline.list_slru_segments(kind, self.lsn)).await?
             {
                 self.add_slru_segment(kind, segno).await?;
             }
@@ -188,56 +180,44 @@ where
 
         // Create tablespace directories
         for ((spcnode, dbnode), has_relmap_file) in
-            self.timeline.list_dbdirs(self.lsn, self.ctx).await?
+            with_ondemand_download(|| self.timeline.list_dbdirs(self.lsn)).await?
         {
             self.add_dbdir(spcnode, dbnode, has_relmap_file).await?;
 
             // Gather and send relational files in each database if full backup is requested.
             if self.full_backup {
-                for rel in self
-                    .timeline
-                    .list_rels(spcnode, dbnode, self.lsn, self.ctx)
-                    .await?
+                for rel in
+                    with_ondemand_download(|| self.timeline.list_rels(spcnode, dbnode, self.lsn))
+                        .await?
                 {
                     self.add_rel(rel).await?;
                 }
             }
         }
-        for xid in self
-            .timeline
-            .list_twophase_files(self.lsn, self.ctx)
-            .await?
-        {
+        for xid in with_ondemand_download(|| self.timeline.list_twophase_files(self.lsn)).await? {
             self.add_twophase_file(xid).await?;
         }
 
         fail_point!("basebackup-before-control-file", |_| {
-            Err(PageReconstructError::from(anyhow!(
-                "failpoint basebackup-before-control-file"
-            )))
+            bail!("failpoint basebackup-before-control-file")
         });
 
         // Generate pg_control and bootstrap WAL segment.
         self.add_pgcontrol_file().await?;
-        self.ar.finish().await.context("could not finish tarball")?;
+        self.ar.finish().await?;
         debug!("all tarred up!");
         Ok(())
     }
 
-    async fn add_rel(&mut self, tag: RelTag) -> Result<(), PageReconstructError> {
-        let nblocks = self
-            .timeline
-            .get_rel_size(tag, self.lsn, false, self.ctx)
-            .await?;
+    async fn add_rel(&mut self, tag: RelTag) -> anyhow::Result<()> {
+        let nblocks =
+            with_ondemand_download(|| self.timeline.get_rel_size(tag, self.lsn, false)).await?;
 
         // If the relation is empty, create an empty file
         if nblocks == 0 {
             let file_name = tag.to_segfile_name(0);
             let header = new_tar_header(&file_name, 0)?;
-            self.ar
-                .append(&header, &mut io::empty())
-                .await
-                .context("could not write empty relfile to tar stream")?;
+            self.ar.append(&header, &mut io::empty()).await?;
             return Ok(());
         }
 
@@ -246,21 +226,20 @@ where
         let mut seg = 0;
         while startblk < nblocks {
             let endblk = std::cmp::min(startblk + RELSEG_SIZE, nblocks);
+
             let mut segment_data: Vec<u8> = vec![];
             for blknum in startblk..endblk {
-                let img = self
-                    .timeline
-                    .get_rel_page_at_lsn(tag, blknum, self.lsn, false, self.ctx)
-                    .await?;
+                let img = with_ondemand_download(|| {
+                    self.timeline
+                        .get_rel_page_at_lsn(tag, blknum, self.lsn, false)
+                })
+                .await?;
                 segment_data.extend_from_slice(&img[..]);
             }
 
             let file_name = tag.to_segfile_name(seg as u32);
             let header = new_tar_header(&file_name, segment_data.len() as u64)?;
-            self.ar
-                .append(&header, segment_data.as_slice())
-                .await
-                .context("could not write relfile segment to tar stream")?;
+            self.ar.append(&header, segment_data.as_slice()).await?;
 
             seg += 1;
             startblk = endblk;
@@ -273,17 +252,17 @@ where
     // Generate SLRU segment files from repository.
     //
     async fn add_slru_segment(&mut self, slru: SlruKind, segno: u32) -> anyhow::Result<()> {
-        let nblocks = self
-            .timeline
-            .get_slru_segment_size(slru, segno, self.lsn, self.ctx)
-            .await?;
+        let nblocks =
+            with_ondemand_download(|| self.timeline.get_slru_segment_size(slru, segno, self.lsn))
+                .await?;
 
         let mut slru_buf: Vec<u8> = Vec::with_capacity(nblocks as usize * BLCKSZ as usize);
         for blknum in 0..nblocks {
-            let img = self
-                .timeline
-                .get_slru_page_at_lsn(slru, segno, blknum, self.lsn, self.ctx)
-                .await?;
+            let img = with_ondemand_download(|| {
+                self.timeline
+                    .get_slru_page_at_lsn(slru, segno, blknum, self.lsn)
+            })
+            .await?;
 
             if slru == SlruKind::Clog {
                 ensure!(img.len() == BLCKSZ as usize || img.len() == BLCKSZ as usize + 8);
@@ -315,10 +294,9 @@ where
         has_relmap_file: bool,
     ) -> anyhow::Result<()> {
         let relmap_img = if has_relmap_file {
-            let img = self
-                .timeline
-                .get_relmap_file(spcnode, dbnode, self.lsn, self.ctx)
-                .await?;
+            let img =
+                with_ondemand_download(|| self.timeline.get_relmap_file(spcnode, dbnode, self.lsn))
+                    .await?;
             ensure!(img.len() == 512);
             Some(img)
         } else {
@@ -351,9 +329,7 @@ where
             // XLOG_TBLSPC_DROP records. But we probably should just
             // throw an error on CREATE TABLESPACE in the first place.
             if !has_relmap_file
-                && self
-                    .timeline
-                    .list_rels(spcnode, dbnode, self.lsn, self.ctx)
+                && with_ondemand_download(|| self.timeline.list_rels(spcnode, dbnode, self.lsn))
                     .await?
                     .is_empty()
             {
@@ -386,10 +362,7 @@ where
     // Extract twophase state files
     //
     async fn add_twophase_file(&mut self, xid: TransactionId) -> anyhow::Result<()> {
-        let img = self
-            .timeline
-            .get_twophase_file(xid, self.lsn, self.ctx)
-            .await?;
+        let img = with_ondemand_download(|| self.timeline.get_twophase_file(xid, self.lsn)).await?;
 
         let mut buf = BytesMut::new();
         buf.extend_from_slice(&img[..]);
@@ -425,14 +398,10 @@ where
             )
             .await?;
 
-        let checkpoint_bytes = self
-            .timeline
-            .get_checkpoint(self.lsn, self.ctx)
+        let checkpoint_bytes = with_ondemand_download(|| self.timeline.get_checkpoint(self.lsn))
             .await
             .context("failed to get checkpoint bytes")?;
-        let pg_control_bytes = self
-            .timeline
-            .get_control_file(self.lsn, self.ctx)
+        let pg_control_bytes = with_ondemand_download(|| self.timeline.get_control_file(self.lsn))
             .await
             .context("failed get control bytes")?;
 

pageserver/src/bin/pageserver.rs

@@ -13,8 +13,8 @@ use tracing::*;
 use metrics::set_build_info_metric;
 use pageserver::{
     config::{defaults::*, PageServerConf},
-    context::{DownloadBehavior, RequestContext, TaskKind},
     http, page_cache, page_service, task_mgr,
+    task_mgr::TaskKind,
     task_mgr::{
         BACKGROUND_RUNTIME, COMPUTE_REQUEST_RUNTIME, MGMT_REQUEST_RUNTIME, WALRECEIVER_RUNTIME,
     },
@@ -303,79 +303,62 @@ fn start_pageserver(conf: &'static PageServerConf) -> anyhow::Result<()> {
     {
         let _rt_guard = MGMT_REQUEST_RUNTIME.enter();
 
-        let mgmt_ctx = RequestContext::new(TaskKind::HttpEndpointListener, DownloadBehavior::Error);
-        let cancellation_token = Box::leak(Box::new(mgmt_ctx.cancellation_token().clone()));
         let router = http::make_router(conf, auth.clone(), remote_storage)?
             .build()
             .map_err(|err| anyhow!(err))?;
         let service = utils::http::RouterService::new(router).unwrap();
         let server = hyper::Server::from_tcp(http_listener)?
             .serve(service)
-            .with_graceful_shutdown(cancellation_token.cancelled());
+            .with_graceful_shutdown(task_mgr::shutdown_watcher());
 
         task_mgr::spawn(
             MGMT_REQUEST_RUNTIME.handle(),
+            TaskKind::HttpEndpointListener,
+            None,
+            None,
             "http endpoint listener",
             true,
             async {
-                match server.await {
-                    Ok(()) => info!("HTTP endpoint listener shut down"),
-                    Err(err) => error!("HTTP endpoint listener shut down with error: {err:?}"),
-                }
+                server.await?;
+                Ok(())
             },
         );
-    }
 
-    if let Some(metric_collection_endpoint) = &conf.metric_collection_endpoint {
-        let metrics_ctx = RequestContext::new(
-            TaskKind::MetricsCollection,
-            DownloadBehavior::Error, // metrics collector shouldn't be downloading anything
-        );
-        task_mgr::spawn(
-            MGMT_REQUEST_RUNTIME.handle(),
-            "consumption metrics collection",
-            true,
-            pageserver::consumption_metrics::collect_metrics(
-                metric_collection_endpoint,
-                conf.metric_collection_interval,
-                conf.id,
-                metrics_ctx,
-            )
-            .instrument(info_span!("metrics_collection")),
-        );
+        if let Some(metric_collection_endpoint) = &conf.metric_collection_endpoint {
+            task_mgr::spawn(
+                MGMT_REQUEST_RUNTIME.handle(),
+                TaskKind::MetricsCollection,
+                None,
+                None,
+                "consumption metrics collection",
+                true,
+                async move {
+                    pageserver::consumption_metrics::collect_metrics(
+                        metric_collection_endpoint,
+                        conf.metric_collection_interval,
+                        conf.id,
+                    )
+                    .instrument(info_span!("metrics_collection"))
+                    .await?;
+                    Ok(())
+                },
+            );
+        }
     }
 
     // Spawn a task to listen for libpq connections. It will spawn further tasks
     // for each connection. We created the listener earlier already.
-    {
-        let libpq_ctx = RequestContext::new(
-            TaskKind::LibpqEndpointListener,
-            // listener task shouldn't need to download anything. (We will
-            // create a separate sub-contexts for each connection, with their
-            // own download behavior. This context is used only to listen and
-            // accept connections.)
-            DownloadBehavior::Error,
-        );
-        task_mgr::spawn(
-            COMPUTE_REQUEST_RUNTIME.handle(),
-            "libpq endpoint listener",
-            true,
-            async move {
-                match page_service::libpq_listener_main(
-                    conf,
-                    auth,
-                    pageserver_listener,
-                    conf.auth_type,
-                    libpq_ctx,
-                )
-                .await
-                {
-                    Ok(()) => info!("libpq endpoint listener shut down"),
-                    Err(err) => error!("libpq endpoint listener shut down with error: {err:?}"),
-                }
-            },
-        );
-    }
+    task_mgr::spawn(
+        COMPUTE_REQUEST_RUNTIME.handle(),
+        TaskKind::LibpqEndpointListener,
+        None,
+        None,
+        "libpq endpoint listener",
+        true,
+        async move {
+            page_service::libpq_listener_main(conf, auth, pageserver_listener, conf.auth_type).await
+        },
+    );
 
     // All started up! Now just sit and wait for shutdown signal.
     signals.handle(|signal| match signal {
@@ -392,7 +375,7 @@ fn start_pageserver(conf: &'static PageServerConf) -> anyhow::Result<()> {
                 "Got {}. Terminating gracefully in fast shutdown mode",
                 signal.name()
             );
-            BACKGROUND_RUNTIME.block_on(task_mgr::shutdown_pageserver(0));
+            BACKGROUND_RUNTIME.block_on(pageserver::shutdown_pageserver(0));
             unreachable!()
         }
     })

pageserver/src/consumption_metrics.rs

@@ -9,8 +9,9 @@ use tracing::*;
 use utils::id::NodeId;
 use utils::id::TimelineId;
 
-use crate::context::RequestContext;
+use crate::task_mgr;
 use crate::tenant::mgr;
+use pageserver_api::models::TenantState;
 use utils::id::TenantId;
 
 use serde::{Deserialize, Serialize};
@@ -137,13 +138,12 @@ struct EventChunk<'a> {
     events: &'a [ConsumptionMetric],
 }
 
-/// Main task that serves metrics collection
+/// Main thread that serves metrics collection
 pub async fn collect_metrics(
     metric_collection_endpoint: &Url,
     metric_collection_interval: Duration,
     node_id: NodeId,
-    metrics_ctx: RequestContext,
-) {
+) -> anyhow::Result<()> {
     let mut ticker = tokio::time::interval(metric_collection_interval);
 
     info!("starting collect_metrics");
@@ -154,15 +154,12 @@ pub async fn collect_metrics(
 
     loop {
         tokio::select! {
-            _ = metrics_ctx.cancelled() => {
+            _ = task_mgr::shutdown_watcher() => {
                 info!("collect_metrics received cancellation request");
-                return;
+                return Ok(());
             },
             _ = ticker.tick() => {
-                if let Err(err) = collect_metrics_task(&client, &mut cached_metrics, metric_collection_endpoint, node_id, &metrics_ctx).await {
-                    // Log the error and continue
-                    error!("metrics collection failed: {err:?}");
-                }
+                collect_metrics_task(&client, &mut cached_metrics, metric_collection_endpoint, node_id).await?;
             }
         }
     }
@@ -177,7 +174,6 @@ pub async fn collect_metrics_task(
     cached_metrics: &mut HashMap<ConsumptionMetricsKey, u64>,
     metric_collection_endpoint: &reqwest::Url,
     node_id: NodeId,
-    ctx: &RequestContext,
 ) -> anyhow::Result<()> {
     let mut current_metrics: Vec<(ConsumptionMetricsKey, u64)> = Vec::new();
     trace!(
@@ -190,27 +186,18 @@ pub async fn collect_metrics_task(
 
     // iterate through list of Active tenants and collect metrics
     for (tenant_id, tenant_state) in tenants {
-        if ctx.is_cancelled() {
+        if tenant_state != TenantState::Active {
             continue;
         }
-        let tenant = mgr::get_tenant(tenant_id).await?;
-        // If the tenant was shut down while while we were looking elsewhere, skip it.
-        let tenant_ctx = match tenant.get_context(ctx) {
-            Ok(ctx) => ctx,
-            Err(_state) => {
-                debug!(
-                    "skipping metrics collection for tenant {tenant_id} because it is not active"
-                );
-                continue;
-            }
-        };
+
+        let tenant = mgr::get_tenant(tenant_id, true).await?;
 
         let mut tenant_resident_size = 0;
 
         // iterate through list of timelines in tenant
         for timeline in tenant.list_timelines().iter() {
             // collect per-timeline metrics only for active timelines
-            if let Ok(timeline_ctx) = timeline.get_context(&tenant_ctx) {
+            if timeline.is_active() {
                 let timeline_written_size = u64::from(timeline.get_last_record_lsn());
 
                 current_metrics.push((
@@ -222,8 +209,7 @@ pub async fn collect_metrics_task(
                     timeline_written_size,
                 ));
 
-                let (timeline_logical_size, is_exact) =
-                    timeline.get_current_logical_size(&timeline_ctx)?;
+                let (timeline_logical_size, is_exact) = timeline.get_current_logical_size()?;
                 // Only send timeline logical size when it is fully calculated.
                 if is_exact {
                     current_metrics.push((

pageserver/src/context.rs

@@ -1,348 +0,0 @@
-//!
-//! Most async functions throughout the pageserver take a `ctx: &RequestContext`
-//! argument. It is used to control desired behaviour of the operation, and to
-//! allow cancelling the operation gracefully.
-//!
-//! # Context hierarchy
-//!
-//! RequestContext's form a hierarchy. For example:
-//!
-//!  listener context (LibpqEndpointListener)
-//!    connection context (PageRequestHandler)
-//!      per-request context (PageRequestHandler)
-//!
-//! The top "listener context" is created at pageserver startup. The tokio
-//! task that listens on the libpq protocol TCP port holds that context. When
-//! it accepts a connection, it spawns a new task to handle that connection
-//! and creates a new per-connection context for it. The mgmt API listener,
-//! background jobs, and other things form separate but similar hierarchies.
-//!
-//! Usually, each tokio task has its own context, but it's not a strict
-//! requirement and some tasks can hold multiple contexts, and converesely,
-//! some contexts are shared by multiple tasks that work together to perform
-//! some operation.
-//!
-//! The hierarchy is not explictly tracked in the RequestContext struct
-//! itself, but only by their cancellation tokens. It's entirely possible for
-//! the parent context to be dropped before its children.
-//!
-//! # Tenant and Timeline registration
-//!
-//! Most operations are performed on a particular Tenant or Timeline. When
-//! operating on a Tenant or Timeline, it's important that the Tenant/Timeline
-//! isn't detached or deleted while there are tasks working on it. To ensure
-//! that, a RequestContext can be registered with a Tenant or Timeline. See
-//! `Tenant::register_context` and `Timeline::register_context` When
-//! shutting down a Tenant or Timeline, the shutdown routine cancels all the
-//! registered contexts, and waits for them to be dropped before completing
-//! the shutdown.
-//!
-//! To enforce that you hold a registered context when operating on a Tenant
-//! or Timeline, most functions take a TimelineRequestContext or
-//! TenantRequestContext reference as argument.
-//!
-//! NOTE: The Tenant / Timeline registration is separate from the context
-//! hierarchy. You can create a new RequestContext with TimelineRequestContext
-//! as the parent, and register it with a different timeline, for example.
-//!
-//! # Notes
-//!
-//! All RequestContexts in the system have a unique ID, and are also tracked
-//! in a global hash table, CONTEXTS.
-//!
-//! - Futures are normally not assumed to be async cancellation-safe. Pass a
-//!   RequestContext as argument and use cancel() on it instead.
-//!
-//! - If you perform an operation that depends on some external actor or the
-//!   network, use the cancellation token to check for cancellation
-//!
-//! - By convention, the appropriate context for current operation is carried in
-//!   a variable called 'ctx'. If a function handles multiple contexts, it's
-//!   best to *not* have a variable called 'ctx', to force you to think which
-//!   one to use in each call.
-//!
-//! # TODO
-//! - include a unique request ID for tracing
-//!
-
-use once_cell::sync::Lazy;
-use tokio_util::sync::CancellationToken;
-use tracing::{info, warn};
-
-use std::collections::HashMap;
-use std::sync::atomic::{AtomicU64, Ordering};
-use std::sync::Mutex;
-
-/// Each RequestContext has a unique context ID. It's just an increasing
-/// number that we assign.
-static NEXT_CONTEXT_ID: AtomicU64 = AtomicU64::new(1);
-
-/// Global registry of contexts
-static CONTEXTS: Lazy<Mutex<HashMap<RequestContextId, (TaskKind, CancellationToken)>>> =
-    Lazy::new(|| Mutex::new(HashMap::new()));
-
-#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
-pub struct RequestContextId(u64);
-
-///
-pub struct RequestContext {
-    context_id: RequestContextId,
-    task_kind: TaskKind,
-
-    download_behavior: DownloadBehavior,
-    cancellation_token: CancellationToken,
-}
-
-/// DownloadBehavior option specifies the behavior if completing the operation
-/// would require downloading a layer file from remote storage.
-#[derive(Clone, Copy, PartialEq, Eq)]
-pub enum DownloadBehavior {
-    /// Download the layer file. It can take a while.
-    Download,
-
-    /// Download the layer file, but print a warning to the log. This should be used
-    /// in code where the layer file is expected to already exist locally.
-    Warn,
-
-    /// Return a PageReconstructError::NeedsDownload error
-    Error,
-}
-
-///
-/// There are many kinds of tasks in the system. Some are associated with a particular
-/// tenant or timeline, while others are global.
-///
-/// The task kind affects the shutdown sequence on pageserver shutdown and on detach
-/// of an individual tenant. For example, when shutting down the pageserver, we shut
-/// down the LibpqEndpointListeners first, so that we don't accept any more client
-/// connections while we perform the rest of the shutdown duties. See
-/// [`Timeline::graceful_shutdown and`] and [`tenant_mgr::shutdown_pageserver`]
-/// for details.
-///
-/// Note that we don't try to limit how many task of a certain kind can be running
-/// at the same time.
-///
-#[derive(Debug, PartialEq, Eq, Clone, Copy)]
-pub enum TaskKind {
-    // libpq listener task. It just accepts connection and spawns a
-    // PageRequestHandler task for each connection.
-    LibpqEndpointListener,
-
-    // HTTP endpoint listener.
-    HttpEndpointListener,
-
-    // Task that handles a single connection. A PageRequestHandler task
-    // starts detached from any particular tenant or timeline, but it can be
-    // associated with one later, after receiving a command from the client.
-    PageRequestHandler,
-
-    // Context for one management API request
-    MgmtRequest,
-
-    // Manages the WAL receiver connection for one timeline. It subscribes to
-    // events from storage_broker, decides which safekeeper to connect to. It spawns a
-    // separate WalReceiverConnection task to handle each connection.
-    WalReceiverManager,
-
-    // Handles a connection to a safekeeper, to stream WAL to a timeline.
-    WalReceiverConnection,
-
-    // Garbage collection worker. One per tenant
-    GarbageCollector,
-
-    // Compaction. One per tenant.
-    Compaction,
-
-    // Initial logical size calculation
-    InitialLogicalSizeCalculation,
-
-    // Task that flushes frozen in-memory layers to disk
-    LayerFlush,
-
-    // Task that uploads a file to remote storage
-    RemoteUploadTask,
-
-    // Task that downloads a file from remote storage
-    RemoteDownloadTask,
-
-    // task that handles the initial downloading of all tenants
-    InitialLoad,
-
-    // task that handles attaching a tenant
-    Attach,
-
-    // task that handles metrics collection
-    MetricsCollection,
-
-    // task that drives downloading layers
-    DownloadAllRemoteLayers,
-
-    // Only used in unit tests
-    UnitTest,
-}
-
-impl Drop for RequestContext {
-    fn drop(&mut self) {
-        CONTEXTS
-            .lock()
-            .unwrap()
-            .remove(&self.context_id)
-            .expect("context is not in global registry");
-    }
-}
-
-impl RequestContext {
-    /// Create a new RequestContext
-    pub fn new(task_kind: TaskKind, download_behavior: DownloadBehavior) -> Self {
-        let cancellation_token = CancellationToken::new();
-        let context_id = RequestContextId(NEXT_CONTEXT_ID.fetch_add(1, Ordering::Relaxed));
-        CONTEXTS
-            .lock()
-            .unwrap()
-            .insert(context_id, (task_kind, cancellation_token.clone()));
-
-        RequestContext {
-            task_kind,
-            context_id,
-            download_behavior,
-            cancellation_token,
-        }
-    }
-
-    /// Create a new RequestContext, as a child of 'parent'.
-    pub fn with_parent(
-        task_kind: TaskKind,
-        download_behavior: DownloadBehavior,
-        parent: &RequestContext,
-    ) -> Self {
-        let cancellation_token = parent.cancellation_token.child_token();
-        let context_id = RequestContextId(NEXT_CONTEXT_ID.fetch_add(1, Ordering::Relaxed));
-        CONTEXTS
-            .lock()
-            .unwrap()
-            .insert(context_id, (task_kind, cancellation_token.clone()));
-
-        RequestContext {
-            task_kind,
-            context_id,
-            download_behavior,
-            cancellation_token,
-        }
-    }
-
-    pub fn context_id(&self) -> RequestContextId {
-        self.context_id
-    }
-
-    pub fn task_kind(&self) -> TaskKind {
-        self.task_kind
-    }
-
-    pub fn download_behavior(&self) -> DownloadBehavior {
-        self.download_behavior
-    }
-
-    pub fn cancellation_token(&self) -> &CancellationToken {
-        &self.cancellation_token
-    }
-
-    pub fn is_cancelled(&self) -> bool {
-        self.cancellation_token.is_cancelled()
-    }
-
-    pub async fn cancelled(&self) {
-        self.cancellation_token.cancelled().await
-    }
-}
-
-///
-/// Cancel all the contexts in 'context_ids' and wait for them to finish.
-///
-/// Whenever we notice that one of the contexts has finished, it is removed
-/// from 'context_ids'. On return, it is empty.
-///
-pub async fn cancel_and_wait(context_ids: &mut Vec<RequestContextId>) {
-    {
-        let contexts = CONTEXTS.lock().unwrap();
-        context_ids.retain(|context_id| {
-            if let Some((task_kind, cancellation_token)) = contexts.get(context_id) {
-                info!("cancelling task {task_kind:?} with ID {context_id:?}");
-                cancellation_token.cancel();
-                true
-            } else {
-                // Already gone
-                false
-            }
-        });
-    }
-    wait_contexts_to_finish(context_ids).await
-}
-
-async fn wait_contexts_to_finish(context_ids: &mut Vec<RequestContextId>) {
-    let mut n = 0;
-    while !context_ids.is_empty() {
-        {
-            let contexts = CONTEXTS.lock().unwrap();
-            while let Some(context_id) = context_ids.last() {
-                if let Some((task_kind, _cancellation_token)) = contexts.get(context_id) {
-                    info!("waiting for task {task_kind:?} with ID {context_id:?} to finish");
-                    break;
-                } else {
-                    context_ids.pop();
-                }
-            }
-        }
-        if !context_ids.is_empty() {
-            crate::exponential_backoff(
-                n,
-                crate::DEFAULT_BASE_BACKOFF_SECONDS,
-                crate::DEFAULT_MAX_BACKOFF_SECONDS,
-            )
-            .await;
-            n += 1;
-        }
-    }
-}
-
-/// Cancel and wait for all tasks of given 'kind' to finish
-pub async fn shutdown_tasks(kind: TaskKind) {
-    let mut context_ids = Vec::new();
-    {
-        let contexts = CONTEXTS.lock().unwrap();
-        for (&context_id, (task_kind, cancellation_token)) in contexts.iter() {
-            if *task_kind == kind {
-                cancellation_token.cancel();
-                context_ids.push(context_id);
-            }
-        }
-    }
-    wait_contexts_to_finish(&mut context_ids).await
-}
-
-/// Cancel all remaining contexts.
-///
-/// This is used as part of pageserver shutdown. We have already shut down all
-/// tasks / contexts, this is just a backstop or sanity check to make sure we
-/// didn't miss anything. Hence, also print a warning for any remaining tasks.
-pub async fn shutdown_all_tasks() {
-    loop {
-        let mut context_ids = Vec::new();
-        {
-            let contexts = CONTEXTS.lock().unwrap();
-
-            if contexts.is_empty() {
-                return;
-            }
-
-            for (&context_id, (task_kind, cancellation_token)) in contexts.iter() {
-                cancellation_token.cancel();
-                context_ids.push(context_id);
-                warn!(
-                    "unexpected task of kind {:?} with ID {:?} still running",
-                    *task_kind, context_id
-                );
-            }
-        }
-        wait_contexts_to_finish(&mut context_ids).await
-    }
-}

pageserver/src/http/routes.rs

@@ -4,16 +4,16 @@ use anyhow::{anyhow, Context, Result};
 use hyper::StatusCode;
 use hyper::{Body, Request, Response, Uri};
 use remote_storage::GenericRemoteStorage;
+use tokio_util::sync::CancellationToken;
 use tracing::*;
 
 use super::models::{
     StatusResponse, TenantConfigRequest, TenantCreateRequest, TenantCreateResponse, TenantInfo,
     TimelineCreateRequest, TimelineInfo,
 };
-use crate::context::{DownloadBehavior, RequestContext, TaskKind};
 use crate::pgdatadir_mapping::LsnForTimestamp;
 use crate::tenant::config::TenantConfOpt;
-use crate::tenant::{PageReconstructError, Timeline, TimelineRequestContext};
+use crate::tenant::{with_ondemand_download, Timeline};
 use crate::{config::PageServerConf, tenant::mgr};
 use utils::{
     auth::JwtAuth,
@@ -77,50 +77,29 @@ fn check_permission(request: &Request<Body>, tenant_id: Option<TenantId>) -> Res
     })
 }
 
-fn apierror_from_prerror(err: PageReconstructError) -> ApiError {
-    match err {
-        PageReconstructError::Other(err) => ApiError::InternalServerError(err),
-        PageReconstructError::NeedsDownload(_, _) => {
-            // This shouldn't happen, because we use a RequestContext that requests to
-            // download any missing layer files on-demand.
-            ApiError::InternalServerError(anyhow::anyhow!(
-                "would need to download remote layer file"
-            ))
-        }
-        PageReconstructError::Cancelled => {
-            ApiError::InternalServerError(anyhow::anyhow!("request was cancelled"))
-        }
-        PageReconstructError::WalRedo(err) => {
-            ApiError::InternalServerError(anyhow::Error::new(err))
-        }
-    }
-}
-
 // Helper function to construct a TimelineInfo struct for a timeline
 async fn build_timeline_info(
     timeline: &Arc<Timeline>,
     include_non_incremental_logical_size: bool,
-    ctx: Option<&TimelineRequestContext>,
 ) -> anyhow::Result<TimelineInfo> {
-    let mut info = build_timeline_info_common(timeline, ctx)?;
+    let mut info = build_timeline_info_common(timeline)?;
     if include_non_incremental_logical_size {
-        if let Some(ctx) = ctx {
-            info.current_logical_size_non_incremental = Some(
-                timeline
-                    .get_current_logical_size_non_incremental(info.last_record_lsn, ctx)
-                    .await?,
-            );
-        } else {
-            info!("could not calculate non-incremental size for timeline because it is not active");
-        }
+        // XXX we should be using spawn_ondemand_logical_size_calculation here.
+        // Otherwise, if someone deletes the timeline / detaches the tenant while
+        // we're executing this function, we will outlive the timeline on-disk state.
+        info.current_logical_size_non_incremental = Some(
+            timeline
+                .get_current_logical_size_non_incremental(
+                    info.last_record_lsn,
+                    CancellationToken::new(),
+                )
+                .await?,
+        );
     }
     Ok(info)
 }
 
-fn build_timeline_info_common(
-    timeline: &Arc<Timeline>,
-    ctx: Option<&TimelineRequestContext>,
-) -> anyhow::Result<TimelineInfo> {
+fn build_timeline_info_common(timeline: &Arc<Timeline>) -> anyhow::Result<TimelineInfo> {
     let last_record_lsn = timeline.get_last_record_lsn();
     let (wal_source_connstr, last_received_msg_lsn, last_received_msg_ts) = {
         let guard = timeline.last_received_wal.lock().unwrap();
@@ -140,16 +119,12 @@ fn build_timeline_info_common(
         Lsn(0) => None,
         lsn @ Lsn(_) => Some(lsn),
     };
-    let current_logical_size = if let Some(ctx) = ctx {
-        match timeline.get_current_logical_size(ctx) {
-            Ok((size, _)) => Some(size),
-            Err(err) => {
-                error!("Timeline info creation failed to get current logical size: {err:?}");
-                None
-            }
+    let current_logical_size = match timeline.get_current_logical_size() {
+        Ok((size, _)) => Some(size),
+        Err(err) => {
+            error!("Timeline info creation failed to get current logical size: {err:?}");
+            None
         }
-    } else {
-        None
     };
     let current_physical_size = Some(timeline.layer_size_sum().approximate_is_ok());
     let state = timeline.current_state();
@@ -195,23 +170,20 @@ async fn timeline_create_handler(mut request: Request<Body>) -> Result<Response<
         .new_timeline_id
         .unwrap_or_else(TimelineId::generate);
 
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
-    let (tenant, tenant_ctx) = mgr::get_active_tenant(tenant_id, &ctx)
+    let tenant = mgr::get_tenant(tenant_id, true)
         .await
         .map_err(ApiError::NotFound)?;
     match tenant.create_timeline(
         new_timeline_id,
         request_data.ancestor_timeline_id.map(TimelineId::from),
         request_data.ancestor_start_lsn,
-        request_data.pg_version.unwrap_or(crate::DEFAULT_PG_VERSION),
-        &tenant_ctx,
+        request_data.pg_version.unwrap_or(crate::DEFAULT_PG_VERSION)
     )
     .instrument(info_span!("timeline_create", tenant = %tenant_id, new_timeline = ?request_data.new_timeline_id, timeline_id = %new_timeline_id, lsn=?request_data.ancestor_start_lsn, pg_version=?request_data.pg_version))
     .await {
-        Ok(Some((new_timeline, timeline_ctx))) => {
+        Ok(Some(new_timeline)) => {
             // Created. Construct a TimelineInfo for it.
-            let timeline_info = build_timeline_info_common(&new_timeline, Some(&timeline_ctx))
+            let timeline_info = build_timeline_info_common(&new_timeline)
                 .map_err(ApiError::InternalServerError)?;
             json_response(StatusCode::CREATED, timeline_info)
         }
@@ -226,25 +198,21 @@ async fn timeline_list_handler(request: Request<Body>) -> Result<Response<Body>,
         query_param_present(&request, "include-non-incremental-logical-size");
     check_permission(&request, Some(tenant_id))?;
 
-    let top_ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
     let response_data = async {
-        let (tenant, tenant_ctx) = mgr::get_active_tenant(tenant_id, &top_ctx)
+        let tenant = mgr::get_tenant(tenant_id, true)
             .await
             .map_err(ApiError::NotFound)?;
         let timelines = tenant.list_timelines();
 
         let mut response_data = Vec::with_capacity(timelines.len());
         for timeline in timelines {
-            let timeline_ctx = timeline.get_context(&tenant_ctx).ok();
-            let timeline_info = build_timeline_info(
-                &timeline,
-                include_non_incremental_logical_size,
-                timeline_ctx.as_ref(),
-            )
-            .await
-            .context("Failed to convert tenant timeline {timeline_id} into the local one: {e:?}")
-            .map_err(ApiError::InternalServerError)?;
+            let timeline_info =
+                build_timeline_info(&timeline, include_non_incremental_logical_size)
+                    .await
+                    .context(
+                        "Failed to convert tenant timeline {timeline_id} into the local one: {e:?}",
+                    )
+                    .map_err(ApiError::InternalServerError)?;
 
             response_data.push(timeline_info);
         }
@@ -293,26 +261,19 @@ async fn timeline_detail_handler(request: Request<Body>) -> Result<Response<Body
         query_param_present(&request, "include-non-incremental-logical-size");
     check_permission(&request, Some(tenant_id))?;
 
-    let top_ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
     let timeline_info = async {
-        let (tenant, tenant_ctx) = mgr::get_active_tenant(tenant_id, &top_ctx)
+        let tenant = mgr::get_tenant(tenant_id, true)
             .await
             .map_err(ApiError::NotFound)?;
 
         let timeline = tenant
-            .get_timeline(timeline_id)
+            .get_timeline(timeline_id, false)
             .map_err(ApiError::NotFound)?;
-        let timeline_ctx = timeline.get_context(&tenant_ctx).ok();
 
-        let timeline_info = build_timeline_info(
-            &timeline,
-            include_non_incremental_logical_size,
-            timeline_ctx.as_ref(),
-        )
-        .await
-        .context("Failed to get local timeline info: {e:#}")
-        .map_err(ApiError::InternalServerError)?;
+        let timeline_info = build_timeline_info(&timeline, include_non_incremental_logical_size)
+            .await
+            .context("Failed to get local timeline info: {e:#}")
+            .map_err(ApiError::InternalServerError)?;
 
         Ok::<_, ApiError>(timeline_info)
     }
@@ -333,19 +294,13 @@ async fn get_lsn_by_timestamp_handler(request: Request<Body>) -> Result<Response
         .map_err(ApiError::BadRequest)?;
     let timestamp_pg = postgres_ffi::to_pg_timestamp(timestamp);
 
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
-    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
+    let timeline = mgr::get_tenant(tenant_id, true)
         .await
+        .and_then(|tenant| tenant.get_timeline(timeline_id, true))
         .map_err(ApiError::NotFound)?;
-
-    let (timeline, ctx) = tenant
-        .get_active_timeline(timeline_id, &ctx)
-        .map_err(ApiError::NotFound)?;
-    let result = timeline
-        .find_lsn_for_timestamp(timestamp_pg, &ctx)
+    let result = with_ondemand_download(|| timeline.find_lsn_for_timestamp(timestamp_pg))
         .await
-        .map_err(apierror_from_prerror)?;
+        .map_err(ApiError::InternalServerError)?;
 
     let result = match result {
         LsnForTimestamp::Present(lsn) => format!("{lsn}"),
@@ -385,10 +340,7 @@ async fn timeline_delete_handler(request: Request<Body>) -> Result<Response<Body
     let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
     check_permission(&request, Some(tenant_id))?;
 
-    // deleting shouldn't require downloading anything
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Warn);
-
-    mgr::delete_timeline(tenant_id, timeline_id, &ctx)
+    mgr::delete_timeline(tenant_id, timeline_id)
         .instrument(info_span!("timeline_delete", tenant = %tenant_id, timeline = %timeline_id))
         .await
         // FIXME: Errors from `delete_timeline` can occur for a number of reasons, incuding both
@@ -466,10 +418,8 @@ async fn tenant_status(request: Request<Body>) -> Result<Response<Body>, ApiErro
     let tenant_id: TenantId = parse_request_param(&request, "tenant_id")?;
     check_permission(&request, Some(tenant_id))?;
 
-    let mut _req_ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
     let tenant_info = async {
-        let tenant = mgr::get_tenant(tenant_id).await?;
+        let tenant = mgr::get_tenant(tenant_id, false).await?;
 
         // Calculate total physical size of all timelines
         let mut current_physical_size = 0;
@@ -496,15 +446,13 @@ async fn tenant_size_handler(request: Request<Body>) -> Result<Response<Body>, A
     let tenant_id: TenantId = parse_request_param(&request, "tenant_id")?;
     check_permission(&request, Some(tenant_id))?;
 
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
-    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
+    let tenant = mgr::get_tenant(tenant_id, true)
         .await
         .map_err(ApiError::InternalServerError)?;
 
     // this can be long operation, it currently is not backed by any request coalescing or similar
     let inputs = tenant
-        .gather_size_inputs(&ctx)
+        .gather_size_inputs()
         .await
         .map_err(ApiError::InternalServerError)?;
 
@@ -547,8 +495,6 @@ fn bad_duration<'a>(field_name: &'static str, value: &'a str) -> impl 'a + Fn()
 async fn tenant_create_handler(mut request: Request<Body>) -> Result<Response<Body>, ApiError> {
     check_permission(&request, None)?;
 
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
     let request_data: TenantCreateRequest = json_request(&mut request).await?;
 
     let mut tenant_conf = TenantConfOpt::default();
@@ -637,9 +583,9 @@ async fn tenant_create_handler(mut request: Request<Body>) -> Result<Response<Bo
         Some(tenant) => {
             // We created the tenant. Existing API semantics are that the tenant
             // is Active when this function returns.
-            if let res @ Err(_) = tenant.wait_to_become_active(ctx).await {
+            if let res @ Err(_) = tenant.wait_to_become_active().await {
                 // This shouldn't happen because we just created the tenant directory
-                // in tenant::mgr::create_tenant, and there aren't any remote timelines
+                // in tenant_mgr::create_tenant, and there aren't any remote timelines
                 // to load, so, nothing can really fail during load.
                 // Don't do cleanup because we don't know how we got here.
                 // The tenant will likely be in `Broken` state and subsequent
@@ -661,8 +607,6 @@ async fn tenant_config_handler(mut request: Request<Body>) -> Result<Response<Bo
     let tenant_id = request_data.tenant_id;
     check_permission(&request, Some(tenant_id))?;
 
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
     let mut tenant_conf: TenantConfOpt = Default::default();
     if let Some(gc_period) = request_data.gc_period {
         tenant_conf.gc_period = Some(
@@ -725,7 +669,7 @@ async fn tenant_config_handler(mut request: Request<Body>) -> Result<Response<Bo
     }
 
     let state = get_state(&request);
-    mgr::update_tenant_config(state.conf, tenant_conf, tenant_id, &ctx)
+    mgr::update_tenant_config(state.conf, tenant_conf, tenant_id)
         .instrument(info_span!("tenant_config", tenant = ?tenant_id))
         .await
         // FIXME: `update_tenant_config` can fail because of both user and internal errors.
@@ -777,21 +721,11 @@ async fn timeline_gc_handler(mut request: Request<Body>) -> Result<Response<Body
 
     let gc_req: TimelineGcRequest = json_request(&mut request).await?;
 
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
-    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
-        .await
-        .map_err(ApiError::NotFound)?;
-
-    let gc_horizon = gc_req.gc_horizon.unwrap_or_else(|| tenant.get_gc_horizon());
-    // Use tenant's pitr setting
-    let pitr = tenant.get_pitr_interval();
-
-    fail::fail_point!("immediate_gc_task_pre");
-    let gc_result = tenant
-        .gc_iteration(Some(timeline_id), gc_horizon, pitr, &ctx)
-        .instrument(info_span!("manual_gc", tenant = %tenant_id, timeline = %timeline_id))
+    let wait_task_done = mgr::immediate_gc(tenant_id, timeline_id, gc_req).await?;
+    let gc_result = wait_task_done
         .await
+        .context("wait for gc task")
+        .map_err(ApiError::InternalServerError)?
         .map_err(ApiError::InternalServerError)?;
 
     json_response(StatusCode::OK, gc_result)
@@ -804,20 +738,17 @@ async fn timeline_compact_handler(request: Request<Body>) -> Result<Response<Bod
     let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
     check_permission(&request, Some(tenant_id))?;
 
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
-    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
-        .await
-        .map_err(ApiError::NotFound)?;
-    let (timeline, ctx) = tenant
-        .get_active_timeline(timeline_id, &ctx)
-        .map_err(ApiError::NotFound)?;
-    timeline
-        .compact(&ctx)
-        .instrument(info_span!("manual_compact", tenant = %tenant_id, timeline = %timeline_id))
+    let result_receiver = mgr::immediate_compact(tenant_id, timeline_id)
         .await
+        .context("spawn compaction task")
         .map_err(ApiError::InternalServerError)?;
 
+    let result: anyhow::Result<()> = result_receiver
+        .await
+        .context("receive compaction result")
+        .map_err(ApiError::InternalServerError)?;
+    result.map_err(ApiError::InternalServerError)?;
+
     json_response(StatusCode::OK, ())
 }
 
@@ -828,21 +759,18 @@ async fn timeline_checkpoint_handler(request: Request<Body>) -> Result<Response<
     let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
     check_permission(&request, Some(tenant_id))?;
 
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
-    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
+    let tenant = mgr::get_tenant(tenant_id, true)
         .await
         .map_err(ApiError::NotFound)?;
-    let (timeline, ctx) = tenant
-        .get_active_timeline(timeline_id, &ctx)
+    let timeline = tenant
+        .get_timeline(timeline_id, true)
         .map_err(ApiError::NotFound)?;
     timeline
         .freeze_and_flush()
         .await
         .map_err(ApiError::InternalServerError)?;
     timeline
-        .compact(&ctx)
-        .instrument(info_span!("manual_compact", tenant = %tenant_id, timeline = %timeline_id))
+        .compact()
         .await
         .map_err(ApiError::InternalServerError)?;
 
@@ -856,15 +784,13 @@ async fn timeline_download_remote_layers_handler_post(
     let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
     check_permission(&request, Some(tenant_id))?;
 
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
-    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
+    let tenant = mgr::get_tenant(tenant_id, true)
         .await
         .map_err(ApiError::NotFound)?;
-    let (timeline, ctx) = tenant
-        .get_active_timeline(timeline_id, &ctx)
+    let timeline = tenant
+        .get_timeline(timeline_id, true)
         .map_err(ApiError::NotFound)?;
-    match timeline.spawn_download_all_remote_layers(&ctx).await {
+    match timeline.spawn_download_all_remote_layers().await {
         Ok(st) => json_response(StatusCode::ACCEPTED, st),
         Err(st) => json_response(StatusCode::CONFLICT, st),
     }
@@ -877,13 +803,11 @@ async fn timeline_download_remote_layers_handler_get(
     let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
     check_permission(&request, Some(tenant_id))?;
 
-    let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download);
-
-    let (tenant, ctx) = mgr::get_active_tenant(tenant_id, &ctx)
+    let tenant = mgr::get_tenant(tenant_id, true)
         .await
         .map_err(ApiError::NotFound)?;
-    let (timeline, _ctx) = tenant
-        .get_active_timeline(timeline_id, &ctx)
+    let timeline = tenant
+        .get_timeline(timeline_id, true)
         .map_err(ApiError::NotFound)?;
     let info = timeline
         .get_download_all_remote_layers_task_info()

pageserver/src/import_datadir.rs

@@ -13,7 +13,7 @@ use tracing::*;
 use walkdir::WalkDir;
 
 use crate::pgdatadir_mapping::*;
-use crate::tenant::{Timeline, TimelineRequestContext};
+use crate::tenant::Timeline;
 use crate::walingest::WalIngest;
 use crate::walrecord::DecodedWALRecord;
 use pageserver_api::reltag::{RelTag, SlruKind};
@@ -47,7 +47,6 @@ pub async fn import_timeline_from_postgres_datadir(
     tline: &Timeline,
     pgdata_path: &Path,
     pgdata_lsn: Lsn,
-    ctx: &TimelineRequestContext,
 ) -> Result<()> {
     let mut pg_control: Option<ControlFileData> = None;
 
@@ -70,7 +69,7 @@ pub async fn import_timeline_from_postgres_datadir(
             let mut file = tokio::fs::File::open(absolute_path).await?;
             let len = metadata.len() as usize;
             if let Some(control_file) =
-                import_file(&mut modification, relative_path, &mut file, len, ctx).await?
+                import_file(&mut modification, relative_path, &mut file, len).await?
             {
                 pg_control = Some(control_file);
             }
@@ -100,7 +99,6 @@ pub async fn import_timeline_from_postgres_datadir(
         tline,
         Lsn(pg_control.checkPointCopy.redo),
         pgdata_lsn,
-        ctx,
     )
     .await?;
 
@@ -115,7 +113,6 @@ async fn import_rel(
     dboid: Oid,
     reader: &mut (impl AsyncRead + Send + Sync + Unpin),
     len: usize,
-    ctx: &TimelineRequestContext,
 ) -> anyhow::Result<()> {
     // Does it look like a relation file?
     trace!("importing rel file {}", path.display());
@@ -146,14 +143,7 @@ async fn import_rel(
     // Call put_rel_creation for every segment of the relation,
     // because there is no guarantee about the order in which we are processing segments.
     // ignore "relation already exists" error
-    //
-    // FIXME: use proper error type for this, instead of parsing the error message.
-    // Or better yet, keep track of which relations we've already created
-    // https://github.com/neondatabase/neon/issues/3309
-    if let Err(e) = modification
-        .put_rel_creation(rel, nblocks as u32, ctx)
-        .await
-    {
+    if let Err(e) = modification.put_rel_creation(rel, nblocks as u32) {
         if e.to_string().contains("already exists") {
             debug!("relation {} already exists. we must be extending it", rel);
         } else {
@@ -188,7 +178,7 @@ async fn import_rel(
     //
     // If we process rel segments out of order,
     // put_rel_extend will skip the update.
-    modification.put_rel_extend(rel, blknum, ctx).await?;
+    modification.put_rel_extend(rel, blknum)?;
 
     Ok(())
 }
@@ -201,7 +191,6 @@ async fn import_slru(
     path: &Path,
     reader: &mut (impl AsyncRead + Send + Sync + Unpin),
     len: usize,
-    ctx: &TimelineRequestContext,
 ) -> anyhow::Result<()> {
     info!("importing slru file {path:?}");
 
@@ -217,9 +206,7 @@ async fn import_slru(
 
     ensure!(nblocks <= pg_constants::SLRU_PAGES_PER_SEGMENT as usize);
 
-    modification
-        .put_slru_segment_creation(slru, segno, nblocks as u32, ctx)
-        .await?;
+    modification.put_slru_segment_creation(slru, segno, nblocks as u32)?;
 
     let mut rpageno = 0;
     loop {
@@ -259,15 +246,15 @@ async fn import_wal(
     tline: &Timeline,
     startpoint: Lsn,
     endpoint: Lsn,
-    ctx: &TimelineRequestContext,
 ) -> anyhow::Result<()> {
+    use std::io::Read;
     let mut waldecoder = WalStreamDecoder::new(startpoint, tline.pg_version);
 
     let mut segno = startpoint.segment_number(WAL_SEGMENT_SIZE);
     let mut offset = startpoint.segment_offset(WAL_SEGMENT_SIZE);
     let mut last_lsn = startpoint;
 
-    let mut walingest = WalIngest::new(tline, startpoint, ctx).await?;
+    let mut walingest = WalIngest::new(tline, startpoint).await?;
 
     while last_lsn <= endpoint {
         // FIXME: assume postgresql tli 1 for now
@@ -290,7 +277,6 @@ async fn import_wal(
             file.seek(std::io::SeekFrom::Start(offset as u64))?;
         }
 
-        use std::io::Read;
         let nread = file.read_to_end(&mut buf)?;
         if nread != WAL_SEGMENT_SIZE - offset {
             // Maybe allow this for .partial files?
@@ -305,7 +291,7 @@ async fn import_wal(
         while last_lsn <= endpoint {
             if let Some((lsn, recdata)) = waldecoder.poll_decode()? {
                 walingest
-                    .ingest_record(recdata, lsn, &mut modification, &mut decoded, ctx)
+                    .ingest_record(recdata, lsn, &mut modification, &mut decoded)
                     .await?;
                 last_lsn = lsn;
 
@@ -334,7 +320,6 @@ pub async fn import_basebackup_from_tar(
     tline: &Timeline,
     reader: &mut (impl AsyncRead + Send + Sync + Unpin),
     base_lsn: Lsn,
-    ctx: &TimelineRequestContext,
 ) -> Result<()> {
     info!("importing base at {base_lsn}");
     let mut modification = tline.begin_modification(base_lsn);
@@ -353,7 +338,7 @@ pub async fn import_basebackup_from_tar(
         match header.entry_type() {
             tokio_tar::EntryType::Regular => {
                 if let Some(res) =
-                    import_file(&mut modification, file_path.as_ref(), &mut entry, len, ctx).await?
+                    import_file(&mut modification, file_path.as_ref(), &mut entry, len).await?
                 {
                     // We found the pg_control file.
                     pg_control = Some(res);
@@ -385,14 +370,13 @@ pub async fn import_wal_from_tar(
     reader: &mut (impl AsyncRead + Send + Sync + Unpin),
     start_lsn: Lsn,
     end_lsn: Lsn,
-    ctx: &TimelineRequestContext,
 ) -> Result<()> {
     // Set up walingest mutable state
     let mut waldecoder = WalStreamDecoder::new(start_lsn, tline.pg_version);
     let mut segno = start_lsn.segment_number(WAL_SEGMENT_SIZE);
     let mut offset = start_lsn.segment_offset(WAL_SEGMENT_SIZE);
     let mut last_lsn = start_lsn;
-    let mut walingest = WalIngest::new(tline, start_lsn, ctx).await?;
+    let mut walingest = WalIngest::new(tline, start_lsn).await?;
 
     // Ingest wal until end_lsn
     info!("importing wal until {}", end_lsn);
@@ -441,7 +425,7 @@ pub async fn import_wal_from_tar(
         while last_lsn <= end_lsn {
             if let Some((lsn, recdata)) = waldecoder.poll_decode()? {
                 walingest
-                    .ingest_record(recdata, lsn, &mut modification, &mut decoded, ctx)
+                    .ingest_record(recdata, lsn, &mut modification, &mut decoded)
                     .await?;
                 last_lsn = lsn;
 
@@ -476,7 +460,6 @@ async fn import_file(
     file_path: &Path,
     reader: &mut (impl AsyncRead + Send + Sync + Unpin),
     len: usize,
-    ctx: &TimelineRequestContext,
 ) -> Result<Option<ControlFileData>> {
     let file_name = match file_path.file_name() {
         Some(name) => name.to_string_lossy(),
@@ -509,16 +492,14 @@ async fn import_file(
             }
             "pg_filenode.map" => {
                 let bytes = read_all_bytes(reader).await?;
-                modification
-                    .put_relmap_file(spcnode, dbnode, bytes, ctx)
-                    .await?;
+                modification.put_relmap_file(spcnode, dbnode, bytes)?;
                 debug!("imported relmap file")
             }
             "PG_VERSION" => {
                 debug!("ignored PG_VERSION file");
             }
             _ => {
-                import_rel(modification, file_path, spcnode, dbnode, reader, len, ctx).await?;
+                import_rel(modification, file_path, spcnode, dbnode, reader, len).await?;
                 debug!("imported rel creation");
             }
         }
@@ -534,41 +515,37 @@ async fn import_file(
         match file_name.as_ref() {
             "pg_filenode.map" => {
                 let bytes = read_all_bytes(reader).await?;
-                modification
-                    .put_relmap_file(spcnode, dbnode, bytes, ctx)
-                    .await?;
+                modification.put_relmap_file(spcnode, dbnode, bytes)?;
                 debug!("imported relmap file")
             }
             "PG_VERSION" => {
                 debug!("ignored PG_VERSION file");
             }
             _ => {
-                import_rel(modification, file_path, spcnode, dbnode, reader, len, ctx).await?;
+                import_rel(modification, file_path, spcnode, dbnode, reader, len).await?;
                 debug!("imported rel creation");
             }
         }
     } else if file_path.starts_with("pg_xact") {
         let slru = SlruKind::Clog;
 
-        import_slru(modification, slru, file_path, reader, len, ctx).await?;
+        import_slru(modification, slru, file_path, reader, len).await?;
         debug!("imported clog slru");
     } else if file_path.starts_with("pg_multixact/offsets") {
         let slru = SlruKind::MultiXactOffsets;
 
-        import_slru(modification, slru, file_path, reader, len, ctx).await?;
+        import_slru(modification, slru, file_path, reader, len).await?;
         debug!("imported multixact offsets slru");
     } else if file_path.starts_with("pg_multixact/members") {
         let slru = SlruKind::MultiXactMembers;
 
-        import_slru(modification, slru, file_path, reader, len, ctx).await?;
+        import_slru(modification, slru, file_path, reader, len).await?;
         debug!("imported multixact members slru");
     } else if file_path.starts_with("pg_twophase") {
         let xid = u32::from_str_radix(file_name.as_ref(), 16)?;
 
         let bytes = read_all_bytes(reader).await?;
-        modification
-            .put_twophase_file(xid, Bytes::copy_from_slice(&bytes[..]), ctx)
-            .await?;
+        modification.put_twophase_file(xid, Bytes::copy_from_slice(&bytes[..]))?;
         debug!("imported twophase file");
     } else if file_path.starts_with("pg_wal") {
         debug!("found wal file in base section. ignore it");

pageserver/src/lib.rs

@@ -2,7 +2,6 @@ mod auth;
 pub mod basebackup;
 pub mod config;
 pub mod consumption_metrics;
-pub mod context;
 pub mod http;
 pub mod import_datadir;
 pub mod keyspace;
@@ -22,6 +21,7 @@ pub mod walredo;
 
 use std::path::Path;
 
+use crate::task_mgr::TaskKind;
 use tracing::info;
 
 /// Current storage format version
@@ -41,6 +41,35 @@ pub const DELTA_FILE_MAGIC: u16 = 0x5A61;
 
 static ZERO_PAGE: bytes::Bytes = bytes::Bytes::from_static(&[0u8; 8192]);
 
+pub async fn shutdown_pageserver(exit_code: i32) {
+    // Shut down the libpq endpoint task. This prevents new connections from
+    // being accepted.
+    task_mgr::shutdown_tasks(Some(TaskKind::LibpqEndpointListener), None, None).await;
+
+    // Shut down any page service tasks.
+    task_mgr::shutdown_tasks(Some(TaskKind::PageRequestHandler), None, None).await;
+
+    // Shut down all the tenants. This flushes everything to disk and kills
+    // the checkpoint and GC tasks.
+    tenant::mgr::shutdown_all_tenants().await;
+
+    // Stop syncing with remote storage.
+    //
+    // FIXME: Does this wait for the sync tasks to finish syncing what's queued up?
+    // Should it?
+    task_mgr::shutdown_tasks(Some(TaskKind::RemoteUploadTask), None, None).await;
+
+    // Shut down the HTTP endpoint last, so that you can still check the server's
+    // status while it's shutting down.
+    // FIXME: We should probably stop accepting commands like attach/detach earlier.
+    task_mgr::shutdown_tasks(Some(TaskKind::HttpEndpointListener), None, None).await;
+
+    // There should be nothing left, but let's be sure
+    task_mgr::shutdown_tasks(None, None, None).await;
+    info!("Shut down successfully completed");
+    std::process::exit(exit_code);
+}
+
 const DEFAULT_BASE_BACKOFF_SECONDS: f64 = 0.1;
 const DEFAULT_MAX_BACKOFF_SECONDS: f64 = 3.0;
 

pageserver/src/page_service.rs

@@ -43,27 +43,24 @@ use utils::{
 use crate::auth::check_permission;
 use crate::basebackup;
 use crate::config::PageServerConf;
-use crate::context::{DownloadBehavior, RequestContext, TaskKind};
 use crate::import_datadir::import_wal_from_tar;
 use crate::metrics::{LIVE_CONNECTIONS_COUNT, SMGR_QUERY_TIME};
 use crate::task_mgr;
+use crate::task_mgr::TaskKind;
 use crate::tenant::mgr;
-use crate::tenant::{Tenant, TenantRequestContext, Timeline, TimelineRequestContext};
+use crate::tenant::{Tenant, Timeline};
 use crate::trace::Tracer;
 
 use postgres_ffi::pg_constants::DEFAULTTABLESPACE_OID;
 use postgres_ffi::BLCKSZ;
 
-fn copyin_stream<'a>(
-    pgb: &'a mut PostgresBackend,
-    ctx: &'a RequestContext,
-) -> impl Stream<Item = io::Result<Bytes>> + 'a {
+fn copyin_stream(pgb: &mut PostgresBackend) -> impl Stream<Item = io::Result<Bytes>> + '_ {
     async_stream::try_stream! {
         loop {
             let msg = tokio::select! {
                 biased;
 
-                _ = ctx.cancelled() => {
+                _ = task_mgr::shutdown_watcher() => {
                     // We were requested to shut down.
                     let msg = format!("pageserver is shutting down");
                     let _ = pgb.write_message(&BeMessage::ErrorResponse(&msg, None));
@@ -126,7 +123,6 @@ pub async fn libpq_listener_main(
     auth: Option<Arc<JwtAuth>>,
     listener: TcpListener,
     auth_type: AuthType,
-    listener_ctx: RequestContext,
 ) -> anyhow::Result<()> {
     listener.set_nonblocking(true)?;
     let tokio_listener = tokio::net::TcpListener::from_std(listener)?;
@@ -135,9 +131,8 @@ pub async fn libpq_listener_main(
     while let Some(res) = tokio::select! {
         biased;
 
-        _ = listener_ctx.cancelled() => {
+        _ = task_mgr::shutdown_watcher() => {
             // We were requested to shut down.
-            info!("libpq listener shutting down");
             None
         }
 
@@ -151,33 +146,18 @@ pub async fn libpq_listener_main(
                 debug!("accepted connection from {}", peer_addr);
                 let local_auth = auth.clone();
 
-                let connection_ctx = RequestContext::with_parent(
-                    TaskKind::PageRequestHandler,
-                    DownloadBehavior::Download,
-                    &listener_ctx,
-                );
-
                 // PageRequestHandler tasks are not associated with any particular
                 // timeline in the task manager. In practice most connections will
                 // only deal with a particular timeline, but we don't know which one
                 // yet.
                 task_mgr::spawn(
                     &tokio::runtime::Handle::current(),
+                    TaskKind::PageRequestHandler,
+                    None,
+                    None,
                     "serving compute connection task",
                     false,
-                    async move {
-                        if let Err(err) = page_service_conn_main(
-                            conf,
-                            local_auth,
-                            socket,
-                            auth_type,
-                            connection_ctx,
-                        )
-                        .await
-                        {
-                            error!("connection handler exited with error: {err:?}");
-                        }
-                    },
+                    page_service_conn_main(conf, local_auth, socket, auth_type),
                 );
             }
             Err(err) => {
@@ -197,7 +177,6 @@ async fn page_service_conn_main(
     auth: Option<Arc<JwtAuth>>,
     socket: tokio::net::TcpStream,
     auth_type: AuthType,
-    connection_ctx: RequestContext,
 ) -> anyhow::Result<()> {
     // Immediately increment the gauge, then create a job to decrement it on task exit.
     // One of the pros of `defer!` is that this will *most probably*
@@ -212,13 +191,11 @@ async fn page_service_conn_main(
         .set_nodelay(true)
         .context("could not set TCP_NODELAY")?;
 
-    let cancellation_token = connection_ctx.cancellation_token().clone();
-
-    let mut conn_handler = PageServerHandler::new(conf, auth, connection_ctx);
+    let mut conn_handler = PageServerHandler::new(conf, auth);
     let pgbackend = PostgresBackend::new(socket, auth_type, None)?;
 
     let result = pgbackend
-        .run(&mut conn_handler, || cancellation_token.cancelled())
+        .run(&mut conn_handler, task_mgr::shutdown_watcher)
         .await;
     match result {
         Ok(()) => {
@@ -274,38 +251,35 @@ impl PageRequestMetrics {
     }
 }
 
+#[derive(Debug)]
 struct PageServerHandler {
     _conf: &'static PageServerConf,
     auth: Option<Arc<JwtAuth>>,
     claims: Option<Claims>,
-
-    connection_ctx: RequestContext,
 }
 
 impl PageServerHandler {
-    pub fn new(
-        conf: &'static PageServerConf,
-        auth: Option<Arc<JwtAuth>>,
-        connection_ctx: RequestContext,
-    ) -> Self {
+    pub fn new(conf: &'static PageServerConf, auth: Option<Arc<JwtAuth>>) -> Self {
         PageServerHandler {
             _conf: conf,
             auth,
             claims: None,
-            connection_ctx,
         }
     }
 
     #[instrument(skip(self, pgb))]
     async fn handle_pagerequests(
-        &mut self,
+        &self,
         pgb: &mut PostgresBackend,
         tenant_id: TenantId,
         timeline_id: TimelineId,
     ) -> anyhow::Result<()> {
-        let (tenant, ctx) = get_active_tenant_with_timeout(tenant_id, &self.connection_ctx).await?;
+        // NOTE: pagerequests handler exits when connection is closed,
+        //       so there is no need to reset the association
+        task_mgr::associate_with(Some(tenant_id), Some(timeline_id));
 
         // Make request tracer if needed
+        let tenant = get_active_tenant_with_timeout(tenant_id).await?;
         let mut tracer = if tenant.get_trace_read_requests() {
             let connection_id = ConnectionId::generate();
             let path = tenant
@@ -317,7 +291,7 @@ impl PageServerHandler {
         };
 
         // Check that the timeline exists
-        let (timeline, ctx) = tenant.get_active_timeline(timeline_id, &ctx)?;
+        let timeline = tenant.get_timeline(timeline_id, true)?;
 
         // switch client to COPYBOTH
         pgb.write_message(&BeMessage::CopyBothResponse)?;
@@ -329,7 +303,7 @@ impl PageServerHandler {
             let msg = tokio::select! {
                 biased;
 
-                _ = ctx.cancelled() => {
+                _ = task_mgr::shutdown_watcher() => {
                     // We were requested to shut down.
                     info!("shutdown request received in page handler");
                     break;
@@ -356,27 +330,22 @@ impl PageServerHandler {
 
             let neon_fe_msg = PagestreamFeMessage::parse(&mut copy_data_bytes.reader())?;
 
-            // TODO: We could create a new per-request context here, with unique ID.
-            // Currently we use the same per-timeline context for all requests
-
             let response = match neon_fe_msg {
                 PagestreamFeMessage::Exists(req) => {
                     let _timer = metrics.get_rel_exists.start_timer();
-                    self.handle_get_rel_exists_request(&timeline, &req, &ctx)
-                        .await
+                    self.handle_get_rel_exists_request(&timeline, &req).await
                 }
                 PagestreamFeMessage::Nblocks(req) => {
                     let _timer = metrics.get_rel_size.start_timer();
-                    self.handle_get_nblocks_request(&timeline, &req, &ctx).await
+                    self.handle_get_nblocks_request(&timeline, &req).await
                 }
                 PagestreamFeMessage::GetPage(req) => {
                     let _timer = metrics.get_page_at_lsn.start_timer();
-                    self.handle_get_page_at_lsn_request(&timeline, &req, &ctx)
-                        .await
+                    self.handle_get_page_at_lsn_request(&timeline, &req).await
                 }
                 PagestreamFeMessage::DbSize(req) => {
                     let _timer = metrics.get_db_size.start_timer();
-                    self.handle_db_size_request(&timeline, &req, &ctx).await
+                    self.handle_db_size_request(&timeline, &req).await
                 }
             };
 
@@ -397,7 +366,7 @@ impl PageServerHandler {
 
     #[instrument(skip(self, pgb))]
     async fn handle_import_basebackup(
-        &mut self,
+        &self,
         pgb: &mut PostgresBackend,
         tenant_id: TenantId,
         timeline_id: TimelineId,
@@ -405,12 +374,11 @@ impl PageServerHandler {
         _end_lsn: Lsn,
         pg_version: u32,
     ) -> Result<(), QueryError> {
+        task_mgr::associate_with(Some(tenant_id), Some(timeline_id));
         // Create empty timeline
         info!("creating new timeline");
-        let (tenant, tenant_ctx) =
-            get_active_tenant_with_timeout(tenant_id, &self.connection_ctx).await?;
-        let (timeline, ctx) =
-            tenant.create_empty_timeline(timeline_id, base_lsn, pg_version, &tenant_ctx)?;
+        let tenant = get_active_tenant_with_timeout(tenant_id).await?;
+        let timeline = tenant.create_empty_timeline(timeline_id, base_lsn, pg_version)?;
 
         // TODO mark timeline as not ready until it reaches end_lsn.
         // We might have some wal to import as well, and we should prevent compute
@@ -427,9 +395,9 @@ impl PageServerHandler {
         pgb.write_message(&BeMessage::CopyInResponse)?;
         pgb.flush().await?;
 
-        let mut copyin_stream = Box::pin(copyin_stream(pgb, &ctx));
+        let mut copyin_stream = Box::pin(copyin_stream(pgb));
         timeline
-            .import_basebackup_from_tar(&mut copyin_stream, base_lsn, &ctx)
+            .import_basebackup_from_tar(&mut copyin_stream, base_lsn)
             .await?;
 
         // Drain the rest of the Copy data
@@ -453,15 +421,16 @@ impl PageServerHandler {
 
     #[instrument(skip(self, pgb))]
     async fn handle_import_wal(
-        &mut self,
+        &self,
         pgb: &mut PostgresBackend,
         tenant_id: TenantId,
         timeline_id: TimelineId,
         start_lsn: Lsn,
         end_lsn: Lsn,
     ) -> Result<(), QueryError> {
-        let (timeline, ctx) =
-            get_active_timeline_with_timeout(tenant_id, timeline_id, &self.connection_ctx).await?;
+        task_mgr::associate_with(Some(tenant_id), Some(timeline_id));
+
+        let timeline = get_active_timeline_with_timeout(tenant_id, timeline_id).await?;
         let last_record_lsn = timeline.get_last_record_lsn();
         if last_record_lsn != start_lsn {
             return Err(QueryError::Other(
@@ -476,9 +445,9 @@ impl PageServerHandler {
         info!("importing wal");
         pgb.write_message(&BeMessage::CopyInResponse)?;
         pgb.flush().await?;
-        let mut copyin_stream = Box::pin(copyin_stream(pgb, &ctx));
+        let mut copyin_stream = Box::pin(copyin_stream(pgb));
         let mut reader = tokio_util::io::StreamReader::new(&mut copyin_stream);
-        import_wal_from_tar(&timeline, &mut reader, start_lsn, end_lsn, &ctx).await?;
+        import_wal_from_tar(&timeline, &mut reader, start_lsn, end_lsn).await?;
         info!("wal import complete");
 
         // Drain the rest of the Copy data
@@ -524,7 +493,6 @@ impl PageServerHandler {
         mut lsn: Lsn,
         latest: bool,
         latest_gc_cutoff_lsn: &RcuReadGuard<Lsn>,
-        ctx: &TimelineRequestContext,
     ) -> anyhow::Result<Lsn> {
         if latest {
             // Latest page version was requested. If LSN is given, it is a hint
@@ -548,7 +516,7 @@ impl PageServerHandler {
             if lsn <= last_record_lsn {
                 lsn = last_record_lsn;
             } else {
-                timeline.wait_lsn(lsn, ctx).await?;
+                timeline.wait_lsn(lsn).await?;
                 // Since we waited for 'lsn' to arrive, that is now the last
                 // record LSN. (Or close enough for our purposes; the
                 // last-record LSN can advance immediately after we return
@@ -558,7 +526,7 @@ impl PageServerHandler {
             if lsn == Lsn(0) {
                 anyhow::bail!("invalid LSN(0) in request");
             }
-            timeline.wait_lsn(lsn, ctx).await?;
+            timeline.wait_lsn(lsn).await?;
         }
         anyhow::ensure!(
             lsn >= **latest_gc_cutoff_lsn,
@@ -568,61 +536,60 @@ impl PageServerHandler {
         Ok(lsn)
     }
 
-    #[instrument(skip(self, timeline, req, ctx), fields(rel = %req.rel, req_lsn = %req.lsn))]
+    #[instrument(skip(self, timeline, req), fields(rel = %req.rel, req_lsn = %req.lsn))]
     async fn handle_get_rel_exists_request(
         &self,
         timeline: &Timeline,
         req: &PagestreamExistsRequest,
-        ctx: &TimelineRequestContext,
     ) -> anyhow::Result<PagestreamBeMessage> {
         let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
-        let lsn =
-            Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn, ctx)
-                .await?;
-
-        let exists = timeline
-            .get_rel_exists(req.rel, lsn, req.latest, ctx)
+        let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)
             .await?;
 
+        let exists = crate::tenant::with_ondemand_download(|| {
+            timeline.get_rel_exists(req.rel, lsn, req.latest)
+        })
+        .await?;
+
         Ok(PagestreamBeMessage::Exists(PagestreamExistsResponse {
             exists,
         }))
     }
 
-    #[instrument(skip(self, timeline, req, ctx), fields(rel = %req.rel, req_lsn = %req.lsn))]
+    #[instrument(skip(self, timeline, req), fields(rel = %req.rel, req_lsn = %req.lsn))]
     async fn handle_get_nblocks_request(
         &self,
         timeline: &Timeline,
         req: &PagestreamNblocksRequest,
-        ctx: &TimelineRequestContext,
     ) -> anyhow::Result<PagestreamBeMessage> {
         let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
-        let lsn =
-            Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn, ctx)
-                .await?;
+        let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)
+            .await?;
 
-        let n_blocks = timeline.get_rel_size(req.rel, lsn, req.latest, ctx).await?;
+        let n_blocks = crate::tenant::with_ondemand_download(|| {
+            timeline.get_rel_size(req.rel, lsn, req.latest)
+        })
+        .await?;
 
         Ok(PagestreamBeMessage::Nblocks(PagestreamNblocksResponse {
             n_blocks,
         }))
     }
 
-    #[instrument(skip(self, timeline, req, ctx), fields(dbnode = %req.dbnode, req_lsn = %req.lsn))]
+    #[instrument(skip(self, timeline, req), fields(dbnode = %req.dbnode, req_lsn = %req.lsn))]
     async fn handle_db_size_request(
         &self,
         timeline: &Timeline,
         req: &PagestreamDbSizeRequest,
-        ctx: &TimelineRequestContext,
     ) -> anyhow::Result<PagestreamBeMessage> {
         let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
-        let lsn =
-            Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn, ctx)
-                .await?;
-
-        let total_blocks = timeline
-            .get_db_size(DEFAULTTABLESPACE_OID, req.dbnode, lsn, req.latest, ctx)
+        let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)
             .await?;
+
+        let total_blocks = crate::tenant::with_ondemand_download(|| {
+            timeline.get_db_size(DEFAULTTABLESPACE_OID, req.dbnode, lsn, req.latest)
+        })
+        .await?;
         let db_size = total_blocks as i64 * BLCKSZ as i64;
 
         Ok(PagestreamBeMessage::DbSize(PagestreamDbSizeResponse {
@@ -630,17 +597,15 @@ impl PageServerHandler {
         }))
     }
 
-    #[instrument(skip(self, timeline, req, ctx), fields(rel = %req.rel, blkno = %req.blkno, req_lsn = %req.lsn))]
+    #[instrument(skip(self, timeline, req), fields(rel = %req.rel, blkno = %req.blkno, req_lsn = %req.lsn))]
     async fn handle_get_page_at_lsn_request(
         &self,
         timeline: &Timeline,
         req: &PagestreamGetPageRequest,
-        ctx: &TimelineRequestContext,
     ) -> anyhow::Result<PagestreamBeMessage> {
         let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
-        let lsn =
-            Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn, ctx)
-                .await?;
+        let lsn = Self::wait_or_get_last_lsn(timeline, req.lsn, req.latest, &latest_gc_cutoff_lsn)
+            .await?;
         /*
         // Add a 1s delay to some requests. The delay helps the requests to
         // hit the race condition from github issue #1047 more easily.
@@ -650,9 +615,10 @@ impl PageServerHandler {
         }
         */
 
-        let page = timeline
-            .get_rel_page_at_lsn(req.rel, req.blkno, lsn, req.latest, ctx)
-            .await?;
+        let page = crate::tenant::with_ondemand_download(|| {
+            timeline.get_rel_page_at_lsn(req.rel, req.blkno, lsn, req.latest)
+        })
+        .await?;
 
         Ok(PagestreamBeMessage::GetPage(PagestreamGetPageResponse {
             page,
@@ -661,7 +627,7 @@ impl PageServerHandler {
 
     #[instrument(skip(self, pgb))]
     async fn handle_basebackup_request(
-        &mut self,
+        &self,
         pgb: &mut PostgresBackend,
         tenant_id: TenantId,
         timeline_id: TimelineId,
@@ -670,14 +636,12 @@ impl PageServerHandler {
         full_backup: bool,
     ) -> anyhow::Result<()> {
         // check that the timeline exists
-        let (timeline, ctx) =
-            get_active_timeline_with_timeout(tenant_id, timeline_id, &self.connection_ctx).await?;
-
+        let timeline = get_active_timeline_with_timeout(tenant_id, timeline_id).await?;
         let latest_gc_cutoff_lsn = timeline.get_latest_gc_cutoff_lsn();
         if let Some(lsn) = lsn {
             // Backup was requested at a particular LSN. Wait for it to arrive.
             info!("waiting for {}", lsn);
-            timeline.wait_lsn(lsn, &ctx).await?;
+            timeline.wait_lsn(lsn).await?;
             timeline
                 .check_lsn_is_in_scope(lsn, &latest_gc_cutoff_lsn)
                 .context("invalid basebackup lsn")?;
@@ -687,18 +651,11 @@ impl PageServerHandler {
         pgb.write_message(&BeMessage::CopyOutResponse)?;
         pgb.flush().await?;
 
-        // Send a tarball of the latest layer on the timeline
+        /* Send a tarball of the latest layer on the timeline */
         {
             let mut writer = pgb.copyout_writer();
-            basebackup::send_basebackup_tarball(
-                &mut writer,
-                &timeline,
-                lsn,
-                prev_lsn,
-                full_backup,
-                &ctx,
-            )
-            .await?;
+            basebackup::send_basebackup_tarball(&mut writer, &timeline, lsn, prev_lsn, full_backup)
+                .await?;
         }
 
         pgb.write_message(&BeMessage::CopyDone)?;
@@ -836,9 +793,7 @@ impl postgres_backend_async::Handler for PageServerHandler {
                 .with_context(|| format!("Failed to parse timeline id from {}", params[1]))?;
 
             self.check_permission(Some(tenant_id))?;
-            let (timeline, _ctx) =
-                get_active_timeline_with_timeout(tenant_id, timeline_id, &self.connection_ctx)
-                    .await?;
+            let timeline = get_active_timeline_with_timeout(tenant_id, timeline_id).await?;
 
             let end_of_timeline = timeline.get_last_record_rlsn();
 
@@ -998,8 +953,7 @@ impl postgres_backend_async::Handler for PageServerHandler {
 
             self.check_permission(Some(tenant_id))?;
 
-            let (tenant, _ctx) =
-                get_active_tenant_with_timeout(tenant_id, &self.connection_ctx).await?;
+            let tenant = get_active_tenant_with_timeout(tenant_id).await?;
             pgb.write_message(&BeMessage::RowDescription(&[
                 RowDescriptor::int8_col(b"checkpoint_distance"),
                 RowDescriptor::int8_col(b"checkpoint_timeout"),
@@ -1050,25 +1004,12 @@ impl postgres_backend_async::Handler for PageServerHandler {
 /// If the tenant is Loading, waits for it to become Active, for up to 30 s. That
 /// ensures that queries don't fail immediately after pageserver startup, because
 /// all tenants are still loading.
-async fn get_active_tenant_with_timeout(
-    tenant_id: TenantId,
-    parent_ctx: &RequestContext,
-) -> anyhow::Result<(Arc<Tenant>, TenantRequestContext)> {
-    let child_ctx = RequestContext::with_parent(
-        parent_ctx.task_kind(),
-        parent_ctx.download_behavior(),
-        parent_ctx,
-    );
-
-    let tenant = mgr::get_tenant(tenant_id).await?;
-    match tokio::time::timeout(
-        Duration::from_secs(30),
-        tenant.wait_to_become_active(child_ctx),
-    )
-    .await
-    {
-        Ok(Ok(ctx)) => Ok((tenant, ctx)),
-        Ok(Err(err)) => Err(err),
+async fn get_active_tenant_with_timeout(tenant_id: TenantId) -> anyhow::Result<Arc<Tenant>> {
+    let tenant = mgr::get_tenant(tenant_id, false).await?;
+    match tokio::time::timeout(Duration::from_secs(30), tenant.wait_to_become_active()).await {
+        Ok(wait_result) => wait_result
+            // no .context(), the error message is good enough and some tests depend on it
+            .map(move |()| tenant),
         Err(_) => anyhow::bail!("Timeout waiting for tenant {tenant_id} to become Active"),
     }
 }
@@ -1077,9 +1018,8 @@ async fn get_active_tenant_with_timeout(
 async fn get_active_timeline_with_timeout(
     tenant_id: TenantId,
     timeline_id: TimelineId,
-    ctx: &RequestContext,
-) -> anyhow::Result<(Arc<Timeline>, TimelineRequestContext)> {
-    get_active_tenant_with_timeout(tenant_id, ctx)
+) -> anyhow::Result<Arc<Timeline>> {
+    get_active_tenant_with_timeout(tenant_id)
         .await
-        .and_then(|(tenant, ctx)| tenant.get_active_timeline(timeline_id, &ctx))
+        .and_then(|tenant| tenant.get_timeline(timeline_id, true))
 }

pageserver/src/pgdatadir_mapping.rs

@@ -6,10 +6,11 @@
 //! walingest.rs handles a few things like implicit relation creation and extension.
 //! Clarify that)
 //!
+use super::tenant::PageReconstructResult;
 use crate::keyspace::{KeySpace, KeySpaceAccum};
-use crate::repository::*;
-use crate::tenant::{PageReconstructError, Timeline, TimelineRequestContext};
+use crate::tenant::{with_ondemand_download, Timeline};
 use crate::walrecord::NeonWalRecord;
+use crate::{repository::*, try_no_ondemand_download};
 use anyhow::Context;
 use bytes::{Buf, Bytes};
 use pageserver_api::reltag::{RelTag, SlruKind};
@@ -19,6 +20,7 @@ use postgres_ffi::{Oid, TimestampTz, TransactionId};
 use serde::{Deserialize, Serialize};
 use std::collections::{hash_map, HashMap, HashSet};
 use std::ops::Range;
+use tokio_util::sync::CancellationToken;
 use tracing::{debug, trace, warn};
 use utils::{bin_ser::BeSer, lsn::Lsn};
 
@@ -33,6 +35,14 @@ pub enum LsnForTimestamp {
     NoData(Lsn),
 }
 
+#[derive(Debug, thiserror::Error)]
+pub enum CalculateLogicalSizeError {
+    #[error("cancelled")]
+    Cancelled,
+    #[error(transparent)]
+    Other(#[from] anyhow::Error),
+}
+
 ///
 /// This impl provides all the functionality to store PostgreSQL relations, SLRUs,
 /// and other special kinds of files, in a versioned key-value store. The
@@ -82,83 +92,76 @@ impl Timeline {
     //------------------------------------------------------------------------------
 
     /// Look up given page version.
-    pub async fn get_rel_page_at_lsn(
+    pub fn get_rel_page_at_lsn(
         &self,
         tag: RelTag,
         blknum: BlockNumber,
         lsn: Lsn,
         latest: bool,
-        ctx: &TimelineRequestContext,
-    ) -> Result<Bytes, PageReconstructError> {
+    ) -> PageReconstructResult<Bytes> {
         if tag.relnode == 0 {
-            return Err(PageReconstructError::Other(anyhow::anyhow!(
-                "invalid relnode"
-            )));
+            return PageReconstructResult::from(anyhow::anyhow!("invalid relnode"));
         }
 
-        let nblocks = self.get_rel_size(tag, lsn, latest, ctx).await?;
+        let nblocks = try_no_ondemand_download!(self.get_rel_size(tag, lsn, latest));
         if blknum >= nblocks {
             debug!(
                 "read beyond EOF at {} blk {} at {}, size is {}: returning all-zeros page",
                 tag, blknum, lsn, nblocks
             );
-            return Ok(ZERO_PAGE.clone());
+            return PageReconstructResult::Success(ZERO_PAGE.clone());
         }
 
         let key = rel_block_to_key(tag, blknum);
-        self.get(key, lsn, ctx).await
+        self.get(key, lsn)
     }
 
     // Get size of a database in blocks
-    pub async fn get_db_size(
+    pub fn get_db_size(
         &self,
         spcnode: Oid,
         dbnode: Oid,
         lsn: Lsn,
         latest: bool,
-        ctx: &TimelineRequestContext,
-    ) -> Result<usize, PageReconstructError> {
+    ) -> PageReconstructResult<usize> {
         let mut total_blocks = 0;
 
-        let rels = self.list_rels(spcnode, dbnode, lsn, ctx).await?;
+        let rels = try_no_ondemand_download!(self.list_rels(spcnode, dbnode, lsn));
 
         for rel in rels {
-            let n_blocks = self.get_rel_size(rel, lsn, latest, ctx).await?;
+            let n_blocks = try_no_ondemand_download!(self.get_rel_size(rel, lsn, latest));
             total_blocks += n_blocks as usize;
         }
-        Ok(total_blocks)
+        PageReconstructResult::Success(total_blocks)
     }
 
     /// Get size of a relation file
-    pub async fn get_rel_size(
+    pub fn get_rel_size(
         &self,
         tag: RelTag,
         lsn: Lsn,
         latest: bool,
-        ctx: &TimelineRequestContext,
-    ) -> Result<BlockNumber, PageReconstructError> {
+    ) -> PageReconstructResult<BlockNumber> {
         if tag.relnode == 0 {
-            return Err(PageReconstructError::Other(anyhow::anyhow!(
-                "invalid relnode"
-            )));
+            return PageReconstructResult::from(anyhow::anyhow!("invalid relnode"));
         }
 
         if let Some(nblocks) = self.get_cached_rel_size(&tag, lsn) {
-            return Ok(nblocks);
+            return PageReconstructResult::Success(nblocks);
         }
 
         if (tag.forknum == FSM_FORKNUM || tag.forknum == VISIBILITYMAP_FORKNUM)
-            && !self.get_rel_exists(tag, lsn, latest, ctx).await?
+            && !try_no_ondemand_download!(self.get_rel_exists(tag, lsn, latest))
         {
             // FIXME: Postgres sometimes calls smgrcreate() to create
             // FSM, and smgrnblocks() on it immediately afterwards,
             // without extending it.  Tolerate that by claiming that
             // any non-existent FSM fork has size 0.
-            return Ok(0);
+            return PageReconstructResult::Success(0);
         }
 
         let key = rel_size_to_key(tag);
-        let mut buf = self.get(key, lsn, ctx).await?;
+        let mut buf = try_no_ondemand_download!(self.get(key, lsn));
         let nblocks = buf.get_u32_le();
 
         if latest {
@@ -171,51 +174,47 @@ impl Timeline {
             // associated with most recent value of LSN.
             self.update_cached_rel_size(tag, lsn, nblocks);
         }
-        Ok(nblocks)
+        PageReconstructResult::Success(nblocks)
     }
 
     /// Does relation exist?
-    pub async fn get_rel_exists(
+    pub fn get_rel_exists(
         &self,
         tag: RelTag,
         lsn: Lsn,
         _latest: bool,
-        ctx: &TimelineRequestContext,
-    ) -> Result<bool, PageReconstructError> {
+    ) -> PageReconstructResult<bool> {
         if tag.relnode == 0 {
-            return Err(PageReconstructError::Other(anyhow::anyhow!(
-                "invalid relnode"
-            )));
+            return PageReconstructResult::from(anyhow::anyhow!("invalid relnode"));
         }
 
         // first try to lookup relation in cache
         if let Some(_nblocks) = self.get_cached_rel_size(&tag, lsn) {
-            return Ok(true);
+            return PageReconstructResult::Success(true);
         }
         // fetch directory listing
         let key = rel_dir_to_key(tag.spcnode, tag.dbnode);
-        let buf = self.get(key, lsn, ctx).await?;
+        let buf = try_no_ondemand_download!(self.get(key, lsn));
 
         match RelDirectory::des(&buf).context("deserialization failure") {
             Ok(dir) => {
                 let exists = dir.rels.get(&(tag.relnode, tag.forknum)).is_some();
-                Ok(exists)
+                PageReconstructResult::Success(exists)
             }
-            Err(e) => Err(PageReconstructError::from(e)),
+            Err(e) => PageReconstructResult::from(e),
         }
     }
 
     /// Get a list of all existing relations in given tablespace and database.
-    pub async fn list_rels(
+    pub fn list_rels(
         &self,
         spcnode: Oid,
         dbnode: Oid,
         lsn: Lsn,
-        ctx: &TimelineRequestContext,
-    ) -> Result<HashSet<RelTag>, PageReconstructError> {
+    ) -> PageReconstructResult<HashSet<RelTag>> {
         // fetch directory listing
         let key = rel_dir_to_key(spcnode, dbnode);
-        let buf = self.get(key, lsn, ctx).await?;
+        let buf = try_no_ondemand_download!(self.get(key, lsn));
 
         match RelDirectory::des(&buf).context("deserialization failure") {
             Ok(dir) => {
@@ -227,56 +226,53 @@ impl Timeline {
                         forknum: *forknum,
                     }));
 
-                Ok(rels)
+                PageReconstructResult::Success(rels)
             }
-            Err(e) => Err(PageReconstructError::from(e)),
+            Err(e) => PageReconstructResult::from(e),
         }
     }
 
     /// Look up given SLRU page version.
-    pub async fn get_slru_page_at_lsn(
+    pub fn get_slru_page_at_lsn(
         &self,
         kind: SlruKind,
         segno: u32,
         blknum: BlockNumber,
         lsn: Lsn,
-        ctx: &TimelineRequestContext,
-    ) -> Result<Bytes, PageReconstructError> {
+    ) -> PageReconstructResult<Bytes> {
         let key = slru_block_to_key(kind, segno, blknum);
-        self.get(key, lsn, ctx).await
+        self.get(key, lsn)
     }
 
     /// Get size of an SLRU segment
-    pub async fn get_slru_segment_size(
+    pub fn get_slru_segment_size(
         &self,
         kind: SlruKind,
         segno: u32,
         lsn: Lsn,
-        ctx: &TimelineRequestContext,
-    ) -> Result<BlockNumber, PageReconstructError> {
+    ) -> PageReconstructResult<BlockNumber> {
         let key = slru_segment_size_to_key(kind, segno);
-        let mut buf = self.get(key, lsn, ctx).await?;
-        Ok(buf.get_u32_le())
+        let mut buf = try_no_ondemand_download!(self.get(key, lsn));
+        PageReconstructResult::Success(buf.get_u32_le())
     }
 
     /// Get size of an SLRU segment
-    pub async fn get_slru_segment_exists(
+    pub fn get_slru_segment_exists(
         &self,
         kind: SlruKind,
         segno: u32,
         lsn: Lsn,
-        ctx: &TimelineRequestContext,
-    ) -> Result<bool, PageReconstructError> {
+    ) -> PageReconstructResult<bool> {
         // fetch directory listing
         let key = slru_dir_to_key(kind);
-        let buf = self.get(key, lsn, ctx).await?;
+        let buf = try_no_ondemand_download!(self.get(key, lsn));
 
         match SlruSegmentDirectory::des(&buf).context("deserialization failure") {
             Ok(dir) => {
                 let exists = dir.segments.get(&segno).is_some();
-                Ok(exists)
+                PageReconstructResult::Success(exists)
             }
-            Err(e) => Err(PageReconstructError::from(e)),
+            Err(e) => PageReconstructResult::from(e),
         }
     }
 
@@ -287,11 +283,10 @@ impl Timeline {
     /// so it's not well defined which LSN you get if there were multiple commits
     /// "in flight" at that point in time.
     ///
-    pub async fn find_lsn_for_timestamp(
+    pub fn find_lsn_for_timestamp(
         &self,
         search_timestamp: TimestampTz,
-        ctx: &TimelineRequestContext,
-    ) -> Result<LsnForTimestamp, PageReconstructError> {
+    ) -> PageReconstructResult<LsnForTimestamp> {
         let gc_cutoff_lsn_guard = self.get_latest_gc_cutoff_lsn();
         let min_lsn = *gc_cutoff_lsn_guard;
         let max_lsn = self.get_last_record_lsn();
@@ -307,15 +302,12 @@ impl Timeline {
             // cannot overflow, high and low are both smaller than u64::MAX / 2
             let mid = (high + low) / 2;
 
-            let cmp = self
-                .is_latest_commit_timestamp_ge_than(
-                    search_timestamp,
-                    Lsn(mid * 8),
-                    &mut found_smaller,
-                    &mut found_larger,
-                    ctx,
-                )
-                .await?;
+            let cmp = try_no_ondemand_download!(self.is_latest_commit_timestamp_ge_than(
+                search_timestamp,
+                Lsn(mid * 8),
+                &mut found_smaller,
+                &mut found_larger,
+            ));
 
             if cmp {
                 high = mid;
@@ -327,15 +319,15 @@ impl Timeline {
             (false, false) => {
                 // This can happen if no commit records have been processed yet, e.g.
                 // just after importing a cluster.
-                Ok(LsnForTimestamp::NoData(max_lsn))
+                PageReconstructResult::Success(LsnForTimestamp::NoData(max_lsn))
             }
             (true, false) => {
                 // Didn't find any commit timestamps larger than the request
-                Ok(LsnForTimestamp::Future(max_lsn))
+                PageReconstructResult::Success(LsnForTimestamp::Future(max_lsn))
             }
             (false, true) => {
                 // Didn't find any commit timestamps smaller than the request
-                Ok(LsnForTimestamp::Past(max_lsn))
+                PageReconstructResult::Success(LsnForTimestamp::Past(max_lsn))
             }
             (true, true) => {
                 // low is the LSN of the first commit record *after* the search_timestamp,
@@ -345,7 +337,7 @@ impl Timeline {
                 // Otherwise, if you restore to the returned LSN, the database will
                 // include physical changes from later commits that will be marked
                 // as aborted, and will need to be vacuumed away.
-                Ok(LsnForTimestamp::Present(Lsn((low - 1) * 8)))
+                PageReconstructResult::Success(LsnForTimestamp::Present(Lsn((low - 1) * 8)))
             }
         }
     }
@@ -357,25 +349,26 @@ impl Timeline {
     /// Additionally, sets 'found_smaller'/'found_Larger, if encounters any commits
     /// with a smaller/larger timestamp.
     ///
-    pub async fn is_latest_commit_timestamp_ge_than(
+    pub fn is_latest_commit_timestamp_ge_than(
         &self,
         search_timestamp: TimestampTz,
         probe_lsn: Lsn,
         found_smaller: &mut bool,
         found_larger: &mut bool,
-        ctx: &TimelineRequestContext,
-    ) -> Result<bool, PageReconstructError> {
-        for segno in self
-            .list_slru_segments(SlruKind::Clog, probe_lsn, ctx)
-            .await?
-        {
-            let nblocks = self
-                .get_slru_segment_size(SlruKind::Clog, segno, probe_lsn, ctx)
-                .await?;
+    ) -> PageReconstructResult<bool> {
+        for segno in try_no_ondemand_download!(self.list_slru_segments(SlruKind::Clog, probe_lsn)) {
+            let nblocks = try_no_ondemand_download!(self.get_slru_segment_size(
+                SlruKind::Clog,
+                segno,
+                probe_lsn
+            ));
             for blknum in (0..nblocks).rev() {
-                let clog_page = self
-                    .get_slru_page_at_lsn(SlruKind::Clog, segno, blknum, probe_lsn, ctx)
-                    .await?;
+                let clog_page = try_no_ondemand_download!(self.get_slru_page_at_lsn(
+                    SlruKind::Clog,
+                    segno,
+                    blknum,
+                    probe_lsn
+                ));
 
                 if clog_page.len() == BLCKSZ as usize + 8 {
                     let mut timestamp_bytes = [0u8; 8];
@@ -384,99 +377,76 @@ impl Timeline {
 
                     if timestamp >= search_timestamp {
                         *found_larger = true;
-                        return Ok(true);
+                        return PageReconstructResult::Success(true);
                     } else {
                         *found_smaller = true;
                     }
                 }
             }
         }
-        Ok(false)
+        PageReconstructResult::Success(false)
     }
 
     /// Get a list of SLRU segments
-    pub async fn list_slru_segments(
+    pub fn list_slru_segments(
         &self,
         kind: SlruKind,
         lsn: Lsn,
-        ctx: &TimelineRequestContext,
-    ) -> Result<HashSet<u32>, PageReconstructError> {
+    ) -> PageReconstructResult<HashSet<u32>> {
         // fetch directory entry
         let key = slru_dir_to_key(kind);
 
-        let buf = self.get(key, lsn, ctx).await?;
+        let buf = try_no_ondemand_download!(self.get(key, lsn));
         match SlruSegmentDirectory::des(&buf).context("deserialization failure") {
-            Ok(dir) => Ok(dir.segments),
-            Err(e) => Err(PageReconstructError::from(e)),
+            Ok(dir) => PageReconstructResult::Success(dir.segments),
+            Err(e) => PageReconstructResult::from(e),
         }
     }
 
-    pub async fn get_relmap_file(
+    pub fn get_relmap_file(
         &self,
         spcnode: Oid,
         dbnode: Oid,
         lsn: Lsn,
-        ctx: &TimelineRequestContext,
-    ) -> Result<Bytes, PageReconstructError> {
+    ) -> PageReconstructResult<Bytes> {
         let key = relmap_file_key(spcnode, dbnode);
 
-        let buf = self.get(key, lsn, ctx).await?;
-        Ok(buf)
+        let buf = try_no_ondemand_download!(self.get(key, lsn));
+        PageReconstructResult::Success(buf)
     }
 
-    pub async fn list_dbdirs(
-        &self,
-        lsn: Lsn,
-        ctx: &TimelineRequestContext,
-    ) -> Result<HashMap<(Oid, Oid), bool>, PageReconstructError> {
+    pub fn list_dbdirs(&self, lsn: Lsn) -> PageReconstructResult<HashMap<(Oid, Oid), bool>> {
         // fetch directory entry
-        let buf = self.get(DBDIR_KEY, lsn, ctx).await?;
+        let buf = try_no_ondemand_download!(self.get(DBDIR_KEY, lsn));
 
         match DbDirectory::des(&buf).context("deserialization failure") {
-            Ok(dir) => Ok(dir.dbdirs),
-            Err(e) => Err(PageReconstructError::from(e)),
+            Ok(dir) => PageReconstructResult::Success(dir.dbdirs),
+            Err(e) => PageReconstructResult::from(e),
         }
     }
 
-    pub async fn get_twophase_file(
-        &self,
-        xid: TransactionId,
-        lsn: Lsn,
-        ctx: &TimelineRequestContext,
-    ) -> Result<Bytes, PageReconstructError> {
+    pub fn get_twophase_file(&self, xid: TransactionId, lsn: Lsn) -> PageReconstructResult<Bytes> {
         let key = twophase_file_key(xid);
-        let buf = self.get(key, lsn, ctx).await?;
-        Ok(buf)
+        let buf = try_no_ondemand_download!(self.get(key, lsn));
+        PageReconstructResult::Success(buf)
     }
 
-    pub async fn list_twophase_files(
-        &self,
-        lsn: Lsn,
-        ctx: &TimelineRequestContext,
-    ) -> Result<HashSet<TransactionId>, PageReconstructError> {
+    pub fn list_twophase_files(&self, lsn: Lsn) -> PageReconstructResult<HashSet<TransactionId>> {
         // fetch directory entry
-        let buf = self.get(TWOPHASEDIR_KEY, lsn, ctx).await?;
+        let buf = try_no_ondemand_download!(self.get(TWOPHASEDIR_KEY, lsn));
 
         match TwoPhaseDirectory::des(&buf).context("deserialization failure") {
-            Ok(dir) => Ok(dir.xids),
-            Err(e) => Err(PageReconstructError::from(e)),
+            Ok(dir) => PageReconstructResult::Success(dir.xids),
+            Err(e) => PageReconstructResult::from(e),
         }
     }
 
-    pub async fn get_control_file(
-        &self,
-        lsn: Lsn,
-        ctx: &TimelineRequestContext,
-    ) -> Result<Bytes, PageReconstructError> {
-        self.get(CONTROLFILE_KEY, lsn, ctx).await
+    pub fn get_control_file(&self, lsn: Lsn) -> PageReconstructResult<Bytes> {
+        self.get(CONTROLFILE_KEY, lsn)
     }
 
-    pub async fn get_checkpoint(
-        &self,
-        lsn: Lsn,
-        ctx: &TimelineRequestContext,
-    ) -> Result<Bytes, PageReconstructError> {
-        self.get(CHECKPOINT_KEY, lsn, ctx).await
+    pub fn get_checkpoint(&self, lsn: Lsn) -> PageReconstructResult<Bytes> {
+        self.get(CHECKPOINT_KEY, lsn)
     }
 
     /// Does the same as get_current_logical_size but counted on demand.
@@ -487,20 +457,23 @@ impl Timeline {
     pub async fn get_current_logical_size_non_incremental(
         &self,
         lsn: Lsn,
-        ctx: &TimelineRequestContext,
-    ) -> Result<u64, PageReconstructError> {
+        cancel: CancellationToken,
+    ) -> Result<u64, CalculateLogicalSizeError> {
         // Fetch list of database dirs and iterate them
-        let buf = self.get(DBDIR_KEY, lsn, ctx).await?;
+        let buf = self.get_download(DBDIR_KEY, lsn).await?;
         let dbdir = DbDirectory::des(&buf).context("deserialize db directory")?;
 
         let mut total_size: u64 = 0;
         for (spcnode, dbnode) in dbdir.dbdirs.keys() {
-            for rel in self.list_rels(*spcnode, *dbnode, lsn, ctx).await? {
-                if ctx.is_cancelled() {
-                    return Err(PageReconstructError::Cancelled);
+            for rel in
+                crate::tenant::with_ondemand_download(|| self.list_rels(*spcnode, *dbnode, lsn))
+                    .await?
+            {
+                if cancel.is_cancelled() {
+                    return Err(CalculateLogicalSizeError::Cancelled);
                 }
                 let relsize_key = rel_size_to_key(rel);
-                let mut buf = self.get(relsize_key, lsn, ctx).await?;
+                let mut buf = self.get_download(relsize_key, lsn).await?;
                 let relsize = buf.get_u32_le();
 
                 total_size += relsize as u64;
@@ -513,11 +486,7 @@ impl Timeline {
     /// Get a KeySpace that covers all the Keys that are in use at the given LSN.
     /// Anything that's not listed maybe removed from the underlying storage (from
     /// that LSN forwards).
-    pub async fn collect_keyspace(
-        &self,
-        lsn: Lsn,
-        ctx: &TimelineRequestContext,
-    ) -> anyhow::Result<KeySpace> {
+    pub async fn collect_keyspace(&self, lsn: Lsn) -> anyhow::Result<KeySpace> {
         // Iterate through key ranges, greedily packing them into partitions
         let mut result = KeySpaceAccum::new();
 
@@ -525,7 +494,7 @@ impl Timeline {
         result.add_key(DBDIR_KEY);
 
         // Fetch list of database dirs and iterate them
-        let buf = self.get(DBDIR_KEY, lsn, ctx).await?;
+        let buf = self.get_download(DBDIR_KEY, lsn).await?;
         let dbdir = DbDirectory::des(&buf).context("deserialization failure")?;
 
         let mut dbs: Vec<(Oid, Oid)> = dbdir.dbdirs.keys().cloned().collect();
@@ -534,16 +503,15 @@ impl Timeline {
             result.add_key(relmap_file_key(spcnode, dbnode));
             result.add_key(rel_dir_to_key(spcnode, dbnode));
 
-            let mut rels: Vec<RelTag> = self
-                .list_rels(spcnode, dbnode, lsn, ctx)
-                .await?
-                .iter()
-                .cloned()
-                .collect();
+            let mut rels: Vec<RelTag> =
+                with_ondemand_download(|| self.list_rels(spcnode, dbnode, lsn))
+                    .await?
+                    .into_iter()
+                    .collect();
             rels.sort_unstable();
             for rel in rels {
                 let relsize_key = rel_size_to_key(rel);
-                let mut buf = self.get(relsize_key, lsn, ctx).await?;
+                let mut buf = self.get_download(relsize_key, lsn).await?;
                 let relsize = buf.get_u32_le();
 
                 result.add_range(rel_block_to_key(rel, 0)..rel_block_to_key(rel, relsize));
@@ -559,13 +527,13 @@ impl Timeline {
         ] {
             let slrudir_key = slru_dir_to_key(kind);
             result.add_key(slrudir_key);
-            let buf = self.get(slrudir_key, lsn, ctx).await?;
+            let buf = self.get_download(slrudir_key, lsn).await?;
             let dir = SlruSegmentDirectory::des(&buf).context("deserialization failure")?;
             let mut segments: Vec<u32> = dir.segments.iter().cloned().collect();
             segments.sort_unstable();
             for segno in segments {
                 let segsize_key = slru_segment_size_to_key(kind, segno);
-                let mut buf = self.get(segsize_key, lsn, ctx).await?;
+                let mut buf = self.get_download(segsize_key, lsn).await?;
                 let segsize = buf.get_u32_le();
 
                 result.add_range(
@@ -577,7 +545,7 @@ impl Timeline {
 
         // Then pg_twophase
         result.add_key(TWOPHASEDIR_KEY);
-        let buf = self.get(TWOPHASEDIR_KEY, lsn, ctx).await?;
+        let buf = self.get_download(TWOPHASEDIR_KEY, lsn).await?;
         let twophase_dir = TwoPhaseDirectory::des(&buf).context("deserialization failure")?;
         let mut xids: Vec<TransactionId> = twophase_dir.xids.iter().cloned().collect();
         xids.sort_unstable();
@@ -735,15 +703,9 @@ impl<'a> DatadirModification<'a> {
     }
 
     /// Store a relmapper file (pg_filenode.map) in the repository
-    pub async fn put_relmap_file(
-        &mut self,
-        spcnode: Oid,
-        dbnode: Oid,
-        img: Bytes,
-        ctx: &TimelineRequestContext,
-    ) -> anyhow::Result<()> {
+    pub fn put_relmap_file(&mut self, spcnode: Oid, dbnode: Oid, img: Bytes) -> anyhow::Result<()> {
         // Add it to the directory (if it doesn't exist already)
-        let buf = self.get(DBDIR_KEY, ctx).await?;
+        let buf = self.get(DBDIR_KEY).no_ondemand_download()?;
         let mut dbdir = DbDirectory::des(&buf)?;
 
         let r = dbdir.dbdirs.insert((spcnode, dbnode), true);
@@ -769,14 +731,9 @@ impl<'a> DatadirModification<'a> {
         Ok(())
     }
 
-    pub async fn put_twophase_file(
-        &mut self,
-        xid: TransactionId,
-        img: Bytes,
-        ctx: &TimelineRequestContext,
-    ) -> anyhow::Result<()> {
+    pub fn put_twophase_file(&mut self, xid: TransactionId, img: Bytes) -> anyhow::Result<()> {
         // Add it to the directory entry
-        let buf = self.get(TWOPHASEDIR_KEY, ctx).await?;
+        let buf = self.get(TWOPHASEDIR_KEY).no_ondemand_download()?;
         let mut dir = TwoPhaseDirectory::des(&buf)?;
         if !dir.xids.insert(xid) {
             anyhow::bail!("twophase file for xid {} already exists", xid);
@@ -800,21 +757,16 @@ impl<'a> DatadirModification<'a> {
         Ok(())
     }
 
-    pub async fn drop_dbdir(
-        &mut self,
-        spcnode: Oid,
-        dbnode: Oid,
-        ctx: &TimelineRequestContext,
-    ) -> anyhow::Result<()> {
+    pub fn drop_dbdir(&mut self, spcnode: Oid, dbnode: Oid) -> anyhow::Result<()> {
         let req_lsn = self.tline.get_last_record_lsn();
 
         let total_blocks = self
             .tline
-            .get_db_size(spcnode, dbnode, req_lsn, true, ctx)
-            .await?;
+            .get_db_size(spcnode, dbnode, req_lsn, true)
+            .no_ondemand_download()?;
 
         // Remove entry from dbdir
-        let buf = self.get(DBDIR_KEY, ctx).await?;
+        let buf = self.get(DBDIR_KEY).no_ondemand_download()?;
         let mut dir = DbDirectory::des(&buf)?;
         if dir.dbdirs.remove(&(spcnode, dbnode)).is_some() {
             let buf = DbDirectory::ser(&dir)?;
@@ -837,16 +789,11 @@ impl<'a> DatadirModification<'a> {
     /// Create a relation fork.
     ///
     /// 'nblocks' is the initial size.
-    pub async fn put_rel_creation(
-        &mut self,
-        rel: RelTag,
-        nblocks: BlockNumber,
-        ctx: &TimelineRequestContext,
-    ) -> anyhow::Result<()> {
+    pub fn put_rel_creation(&mut self, rel: RelTag, nblocks: BlockNumber) -> anyhow::Result<()> {
         anyhow::ensure!(rel.relnode != 0, "invalid relnode");
         // It's possible that this is the first rel for this db in this
         // tablespace.  Create the reldir entry for it if so.
-        let mut dbdir = DbDirectory::des(&self.get(DBDIR_KEY, ctx).await?)?;
+        let mut dbdir = DbDirectory::des(&self.get(DBDIR_KEY).no_ondemand_download()?)?;
         let rel_dir_key = rel_dir_to_key(rel.spcnode, rel.dbnode);
         let mut rel_dir = if dbdir.dbdirs.get(&(rel.spcnode, rel.dbnode)).is_none() {
             // Didn't exist. Update dbdir
@@ -858,7 +805,7 @@ impl<'a> DatadirModification<'a> {
             RelDirectory::default()
         } else {
             // reldir already exists, fetch it
-            RelDirectory::des(&self.get(rel_dir_key, ctx).await?)?
+            RelDirectory::des(&self.get(rel_dir_key).no_ondemand_download()?)?
         };
 
         // Add the new relation to the rel directory entry, and write it back
@@ -886,18 +833,17 @@ impl<'a> DatadirModification<'a> {
     }
 
     /// Truncate relation
-    pub async fn put_rel_truncation(
-        &mut self,
-        rel: RelTag,
-        nblocks: BlockNumber,
-        ctx: &TimelineRequestContext,
-    ) -> anyhow::Result<()> {
+    pub fn put_rel_truncation(&mut self, rel: RelTag, nblocks: BlockNumber) -> anyhow::Result<()> {
         anyhow::ensure!(rel.relnode != 0, "invalid relnode");
         let last_lsn = self.tline.get_last_record_lsn();
-        if self.tline.get_rel_exists(rel, last_lsn, true, ctx).await? {
+        if self
+            .tline
+            .get_rel_exists(rel, last_lsn, true)
+            .no_ondemand_download()?
+        {
             let size_key = rel_size_to_key(rel);
             // Fetch the old size first
-            let old_size = self.get(size_key, ctx).await?.get_u32_le();
+            let old_size = self.get(size_key).no_ondemand_download()?.get_u32_le();
 
             // Update the entry with the new size.
             let buf = nblocks.to_le_bytes();
@@ -917,17 +863,12 @@ impl<'a> DatadirModification<'a> {
 
     /// Extend relation
     /// If new size is smaller, do nothing.
-    pub async fn put_rel_extend(
-        &mut self,
-        rel: RelTag,
-        nblocks: BlockNumber,
-        ctx: &TimelineRequestContext,
-    ) -> anyhow::Result<()> {
+    pub fn put_rel_extend(&mut self, rel: RelTag, nblocks: BlockNumber) -> anyhow::Result<()> {
         anyhow::ensure!(rel.relnode != 0, "invalid relnode");
 
         // Put size
         let size_key = rel_size_to_key(rel);
-        let old_size = self.get(size_key, ctx).await?.get_u32_le();
+        let old_size = self.get(size_key).no_ondemand_download()?.get_u32_le();
 
         // only extend relation here. never decrease the size
         if nblocks > old_size {
@@ -943,16 +884,12 @@ impl<'a> DatadirModification<'a> {
     }
 
     /// Drop a relation.
-    pub async fn put_rel_drop(
-        &mut self,
-        rel: RelTag,
-        ctx: &TimelineRequestContext,
-    ) -> anyhow::Result<()> {
+    pub fn put_rel_drop(&mut self, rel: RelTag) -> anyhow::Result<()> {
         anyhow::ensure!(rel.relnode != 0, "invalid relnode");
 
         // Remove it from the directory entry
         let dir_key = rel_dir_to_key(rel.spcnode, rel.dbnode);
-        let buf = self.get(dir_key, ctx).await?;
+        let buf = self.get(dir_key).no_ondemand_download()?;
         let mut dir = RelDirectory::des(&buf)?;
 
         if dir.rels.remove(&(rel.relnode, rel.forknum)) {
@@ -963,7 +900,7 @@ impl<'a> DatadirModification<'a> {
 
         // update logical size
         let size_key = rel_size_to_key(rel);
-        let old_size = self.get(size_key, ctx).await?.get_u32_le();
+        let old_size = self.get(size_key).no_ondemand_download()?.get_u32_le();
         self.pending_nblocks -= old_size as i64;
 
         // Remove enty from relation size cache
@@ -975,16 +912,15 @@ impl<'a> DatadirModification<'a> {
         Ok(())
     }
 
-    pub async fn put_slru_segment_creation(
+    pub fn put_slru_segment_creation(
         &mut self,
         kind: SlruKind,
         segno: u32,
         nblocks: BlockNumber,
-        ctx: &TimelineRequestContext,
     ) -> anyhow::Result<()> {
         // Add it to the directory entry
         let dir_key = slru_dir_to_key(kind);
-        let buf = self.get(dir_key, ctx).await?;
+        let buf = self.get(dir_key).no_ondemand_download()?;
         let mut dir = SlruSegmentDirectory::des(&buf)?;
 
         if !dir.segments.insert(segno) {
@@ -1020,15 +956,10 @@ impl<'a> DatadirModification<'a> {
     }
 
     /// This method is used for marking truncated SLRU files
-    pub async fn drop_slru_segment(
-        &mut self,
-        kind: SlruKind,
-        segno: u32,
-        ctx: &TimelineRequestContext,
-    ) -> anyhow::Result<()> {
+    pub fn drop_slru_segment(&mut self, kind: SlruKind, segno: u32) -> anyhow::Result<()> {
         // Remove it from the directory entry
         let dir_key = slru_dir_to_key(kind);
-        let buf = self.get(dir_key, ctx).await?;
+        let buf = self.get(dir_key).no_ondemand_download()?;
         let mut dir = SlruSegmentDirectory::des(&buf)?;
 
         if !dir.segments.remove(&segno) {
@@ -1052,13 +983,9 @@ impl<'a> DatadirModification<'a> {
     }
 
     /// This method is used for marking truncated SLRU files
-    pub async fn drop_twophase_file(
-        &mut self,
-        xid: TransactionId,
-        ctx: &TimelineRequestContext,
-    ) -> anyhow::Result<()> {
+    pub fn drop_twophase_file(&mut self, xid: TransactionId) -> anyhow::Result<()> {
         // Remove it from the directory entry
-        let buf = self.get(TWOPHASEDIR_KEY, ctx).await?;
+        let buf = self.get(TWOPHASEDIR_KEY).no_ondemand_download()?;
         let mut dir = TwoPhaseDirectory::des(&buf)?;
 
         if !dir.xids.remove(&xid) {
@@ -1152,11 +1079,7 @@ impl<'a> DatadirModification<'a> {
 
     // Internal helper functions to batch the modifications
 
-    async fn get(
-        &self,
-        key: Key,
-        ctx: &TimelineRequestContext,
-    ) -> Result<Bytes, PageReconstructError> {
+    fn get(&self, key: Key) -> PageReconstructResult<Bytes> {
         // Have we already updated the same key? Read the pending updated
         // version in that case.
         //
@@ -1164,20 +1087,18 @@ impl<'a> DatadirModification<'a> {
         // value that has been removed, deletion only avoids leaking storage.
         if let Some(value) = self.pending_updates.get(&key) {
             if let Value::Image(img) = value {
-                Ok(img.clone())
+                PageReconstructResult::Success(img.clone())
             } else {
                 // Currently, we never need to read back a WAL record that we
                 // inserted in the same "transaction". All the metadata updates
                 // work directly with Images, and we never need to read actual
                 // data pages. We could handle this if we had to, by calling
                 // the walredo manager, but let's keep it simple for now.
-                Err(PageReconstructError::from(anyhow::anyhow!(
-                    "unexpected pending WAL record"
-                )))
+                PageReconstructResult::from(anyhow::anyhow!("unexpected pending WAL record"))
             }
         } else {
             let lsn = Lsn::max(self.tline.get_last_record_lsn(), self.lsn);
-            self.tline.get(key, lsn, ctx).await
+            self.tline.get(key, lsn)
         }
     }
 
@@ -1584,18 +1505,17 @@ fn is_slru_block_key(key: Key) -> bool {
 
 #[cfg(test)]
 pub fn create_test_timeline(
-    tenant: &std::sync::Arc<crate::tenant::Tenant>,
+    tenant: &crate::tenant::Tenant,
     timeline_id: utils::id::TimelineId,
     pg_version: u32,
-    tenant_ctx: &crate::tenant::TenantRequestContext,
-) -> anyhow::Result<(std::sync::Arc<Timeline>, TimelineRequestContext)> {
-    let (tline, timeline_ctx) =
-        tenant.create_empty_timeline(timeline_id, Lsn(8), pg_version, tenant_ctx)?;
-    let tline = tline.initialize(&timeline_ctx)?;
+) -> anyhow::Result<std::sync::Arc<Timeline>> {
+    let tline = tenant
+        .create_empty_timeline(timeline_id, Lsn(8), pg_version)?
+        .initialize()?;
     let mut m = tline.begin_modification(Lsn(8));
     m.init_empty()?;
     m.commit()?;
-    Ok((tline, timeline_ctx))
+    Ok(tline)
 }
 
 #[allow(clippy::bool_assert_comparison)]

pageserver/src/repository.rs

@@ -37,6 +37,17 @@ 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/task_mgr.rs

@@ -1,21 +1,59 @@
 //!
-//! This module provides some helpers for spawning tokio tasks in the pageserver.
+//! This module provides centralized handling of tokio tasks in the Page Server.
 //!
-//! Mostly just a wrapper around tokio::spawn, with some code to handle panics.
+//! We provide a few basic facilities:
+//! - A global registry of tasks that lists what kind of tasks they are, and
+//!   which tenant or timeline they are working on
+//!
+//! - The ability to request a task to shut down.
+//!
+//!
+//! # How it works?
+//!
+//! There is a global hashmap of all the tasks (`TASKS`). Whenever a new
+//! task is spawned, a PageServerTask entry is added there, and when a
+//! task dies, it removes itself from the hashmap. If you want to kill a
+//! task, you can scan the hashmap to find it.
+//!
+//! # Task shutdown
+//!
+//! To kill a task, we rely on co-operation from the victim. Each task is
+//! expected to periodically call the `is_shutdown_requested()` function, and
+//! if it returns true, exit gracefully. In addition to that, when waiting for
+//! the network or other long-running operation, you can use
+//! `shutdown_watcher()` function to get a Future that will become ready if
+//! the current task has been requested to shut down. You can use that with
+//! Tokio select!().
+//!
+//! TODO: This would be a good place to also handle panics in a somewhat sane way.
+//! Depending on what task panics, we might want to kill the whole server, or
+//! only a single tenant or timeline.
 //!
 
+// Clippy 1.60 incorrectly complains about the tokio::task_local!() macro.
+// Silence it. See https://github.com/rust-lang/rust-clippy/issues/9224.
+#![allow(clippy::declare_interior_mutable_const)]
+
+use std::collections::HashMap;
+use std::fmt;
 use std::future::Future;
-use std::panic::{resume_unwind, AssertUnwindSafe};
+use std::panic::AssertUnwindSafe;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
 
 use futures::FutureExt;
 use tokio::runtime::Runtime;
 use tokio::task::JoinHandle;
+use tokio::task_local;
+use tokio_util::sync::CancellationToken;
 
-use tracing::{debug, error, info};
+use tracing::{debug, error, info, warn};
 
 use once_cell::sync::Lazy;
 
-use crate::context::{self, TaskKind};
+use utils::id::{TenantId, TimelineId};
+
+use crate::shutdown_pageserver;
 
 //
 // There are four runtimes:
@@ -54,6 +92,10 @@ use crate::context::{self, TaskKind};
 // runtime. If a GetPage request comes in before the load of a tenant has finished, the
 // GetPage request will wait for the tenant load to finish.
 //
+// The core Timeline code is synchronous, and uses a bunch of std Mutexes and RWLocks to
+// protect data structures. Let's keep it that way. Synchronous code is easier to debug
+// and analyze, and there's a lot of hairy, low-level, performance critical code there.
+//
 // It's nice to have different runtimes, so that you can quickly eyeball how much CPU
 // time each class of operations is taking, with 'top -H' or similar.
 //
@@ -93,81 +135,355 @@ pub static BACKGROUND_RUNTIME: Lazy<Runtime> = Lazy::new(|| {
         .expect("Failed to create background op runtime")
 });
 
-/// Launch a new task
-///
-/// This is a wrapper around tokio::spawn. One difference is that the Future
-/// is marked to return nothing to avoid silently swallowing errors. This
-/// forces the future to handle errors by itself. If you need the return
-/// value, you could create another function that passes it through, but we
-/// don't have a need for that currently.
-///
-/// If shutdown_process_on_panic is set to true, panic of the task will lead
-/// to shutdown of entire process. Otherwise we log the panic and continue.
-pub fn spawn<F>(
-    runtime: &tokio::runtime::Handle,
-    name: &str,
-    shutdown_process_on_panic: bool,
-    future: F,
-) -> JoinHandle<F::Output>
-where
-    F: Future<Output = ()> + Send + 'static,
-{
-    let task_name = name.to_string();
-    runtime.spawn(task_wrapper(task_name, shutdown_process_on_panic, future))
+#[derive(Debug, Clone, Copy)]
+pub struct PageserverTaskId(u64);
+
+impl fmt::Display for PageserverTaskId {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
+    }
 }
 
-/// This wrapper function runs in a newly-spawned task. To handle panics.
-async fn task_wrapper<F, R>(task_name: String, shutdown_process_on_panic: bool, future: F) -> R
+/// Each task that we track is associated with a "task ID". It's just an
+/// increasing number that we assign. Note that it is different from tokio::task::Id.
+static NEXT_TASK_ID: AtomicU64 = AtomicU64::new(1);
+
+/// Global registry of tasks
+static TASKS: Lazy<Mutex<HashMap<u64, Arc<PageServerTask>>>> =
+    Lazy::new(|| Mutex::new(HashMap::new()));
+
+task_local! {
+    // This is a cancellation token which will be cancelled when a task needs to shut down. The
+    // root token is kept in the global registry, so that anyone can send the signal to request
+    // task shutdown.
+    static SHUTDOWN_TOKEN: CancellationToken;
+
+    // Each task holds reference to its own PageServerTask here.
+    static CURRENT_TASK: Arc<PageServerTask>;
+}
+
+///
+/// There are many kinds of tasks in the system. Some are associated with a particular
+/// tenant or timeline, while others are global.
+///
+/// Note that we don't try to limit how many task of a certain kind can be running
+/// at the same time.
+///
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+pub enum TaskKind {
+    // libpq listener task. It just accepts connection and spawns a
+    // PageRequestHandler task for each connection.
+    LibpqEndpointListener,
+
+    // HTTP endpoint listener.
+    HttpEndpointListener,
+
+    // Task that handles a single connection. A PageRequestHandler task
+    // starts detached from any particular tenant or timeline, but it can be
+    // associated with one later, after receiving a command from the client.
+    PageRequestHandler,
+
+    // Manages the WAL receiver connection for one timeline. It subscribes to
+    // events from storage_broker, decides which safekeeper to connect to. It spawns a
+    // separate WalReceiverConnection task to handle each connection.
+    WalReceiverManager,
+
+    // Handles a connection to a safekeeper, to stream WAL to a timeline.
+    WalReceiverConnection,
+
+    // Garbage collection worker. One per tenant
+    GarbageCollector,
+
+    // Compaction. One per tenant.
+    Compaction,
+
+    // Initial logical size calculation
+    InitialLogicalSizeCalculation,
+
+    // Task that flushes frozen in-memory layers to disk
+    LayerFlushTask,
+
+    // Task that uploads a file to remote storage
+    RemoteUploadTask,
+
+    // Task that downloads a file from remote storage
+    RemoteDownloadTask,
+
+    // task that handles the initial downloading of all tenants
+    InitialLoad,
+
+    // task that handles attaching a tenant
+    Attach,
+
+    // task that handhes metrics collection
+    MetricsCollection,
+
+    // task that drives downloading layers
+    DownloadAllRemoteLayers,
+}
+
+#[derive(Default)]
+struct MutableTaskState {
+    /// Tenant and timeline that this task is associated with.
+    tenant_id: Option<TenantId>,
+    timeline_id: Option<TimelineId>,
+
+    /// Handle for waiting for the task to exit. It can be None, if the
+    /// the task has already exited.
+    join_handle: Option<JoinHandle<()>>,
+}
+
+struct PageServerTask {
+    #[allow(dead_code)] // unused currently
+    task_id: PageserverTaskId,
+
+    kind: TaskKind,
+
+    name: String,
+
+    // To request task shutdown, just cancel this token.
+    cancel: CancellationToken,
+
+    mutable: Mutex<MutableTaskState>,
+}
+
+/// Launch a new task
+/// Note: if shutdown_process_on_error is set to true failure
+///   of the task will lead to shutdown of entire process
+pub fn spawn<F>(
+    runtime: &tokio::runtime::Handle,
+    kind: TaskKind,
+    tenant_id: Option<TenantId>,
+    timeline_id: Option<TimelineId>,
+    name: &str,
+    shutdown_process_on_error: bool,
+    future: F,
+) -> PageserverTaskId
 where
-    F: Future<Output = R> + Send + 'static,
+    F: Future<Output = anyhow::Result<()>> + Send + 'static,
+{
+    let cancel = CancellationToken::new();
+    let task_id = NEXT_TASK_ID.fetch_add(1, Ordering::Relaxed);
+    let task = Arc::new(PageServerTask {
+        task_id: PageserverTaskId(task_id),
+        kind,
+        name: name.to_string(),
+        cancel: cancel.clone(),
+        mutable: Mutex::new(MutableTaskState {
+            tenant_id,
+            timeline_id,
+            join_handle: None,
+        }),
+    });
+
+    TASKS.lock().unwrap().insert(task_id, Arc::clone(&task));
+
+    let mut task_mut = task.mutable.lock().unwrap();
+
+    let task_name = name.to_string();
+    let task_cloned = Arc::clone(&task);
+    let join_handle = runtime.spawn(task_wrapper(
+        task_name,
+        task_id,
+        task_cloned,
+        cancel,
+        shutdown_process_on_error,
+        future,
+    ));
+    task_mut.join_handle = Some(join_handle);
+    drop(task_mut);
+
+    // The task is now running. Nothing more to do here
+    PageserverTaskId(task_id)
+}
+
+/// This wrapper function runs in a newly-spawned task. It initializes the
+/// task-local variables and calls the payload function.
+async fn task_wrapper<F>(
+    task_name: String,
+    task_id: u64,
+    task: Arc<PageServerTask>,
+    shutdown_token: CancellationToken,
+    shutdown_process_on_error: bool,
+    future: F,
+) where
+    F: Future<Output = anyhow::Result<()>> + Send + 'static,
 {
     debug!("Starting task '{}'", task_name);
 
-    // We use AssertUnwindSafe here so that the payload function
-    // doesn't need to be UnwindSafe. We don't do anything after the
-    // unwinding that would expose us to unwind-unsafe behavior.
-    let result = AssertUnwindSafe(future).catch_unwind().await;
+    let result = SHUTDOWN_TOKEN
+        .scope(
+            shutdown_token,
+            CURRENT_TASK.scope(task, {
+                // We use AssertUnwindSafe here so that the payload function
+                // doesn't need to be UnwindSafe. We don't do anything after the
+                // unwinding that would expose us to unwind-unsafe behavior.
+                AssertUnwindSafe(future).catch_unwind()
+            }),
+        )
+        .await;
+    task_finish(result, task_name, task_id, shutdown_process_on_error).await;
+}
 
-    // Handle panics
-    match result {
-        Ok(result) => {
-            debug!("Task '{}' exited normally", task_name);
-            result
-        }
-        Err(err) => {
-            if shutdown_process_on_panic {
-                error!("Shutting down: task '{}' panicked: {:?}", task_name, err);
-                shutdown_pageserver(1).await;
-                unreachable!();
-            } else {
-                error!("Task '{}'  panicked: {:?}", task_name, err);
-                resume_unwind(err);
+async fn task_finish(
+    result: std::result::Result<
+        anyhow::Result<()>,
+        std::boxed::Box<dyn std::any::Any + std::marker::Send>,
+    >,
+    task_name: String,
+    task_id: u64,
+    shutdown_process_on_error: bool,
+) {
+    // Remove our entry from the global hashmap.
+    let task = TASKS
+        .lock()
+        .unwrap()
+        .remove(&task_id)
+        .expect("no task in registry");
+
+    let mut shutdown_process = false;
+    {
+        let task_mut = task.mutable.lock().unwrap();
+
+        match result {
+            Ok(Ok(())) => {
+                debug!("Task '{}' exited normally", task_name);
             }
+            Ok(Err(err)) => {
+                if shutdown_process_on_error {
+                    error!(
+                        "Shutting down: task '{}' tenant_id: {:?}, timeline_id: {:?} exited with error: {:?}",
+                        task_name, task_mut.tenant_id, task_mut.timeline_id, err
+                    );
+                    shutdown_process = true;
+                } else {
+                    error!(
+                        "Task '{}' tenant_id: {:?}, timeline_id: {:?} exited with error: {:?}",
+                        task_name, task_mut.tenant_id, task_mut.timeline_id, err
+                    );
+                }
+            }
+            Err(err) => {
+                if shutdown_process_on_error {
+                    error!(
+                        "Shutting down: task '{}' tenant_id: {:?}, timeline_id: {:?} panicked: {:?}",
+                        task_name, task_mut.tenant_id, task_mut.timeline_id, err
+                    );
+                    shutdown_process = true;
+                } else {
+                    error!(
+                        "Task '{}' tenant_id: {:?}, timeline_id: {:?} panicked: {:?}",
+                        task_name, task_mut.tenant_id, task_mut.timeline_id, err
+                    );
+                }
+            }
+        }
+    }
+
+    if shutdown_process {
+        shutdown_pageserver(1).await;
+    }
+}
+
+// expected to be called from the task of the given id.
+pub fn associate_with(tenant_id: Option<TenantId>, timeline_id: Option<TimelineId>) {
+    CURRENT_TASK.with(|ct| {
+        let mut task_mut = ct.mutable.lock().unwrap();
+        task_mut.tenant_id = tenant_id;
+        task_mut.timeline_id = timeline_id;
+    });
+}
+
+/// Is there a task running that matches the criteria
+
+/// Signal and wait for tasks to shut down.
+///
+///
+/// The arguments are used to select the tasks to kill. Any None arguments are
+/// ignored. For example, to shut down all WalReceiver tasks:
+///
+///   shutdown_tasks(Some(TaskKind::WalReceiver), None, None)
+///
+/// Or to shut down all tasks for given timeline:
+///
+///   shutdown_tasks(None, Some(tenant_id), Some(timeline_id))
+///
+pub async fn shutdown_tasks(
+    kind: Option<TaskKind>,
+    tenant_id: Option<TenantId>,
+    timeline_id: Option<TimelineId>,
+) {
+    let mut victim_tasks = Vec::new();
+
+    {
+        let tasks = TASKS.lock().unwrap();
+        for task in tasks.values() {
+            let task_mut = task.mutable.lock().unwrap();
+            if (kind.is_none() || Some(task.kind) == kind)
+                && (tenant_id.is_none() || task_mut.tenant_id == tenant_id)
+                && (timeline_id.is_none() || task_mut.timeline_id == timeline_id)
+            {
+                task.cancel.cancel();
+                victim_tasks.push(Arc::clone(task));
+            }
+        }
+    }
+
+    for task in victim_tasks {
+        let join_handle = {
+            let mut task_mut = task.mutable.lock().unwrap();
+            info!("waiting for {} to shut down", task.name);
+            let join_handle = task_mut.join_handle.take();
+            drop(task_mut);
+            join_handle
+        };
+        if let Some(join_handle) = join_handle {
+            let _ = join_handle.await;
+        } else {
+            // Possibly one of:
+            //  * The task had not even fully started yet.
+            //  * It was shut down concurrently and already exited
         }
     }
 }
 
-///
-/// Perform pageserver shutdown. This is called on receiving a signal,
-/// or if one of the tasks marked as 'shutdown_process_on_error' dies.
-///
-/// This never returns.
-pub async fn shutdown_pageserver(exit_code: i32) {
-    // Shut down the libpq endpoint task. This prevents new connections from
-    // being accepted.
-    context::shutdown_tasks(TaskKind::LibpqEndpointListener).await;
-
-    // Shut down all tenants gracefully
-    crate::tenant::mgr::shutdown_all_tenants().await;
-
-    // Shut down the HTTP endpoint last, so that you can still check the server's
-    // status while it's shutting down.
-    // FIXME: We should probably stop accepting commands like attach/detach earlier.
-    context::shutdown_tasks(TaskKind::HttpEndpointListener).await;
-
-    // There should be nothing left, but let's be sure
-    context::shutdown_all_tasks().await;
-
-    info!("Shut down successfully completed");
-    std::process::exit(exit_code);
+pub fn current_task_kind() -> Option<TaskKind> {
+    CURRENT_TASK.try_with(|ct| ct.kind).ok()
+}
+
+pub fn current_task_id() -> Option<PageserverTaskId> {
+    CURRENT_TASK.try_with(|ct| ct.task_id).ok()
+}
+
+/// A Future that can be used to check if the current task has been requested to
+/// shut down.
+pub async fn shutdown_watcher() {
+    let token = SHUTDOWN_TOKEN
+        .try_with(|t| t.clone())
+        .expect("shutdown_requested() called in an unexpected task or thread");
+
+    token.cancelled().await;
+}
+
+/// Clone the current task's cancellation token, which can be moved across tasks.
+///
+/// When the task which is currently executing is shutdown, the cancellation token will be
+/// cancelled. It can however be moved to other tasks, such as `tokio::task::spawn_blocking` or
+/// `tokio::task::JoinSet::spawn`.
+pub fn shutdown_token() -> CancellationToken {
+    SHUTDOWN_TOKEN
+        .try_with(|t| t.clone())
+        .expect("shutdown_token() called in an unexpected task or thread")
+}
+
+/// Has the current task been requested to shut down?
+pub fn is_shutdown_requested() -> bool {
+    if let Ok(cancel) = SHUTDOWN_TOKEN.try_with(|t| t.clone()) {
+        cancel.is_cancelled()
+    } else {
+        if !cfg!(test) {
+            warn!("is_shutdown_requested() called in an unexpected task or thread");
+        }
+        false
+    }
 }

pageserver/src/tenant/layer_map.rs

@@ -10,23 +10,21 @@
 //! 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::{range_eq, range_overlaps};
-use amplify_num::i256;
+use crate::tenant::storage_layer::InMemoryLayer;
+use crate::tenant::storage_layer::Layer;
 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 super::storage_layer::{InMemoryLayer, Layer};
+use historic_layer_coverage::BufferedHistoricLayerCoverage;
 
 ///
 /// LayerMap tracks what layers exist on a timeline.
@@ -51,194 +49,26 @@ pub struct LayerMap<L: ?Sized> {
     ///
     pub frozen_layers: VecDeque<Arc<InMemoryLayer>>,
 
-    /// All the historic layers are kept here
-    historic_layers: RTree<LayerRTreeObject<L>>,
+    /// Index of the historic layers optimized for search
+    historic: BufferedHistoricLayerCoverage<Arc<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.
     l0_delta_layers: Vec<Arc<L>>,
 }
 
-impl<L: ?Sized> Default for LayerMap<L> {
+impl<L: ?Sized + PartialEq> Default for LayerMap<L> {
     fn default() -> Self {
         Self {
             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>,
@@ -247,141 +77,120 @@ pub struct SearchResult<L: ?Sized> {
 
 impl<L> LayerMap<L>
 where
-    L: ?Sized + Layer,
+    L: ?Sized + Layer + PartialEq,
 {
     ///
-    /// Find the latest layer that covers the given 'key', with lsn <
-    /// 'end_lsn'.
+    /// Find the latest layer (by lsn.end) that covers the given
+    /// 'key', with lsn.start < 'end_lsn'.
     ///
-    /// 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.
+    /// 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.
     ///
     /// 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>> {
-        // 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);
-            }
-        }
+        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());
 
-        // 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;
+        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,
+                })
             }
-            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), None) => {
+                let lsn_floor = delta.get_lsn_range().start;
+                Some(SearchResult {
+                    layer: delta,
+                    lsn_floor,
+                })
             }
-            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));
-                    }
+            (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,
+                    })
                 } else {
-                    latest_delta.replace(Arc::clone(l));
+                    let lsn_floor =
+                        std::cmp::max(delta.get_lsn_range().start, image.get_lsn_range().start + 1);
+                    Some(SearchResult {
+                        layer: delta,
+                        lsn_floor,
+                    })
                 }
             }
         }
-        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.clone());
+            self.l0_delta_layers.push(layer);
         }
-        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();
 
@@ -394,96 +203,51 @@ 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?
+    /// 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?
     ///
     /// This is used for garbage collection, to determine if an old layer can
     /// be deleted.
-    pub fn image_layer_exists(
-        &self,
-        key_range: &Range<Key>,
-        lsn_range: &Range<Lsn>,
-    ) -> Result<bool> {
-        let mut range_remain = key_range.clone();
+    pub fn image_layer_exists(&self, key: &Range<Key>, lsn: &Range<Lsn>) -> Result<bool> {
+        if key.is_empty() {
+            return Ok(true);
+        }
 
-        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 version = match self.historic.get().unwrap().get_version(lsn.end.0) {
+            Some(v) => v,
+            None => return Ok(false),
+        };
 
-                    if img_key_end >= range_remain.end {
-                        return Ok(true);
-                    }
-                    range_remain.start = img_key_end;
-                }
-            }
+        let start = key.start.to_i128();
+        let end = key.end.to_i128();
 
-            if !made_progress {
+        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) {
                 return Ok(false);
             }
         }
+
+        Ok(true)
     }
 
     pub fn iter_historic_layers(&self) -> impl '_ + Iterator<Item = Arc<L>> {
-        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
+        self.historic.iter()
     }
 
     ///
@@ -499,87 +263,239 @@ where
         key_range: &Range<Key>,
         lsn: Lsn,
     ) -> Result<Vec<(Range<Key>, Option<Arc<L>>)>> {
-        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);
-            }
+        // 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();
         }
-        points.push(key_range.end);
 
-        points.sort();
-        points.dedup();
+        // Add the final interval
+        let kr = Key::from_i128(current_key)..Key::from_i128(end);
+        coverage.push((kr, current_val.take()));
 
-        // 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)
+        Ok(coverage)
     }
 
-    /// 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 {
+    /// 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) {
             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));
 
-            // 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;
+        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);
+                    }
+                }
             }
 
-            result += 1;
+            current_key = change_key;
+            current_val = change_val.clone();
         }
-        Ok(result)
+
+        // 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()
     }
 
     /// Return all L0 delta layers
@@ -603,8 +519,8 @@ where
         }
 
         println!("historic_layers:");
-        for e in self.historic_layers.iter() {
-            e.layer.dump(verbose)?;
+        for layer in self.iter_historic_layers() {
+            layer.dump(verbose)?;
         }
         println!("End dump LayerMap");
         Ok(())

pageserver/src/tenant/layer_map/historic_layer_coverage.rs

@@ -0,0 +1,418 @@
+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 + PartialEq> Default for HistoricLayerCoverage<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<Value: Clone + PartialEq> 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 + PartialEq> Default for BufferedHistoricLayerCoverage<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<Value: Clone + PartialEq> 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

@@ -0,0 +1,200 @@
+use std::ops::Range;
+
+// TODO the `im` crate has 20x more downloads and also has
+// persistent/immutable BTree. It also runs a bit faster but
+// results are not the same on some tests.
+use rpds::RedBlackTreeMapSync;
+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.
+    ///
+    /// We use the Sync version of the map because we want Self to
+    /// be Sync. Using nonsync might be faster, if we can work with
+    /// that.
+    nodes: RedBlackTreeMapSync<i128, Option<(u64, Value)>>,
+
+    im: OrdMap<i128, Option<(u64, Value)>>,
+}
+
+impl<T: Clone + PartialEq> Default for LayerCoverage<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<Value: Clone + PartialEq> LayerCoverage<Value> {
+    pub fn new() -> Self {
+        Self {
+            nodes: RedBlackTreeMapSync::default(),
+            im: 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_mut(key, value);
+
+        let im_value = match self.im.range(..=key).last() {
+            Some((_, Some(v))) => Some(v.clone()),
+            Some((_, None)) => None,
+            None => None,
+        };
+        self.im.insert(key, im_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_mut(k, Some((lsn.end, value.clone())));
+        }
+        for k in to_remove {
+            self.nodes.remove_mut(&k);
+        }
+
+        let mut to_update = Vec::new();
+        let mut to_remove = Vec::new();
+        let mut prev_covered = false;
+        for (k, node) in self.im.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.im.insert(k, Some((lsn.end, value.clone())));
+        }
+        for k in to_remove {
+            self.im.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> {
+        let res = self.nodes
+            .range(..=key)
+            .rev()
+            .next()?
+            .1
+            .as_ref()
+            .map(|(_, v)| v.clone());
+        let im_res = self.im
+            .range(..=key)
+            .rev()
+            .next()?
+            .1
+            .as_ref()
+            .map(|(_, v)| v.clone());
+        if res != im_res {
+            panic!("aha");
+        }
+        im_res
+    }
+
+    /// 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(),
+            im: self.im.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 + PartialEq> Default for LayerCoverageTuple<T> {
+    fn default() -> Self {
+        Self {
+            image_coverage: LayerCoverage::default(),
+            delta_coverage: LayerCoverage::default(),
+        }
+    }
+}
+
+impl<Value: Clone + PartialEq> LayerCoverageTuple<Value> {
+    pub fn clone(&self) -> Self {
+        Self {
+            image_coverage: self.image_coverage.clone(),
+            delta_coverage: self.delta_coverage.clone(),
+        }
+    }
+}

pageserver/src/tenant/mgr.rs

@@ -8,8 +8,6 @@ use std::sync::Arc;
 use tokio::fs;
 
 use anyhow::Context;
-use futures::stream::FuturesUnordered;
-use futures::StreamExt;
 use once_cell::sync::Lazy;
 use tokio::sync::RwLock;
 use tracing::*;
@@ -18,9 +16,9 @@ use remote_storage::GenericRemoteStorage;
 use utils::crashsafe;
 
 use crate::config::PageServerConf;
-use crate::context::RequestContext;
+use crate::task_mgr::{self, TaskKind};
 use crate::tenant::config::TenantConfOpt;
-use crate::tenant::{Tenant, TenantRequestContext, TenantState};
+use crate::tenant::{Tenant, TenantState};
 use crate::IGNORED_TENANT_FILE_NAME;
 
 use utils::fs_ext::PathExt;
@@ -183,11 +181,25 @@ pub async fn shutdown_all_tenants() {
         tenants_to_shut_down
     };
 
-    let mut shutdown_futures: FuturesUnordered<_> = FuturesUnordered::new();
-    for tenant in tenants_to_shut_down.iter() {
-        shutdown_futures.push(tenant.graceful_shutdown(true));
+    // Shut down all existing walreceiver connections and stop accepting the new ones.
+    task_mgr::shutdown_tasks(Some(TaskKind::WalReceiverManager), None, None).await;
+
+    // Ok, no background tasks running anymore. Flush any remaining data in
+    // memory to disk.
+    //
+    // We assume that any incoming connections that might request pages from
+    // the tenant have already been terminated by the caller, so there
+    // should be no more activity in any of the repositories.
+    //
+    // On error, log it but continue with the shutdown for other tenants.
+    for tenant in tenants_to_shut_down {
+        let tenant_id = tenant.tenant_id();
+        debug!("shutdown tenant {tenant_id}");
+
+        if let Err(err) = tenant.freeze_and_flush().await {
+            error!("Could not checkpoint tenant {tenant_id} during shutdown: {err:?}");
+        }
     }
-    while let Some(_result) = shutdown_futures.next().await {}
 }
 
 pub async fn create_tenant(
@@ -222,47 +234,36 @@ pub async fn update_tenant_config(
     conf: &'static PageServerConf,
     tenant_conf: TenantConfOpt,
     tenant_id: TenantId,
-    ctx: &RequestContext,
 ) -> anyhow::Result<()> {
     info!("configuring tenant {tenant_id}");
-    let (tenant, _ctx) = get_active_tenant(tenant_id, ctx).await?;
-
-    tenant.update_tenant_config(tenant_conf);
+    get_tenant(tenant_id, true)
+        .await?
+        .update_tenant_config(tenant_conf);
     Tenant::persist_tenant_config(&conf.tenant_config_path(tenant_id), tenant_conf, false)?;
     Ok(())
 }
 
 /// Gets the tenant from the in-memory data, erroring if it's absent or is not fitting to the query.
 /// `active_only = true` allows to query only tenants that are ready for operations, erroring on other kinds of tenants.
-pub async fn get_active_tenant(
-    tenant_id: TenantId,
-    parent_ctx: &RequestContext,
-) -> anyhow::Result<(Arc<Tenant>, TenantRequestContext)> {
-    let tenant = get_tenant(tenant_id).await?;
-    let tenant_ctx = match tenant.get_context(parent_ctx) {
-        Ok(ctx) => ctx,
-        Err(state) => anyhow::bail!("Tenant {} is not active, state: {:?}", tenant_id, state,),
-    };
-    Ok((tenant, tenant_ctx))
-}
-
-pub async fn get_tenant(tenant_id: TenantId) -> anyhow::Result<Arc<Tenant>> {
+pub async fn get_tenant(tenant_id: TenantId, active_only: bool) -> anyhow::Result<Arc<Tenant>> {
     let m = TENANTS.read().await;
     let tenant = m
         .get(&tenant_id)
         .with_context(|| format!("Tenant {tenant_id} not found in the local state"))?;
-
-    Ok(Arc::clone(tenant))
+    if active_only && !tenant.is_active() {
+        anyhow::bail!(
+            "Tenant {tenant_id} is not active. Current state: {:?}",
+            tenant.current_state()
+        )
+    } else {
+        Ok(Arc::clone(tenant))
+    }
 }
 
-pub async fn delete_timeline(
-    tenant_id: TenantId,
-    timeline_id: TimelineId,
-    ctx: &RequestContext,
-) -> anyhow::Result<()> {
-    match get_active_tenant(tenant_id, ctx).await {
-        Ok((tenant, ctx)) => {
-            tenant.delete_timeline(timeline_id, &ctx).await?;
+pub async fn delete_timeline(tenant_id: TenantId, timeline_id: TimelineId) -> anyhow::Result<()> {
+    match get_tenant(tenant_id, true).await {
+        Ok(tenant) => {
+            tenant.delete_timeline(timeline_id).await?;
         }
         Err(e) => anyhow::bail!("Cannot access tenant {tenant_id} in local tenant state: {e:?}"),
     }
@@ -394,31 +395,27 @@ where
     // The exclusive lock here ensures we don't miss the tenant state updates before trying another removal.
     // tenant-wde cleanup operations may take some time (removing the entire tenant directory), we want to
     // avoid holding the lock for the entire process.
-    let tenant = {
+    {
         let tenants_accessor = TENANTS.write().await;
         match tenants_accessor.get(&tenant_id) {
             Some(tenant) => match tenant.current_state() {
                 TenantState::Attaching
                 | TenantState::Loading
                 | TenantState::Broken
-                | TenantState::Active => {
-                    tenant.set_stopping();
-                    Arc::clone(tenant)
-                }
+                | TenantState::Active => tenant.set_stopping(),
                 TenantState::Stopping => {
                     anyhow::bail!("Tenant {tenant_id} is stopping already")
                 }
             },
             None => anyhow::bail!("Tenant not found for id {tenant_id}"),
         }
-    };
+    }
 
-    // Shut down all tenant and timeline tasks.
-    tenant.graceful_shutdown(true).await;
+    // shutdown all tenant and timeline tasks: gc, compaction, page service)
+    // No new tasks will be started for this tenant because it's in `Stopping` state.
+    // Hence, once we're done here, the `tenant_cleanup` callback can mutate tenant on-disk state freely.
+    task_mgr::shutdown_tasks(None, Some(tenant_id), None).await;
 
-    // All tasks that operated on the tenant or any of its timelines have no finished,
-    // and they are in Stopped state so that new ones cannot appear anymore. Proceed
-    // with the cleanup.
     match tenant_cleanup
         .await
         .with_context(|| format!("Failed to run cleanup for tenant {tenant_id}"))
@@ -440,3 +437,108 @@ where
         }
     }
 }
+
+#[cfg(feature = "testing")]
+use {
+    crate::repository::GcResult, pageserver_api::models::TimelineGcRequest,
+    utils::http::error::ApiError,
+};
+
+#[cfg(feature = "testing")]
+pub async fn immediate_gc(
+    tenant_id: TenantId,
+    timeline_id: TimelineId,
+    gc_req: TimelineGcRequest,
+) -> Result<tokio::sync::oneshot::Receiver<Result<GcResult, anyhow::Error>>, ApiError> {
+    let guard = TENANTS.read().await;
+
+    let tenant = guard
+        .get(&tenant_id)
+        .map(Arc::clone)
+        .with_context(|| format!("Tenant {tenant_id} not found"))
+        .map_err(ApiError::NotFound)?;
+
+    let gc_horizon = gc_req.gc_horizon.unwrap_or_else(|| tenant.get_gc_horizon());
+    // Use tenant's pitr setting
+    let pitr = tenant.get_pitr_interval();
+
+    // Run in task_mgr to avoid race with detach operation
+    let (task_done, wait_task_done) = tokio::sync::oneshot::channel();
+    task_mgr::spawn(
+        &tokio::runtime::Handle::current(),
+        TaskKind::GarbageCollector,
+        Some(tenant_id),
+        Some(timeline_id),
+        &format!("timeline_gc_handler garbage collection run for tenant {tenant_id} timeline {timeline_id}"),
+        false,
+        async move {
+            fail::fail_point!("immediate_gc_task_pre");
+            let result = tenant
+                .gc_iteration(Some(timeline_id), gc_horizon, pitr)
+                .instrument(info_span!("manual_gc", tenant = %tenant_id, timeline = %timeline_id))
+                .await;
+                // FIXME: `gc_iteration` can return an error for multiple reasons; we should handle it
+                // better once the types support it.
+            match task_done.send(result) {
+                Ok(_) => (),
+                Err(result) => error!("failed to send gc result: {result:?}"),
+            }
+            Ok(())
+        }
+    );
+
+    // drop the guard until after we've spawned the task so that timeline shutdown will wait for the task
+    drop(guard);
+
+    Ok(wait_task_done)
+}
+
+#[cfg(feature = "testing")]
+pub async fn immediate_compact(
+    tenant_id: TenantId,
+    timeline_id: TimelineId,
+) -> Result<tokio::sync::oneshot::Receiver<anyhow::Result<()>>, ApiError> {
+    let guard = TENANTS.read().await;
+
+    let tenant = guard
+        .get(&tenant_id)
+        .map(Arc::clone)
+        .with_context(|| format!("Tenant {tenant_id} not found"))
+        .map_err(ApiError::NotFound)?;
+
+    let timeline = tenant
+        .get_timeline(timeline_id, true)
+        .map_err(ApiError::NotFound)?;
+
+    // Run in task_mgr to avoid race with detach operation
+    let (task_done, wait_task_done) = tokio::sync::oneshot::channel();
+    task_mgr::spawn(
+        &tokio::runtime::Handle::current(),
+        TaskKind::Compaction,
+        Some(tenant_id),
+        Some(timeline_id),
+        &format!(
+            "timeline_compact_handler compaction run for tenant {tenant_id} timeline {timeline_id}"
+        ),
+        false,
+        async move {
+            let result = timeline
+                .compact()
+                .instrument(
+                    info_span!("manual_compact", tenant = %tenant_id, timeline = %timeline_id),
+                )
+                .await;
+
+            match task_done.send(result) {
+                Ok(_) => (),
+                Err(result) => error!("failed to send compaction result: {result:?}"),
+            }
+            Ok(())
+        },
+    );
+
+    // drop the guard until after we've spawned the task so that timeline shutdown will wait for the task
+    drop(guard);
+
+    Ok(wait_task_done)
+}

pageserver/src/tenant/remote_timeline_client.rs

@@ -16,7 +16,7 @@
 //! unless the pageserver is configured without remote storage.
 //!
 //! We allocate the client instance in [Timeline][`crate::tenant::Timeline`], i.e.,
-//! either in [`crate::tenant::mgr`] during startup or when creating a new
+//! either in [`crate::tenant_mgr`] during startup or when creating a new
 //! timeline.
 //! However, the client does not become ready for use until we've initialized its upload queue:
 //!
@@ -214,7 +214,6 @@ use anyhow::ensure;
 use remote_storage::{DownloadError, GenericRemoteStorage};
 use std::ops::DerefMut;
 use tokio::runtime::Runtime;
-use tokio_util::sync::CancellationToken;
 use tracing::{debug, info, warn};
 use tracing::{info_span, Instrument};
 use utils::lsn::Lsn;
@@ -226,12 +225,12 @@ use crate::tenant::remote_timeline_client::index::LayerFileMetadata;
 use crate::{
     config::PageServerConf,
     task_mgr,
+    task_mgr::TaskKind,
     task_mgr::BACKGROUND_RUNTIME,
     tenant::metadata::TimelineMetadata,
     tenant::upload_queue::{
         UploadOp, UploadQueue, UploadQueueInitialized, UploadQueueStopped, UploadTask,
     },
-    tenant::TimelineRequestContext,
     {exponential_backoff, DEFAULT_BASE_BACKOFF_SECONDS, DEFAULT_MAX_BACKOFF_SECONDS},
 };
 
@@ -314,50 +313,25 @@ impl RemoteTimelineClient {
     /// Initialize the upload queue for a remote storage that already received
     /// an index file upload, i.e., it's not empty.
     /// The given `index_part` must be the one on the remote.
-    pub fn init_upload_queue(
-        self: &Arc<Self>,
-        index_part: &IndexPart,
-        upload_ctx: TimelineRequestContext,
-    ) -> anyhow::Result<()> {
-        let cancellation_token = upload_ctx.cancellation_token().clone();
+    pub fn init_upload_queue(&self, index_part: &IndexPart) -> anyhow::Result<()> {
         let mut upload_queue = self.upload_queue.lock().unwrap();
-        upload_queue.initialize_with_current_remote_index_part(index_part, upload_ctx)?;
+        upload_queue.initialize_with_current_remote_index_part(index_part)?;
         self.update_remote_physical_size_gauge(Some(index_part));
-        self.spawn_cancellation_watch(cancellation_token);
         Ok(())
     }
 
     /// Initialize the upload queue for the case where the remote storage is empty,
     /// i.e., it doesn't have an `IndexPart`.
     pub fn init_upload_queue_for_empty_remote(
-        self: &Arc<Self>,
+        &self,
         local_metadata: &TimelineMetadata,
-        upload_ctx: TimelineRequestContext,
     ) -> anyhow::Result<()> {
-        let cancellation_token = upload_ctx.cancellation_token().clone();
         let mut upload_queue = self.upload_queue.lock().unwrap();
-        upload_queue.initialize_empty_remote(local_metadata, upload_ctx)?;
+        upload_queue.initialize_empty_remote(local_metadata)?;
         self.update_remote_physical_size_gauge(None);
-        self.spawn_cancellation_watch(cancellation_token);
         Ok(())
     }
 
-    /// Spawn a task that calls `stop` on cancellation. It's important that we
-    /// stop the upload queue promptly, because it holds onto the RequestContext,
-    /// which in turn prevents the Timeline from shutting down.
-    fn spawn_cancellation_watch(self: &Arc<Self>, cancellation_token: CancellationToken) {
-        let self_rc = Arc::clone(self);
-        task_mgr::spawn(
-            self.runtime.handle(),
-            "remote upload queue cancellation watch",
-            false,
-            async move {
-                cancellation_token.cancelled().await;
-                self_rc.stop();
-            },
-        );
-    }
-
     pub fn last_uploaded_consistent_lsn(&self) -> Option<Lsn> {
         match &*self.upload_queue.lock().unwrap() {
             UploadQueue::Uninitialized => None,
@@ -651,10 +625,7 @@ impl RemoteTimelineClient {
     ///
     /// Wait for all previously scheduled uploads/deletions to complete
     ///
-    pub async fn wait_completion(
-        self: &Arc<Self>,
-        ctx: &TimelineRequestContext,
-    ) -> anyhow::Result<()> {
+    pub async fn wait_completion(self: &Arc<Self>) -> anyhow::Result<()> {
         let (sender, mut receiver) = tokio::sync::watch::channel(());
         let barrier_op = UploadOp::Barrier(sender);
 
@@ -668,16 +639,9 @@ impl RemoteTimelineClient {
             self.launch_queued_tasks(upload_queue);
         }
 
-        tokio::select! {
-            result = receiver.changed() => {
-                if result.is_err() {
-                    anyhow::bail!("wait_completion aborted because upload queue was stopped");
-                }
-            },
-            _ = ctx.cancelled() => {
-                anyhow::bail!("request cancelled while waiting on uploads to finish");
-            },
-        };
+        if receiver.changed().await.is_err() {
+            anyhow::bail!("wait_completion aborted because upload queue was stopped");
+        }
         Ok(())
     }
 
@@ -755,15 +719,16 @@ impl RemoteTimelineClient {
 
             // Spawn task to perform the task
             let self_rc = Arc::clone(self);
-
-            let cancellation_token = upload_queue.upload_ctx.cancellation_token().clone();
-
             task_mgr::spawn(
                 self.runtime.handle(),
+                TaskKind::RemoteUploadTask,
+                Some(self.tenant_id),
+                Some(self.timeline_id),
                 "remote upload",
                 false,
                 async move {
-                    self_rc.perform_upload_task(task, cancellation_token).await;
+                    self_rc.perform_upload_task(task).await;
+                    Ok(())
                 }
                 .instrument(info_span!(parent: None, "remote_upload", tenant = %self.tenant_id, timeline = %self.timeline_id, upload_task_id = %task_id)),
             );
@@ -783,11 +748,7 @@ impl RemoteTimelineClient {
     /// The task can be shut down, however. That leads to stopping the whole
     /// queue.
     ///
-    async fn perform_upload_task(
-        self: &Arc<Self>,
-        task: Arc<UploadTask>,
-        cancellation_token: CancellationToken,
-    ) {
+    async fn perform_upload_task(self: &Arc<Self>, task: Arc<UploadTask>) {
         // Loop to retry until it completes.
         loop {
             // If we're requested to shut down, close up shop and exit.
@@ -795,11 +756,11 @@ impl RemoteTimelineClient {
             // Note: We only check for the shutdown requests between retries, so
             // if a shutdown request arrives while we're busy uploading, in the
             // upload::upload:*() call below, we will wait not exit until it has
-            // finished. We probably could cancel the upload by simply dropping
+            // finisheed. We probably could cancel the upload by simply dropping
             // the Future, but we're not 100% sure if the remote storage library
             // is cancellation safe, so we don't dare to do that. Hopefully, the
             // upload finishes or times out soon enough.
-            if cancellation_token.is_cancelled() {
+            if task_mgr::is_shutdown_requested() {
                 info!("upload task cancelled by shutdown request");
                 self.calls_unfinished_metric_end(&task.op);
                 self.stop();
@@ -897,7 +858,7 @@ impl RemoteTimelineClient {
 
                     // sleep until it's time to retry, or we're cancelled
                     tokio::select! {
-                        _ = cancellation_token.cancelled() => { },
+                        _ = task_mgr::shutdown_watcher() => { },
                         _ = exponential_backoff(
                             retries,
                             DEFAULT_BASE_BACKOFF_SECONDS,
@@ -1049,9 +1010,7 @@ impl RemoteTimelineClient {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::context::{DownloadBehavior, RequestContext, TaskKind};
     use crate::tenant::harness::{TenantHarness, TIMELINE_ID};
-    use crate::DEFAULT_PG_VERSION;
     use remote_storage::{RemoteStorageConfig, RemoteStorageKind};
     use std::{collections::HashSet, path::Path};
     use utils::lsn::Lsn;
@@ -1070,7 +1029,7 @@ mod tests {
             Lsn(0),
             // Any version will do
             // but it should be consistent with the one in the tests
-            DEFAULT_PG_VERSION,
+            crate::DEFAULT_PG_VERSION,
         );
 
         // go through serialize + deserialize to fix the header, including checksum
@@ -1105,19 +1064,9 @@ mod tests {
     // Test scheduling
     #[test]
     fn upload_scheduling() -> anyhow::Result<()> {
-        // Use a current-thread runtime in the test
-        let runtime = Box::leak(Box::new(
-            tokio::runtime::Builder::new_current_thread()
-                .enable_all()
-                .build()?,
-        ));
-        let _entered = runtime.enter();
-
         let harness = TenantHarness::create("upload_scheduling")?;
-        let (tenant, tenant_ctx) = runtime.block_on(harness.load());
-        let (_timeline, timeline_ctx) =
-            tenant.create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION, &tenant_ctx)?;
         let timeline_path = harness.timeline_path(&TIMELINE_ID);
+        std::fs::create_dir_all(&timeline_path)?;
 
         let remote_fs_dir = harness.conf.workdir.join("remote_fs");
         std::fs::create_dir_all(remote_fs_dir)?;
@@ -1135,6 +1084,14 @@ mod tests {
             storage: RemoteStorageKind::LocalFs(remote_fs_dir.clone()),
         };
 
+        // Use a current-thread runtime in the test
+        let runtime = Box::leak(Box::new(
+            tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()?,
+        ));
+        let _entered = runtime.enter();
+
         // Test outline:
         //
         // Schedule upload of a bunch of layers. Check that they are started immediately, not queued
@@ -1170,11 +1127,7 @@ mod tests {
         println!("remote_timeline_dir: {}", remote_timeline_dir.display());
 
         let metadata = dummy_metadata(Lsn(0x10));
-        let upload_ctx = timeline_ctx.register_another(RequestContext::new(
-            TaskKind::RemoteUploadTask,
-            DownloadBehavior::Error,
-        ));
-        client.init_upload_queue_for_empty_remote(&metadata, upload_ctx)?;
+        client.init_upload_queue_for_empty_remote(&metadata)?;
 
         // Create a couple of dummy files,  schedule upload for them
         let content_foo = dummy_contents("foo");
@@ -1214,7 +1167,7 @@ mod tests {
         }
 
         // Wait for the uploads to finish
-        runtime.block_on(client.wait_completion(&timeline_ctx))?;
+        runtime.block_on(client.wait_completion())?;
         {
             let mut guard = client.upload_queue.lock().unwrap();
             let upload_queue = guard.initialized_mut().unwrap();
@@ -1251,7 +1204,7 @@ mod tests {
         assert_remote_files(&["foo", "bar", "index_part.json"], &remote_timeline_dir);
 
         // Finish them
-        runtime.block_on(client.wait_completion(&timeline_ctx))?;
+        runtime.block_on(client.wait_completion())?;
 
         assert_remote_files(&["bar", "baz", "index_part.json"], &remote_timeline_dir);
 

pageserver/src/tenant/size.rs

@@ -3,9 +3,10 @@ use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use anyhow::Context;
+use tokio::sync::oneshot::error::RecvError;
 use tokio::sync::Semaphore;
 
-use crate::tenant::{PageReconstructError, TenantRequestContext, TimelineRequestContext};
+use crate::pgdatadir_mapping::CalculateLogicalSizeError;
 
 use super::Tenant;
 use utils::id::TimelineId;
@@ -62,14 +63,13 @@ pub(super) async fn gather_inputs(
     tenant: &Tenant,
     limit: &Arc<Semaphore>,
     logical_size_cache: &mut HashMap<(TimelineId, Lsn), u64>,
-    tenant_ctx: &TenantRequestContext,
 ) -> anyhow::Result<ModelInputs> {
     // with joinset, on drop, all of the tasks will just be de-scheduled, which we can use to
     // our advantage with `?` error handling.
     let mut joinset = tokio::task::JoinSet::new();
 
     let timelines = tenant
-        .refresh_gc_info(tenant_ctx)
+        .refresh_gc_info()
         .await
         .context("Failed to refresh gc_info before gathering inputs")?;
 
@@ -97,21 +97,9 @@ pub(super) async fn gather_inputs(
     // used to determine the `retention_period` for the size model
     let mut max_cutoff_distance = None;
 
-    let mut ctx_dropguards: Vec<tokio_util::sync::DropGuard> = Vec::new();
-
     for timeline in timelines {
         let last_record_lsn = timeline.get_last_record_lsn();
 
-        let ctx = match timeline.get_context(tenant_ctx) {
-            Ok(ctx) => ctx,
-            Err(state) => {
-                info!("skipping tenant size calculation for timeline because it is in {state:?} state");
-                continue;
-            }
-        };
-        ctx_dropguards.push(ctx.cancellation_token().clone().drop_guard());
-        let ctx = Arc::new(ctx);
-
         let (interesting_lsns, horizon_cutoff, pitr_cutoff, next_gc_cutoff) = {
             // there's a race between the update (holding tenant.gc_lock) and this read but it
             // might not be an issue, because it's not for Timeline::gc
@@ -181,23 +169,19 @@ pub(super) async fn gather_inputs(
             timeline_id: timeline.timeline_id,
         });
 
-        for (lsn, _kind) in interesting_lsns.iter() {
-            let lsn = *lsn;
-            if let Some(size) = logical_size_cache.get(&(timeline.timeline_id, lsn)) {
+        for (lsn, _kind) in &interesting_lsns {
+            if let Some(size) = logical_size_cache.get(&(timeline.timeline_id, *lsn)) {
                 updates.push(Update {
-                    lsn,
+                    lsn: *lsn,
                     timeline_id: timeline.timeline_id,
                     command: Command::Update(*size),
                 });
 
-                needed_cache.insert((timeline.timeline_id, lsn));
+                needed_cache.insert((timeline.timeline_id, *lsn));
             } else {
                 let timeline = Arc::clone(&timeline);
                 let parallel_size_calcs = Arc::clone(limit);
-                let ctx_clone = Arc::clone(&ctx);
-                joinset.spawn(async move {
-                    calculate_logical_size(parallel_size_calcs, timeline, lsn, &ctx_clone).await
-                });
+                joinset.spawn(calculate_logical_size(parallel_size_calcs, timeline, *lsn));
             }
         }
 
@@ -373,7 +357,7 @@ enum LsnKind {
 struct TimelineAtLsnSizeResult(
     Arc<crate::tenant::Timeline>,
     utils::lsn::Lsn,
-    Result<u64, PageReconstructError>,
+    Result<u64, CalculateLogicalSizeError>,
 );
 
 #[instrument(skip_all, fields(timeline_id=%timeline.timeline_id, lsn=%lsn))]
@@ -381,13 +365,14 @@ async fn calculate_logical_size(
     limit: Arc<tokio::sync::Semaphore>,
     timeline: Arc<crate::tenant::Timeline>,
     lsn: utils::lsn::Lsn,
-    ctx: &TimelineRequestContext,
-) -> Result<TimelineAtLsnSizeResult, PageReconstructError> {
+) -> Result<TimelineAtLsnSizeResult, RecvError> {
     let _permit = tokio::sync::Semaphore::acquire_owned(limit)
         .await
-        .expect("global semaphore should not have been closed");
+        .expect("global semaphore should not had been closed");
 
-    let size_res = timeline.calculate_logical_size(lsn, ctx).await;
+    let size_res = timeline
+        .spawn_ondemand_logical_size_calculation(lsn)
+        .await?;
     Ok(TimelineAtLsnSizeResult(timeline, lsn, size_res))
 }
 

pageserver/src/tenant/storage_layer.rs

@@ -109,7 +109,7 @@ pub trait Layer: Send + Sync {
     /// See PageReconstructResult for possible return values. The collected data
     /// is appended to reconstruct_data; the caller should pass an empty struct
     /// on first call, or a struct with a cached older image of the page if one
-    /// is available. If this returns ValueReconstructResult::Continue, look up
+    /// is available. If this returns PageReconstructResult::Continue, look up
     /// the predecessor layer and call again with the same 'reconstruct_data' to
     /// collect more data.
     fn get_value_reconstruct_data(
@@ -187,6 +187,12 @@ pub trait PersistentLayer: Layer {
     fn file_size(&self) -> Option<u64>;
 }
 
+impl PartialEq for dyn PersistentLayer {
+    fn eq(&self, other: &Self) -> bool {
+        self.filename().eq(&other.filename())
+    }
+}
+
 pub fn downcast_remote_layer(
     layer: &Arc<dyn PersistentLayer>,
 ) -> Option<std::sync::Arc<RemoteLayer>> {
@@ -196,3 +202,56 @@ 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 PartialEq for LayerDescriptor {
+    fn eq(&self, other: &Self) -> bool {
+        self.short_id == other.short_id
+    }
+}
+
+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/tasks.rs

@@ -1,39 +1,45 @@
 //! This module contains functions to serve per-tenant background processes,
 //! such as compaction and GC
 
+use std::ops::ControlFlow;
 use std::sync::Arc;
 use std::time::Duration;
 
-use crate::context::{DownloadBehavior, RequestContext, TaskKind};
 use crate::metrics::TENANT_TASK_EVENTS;
 use crate::task_mgr;
-use crate::task_mgr::BACKGROUND_RUNTIME;
-use crate::tenant::Tenant;
+use crate::task_mgr::{TaskKind, BACKGROUND_RUNTIME};
+use crate::tenant::mgr;
+use crate::tenant::{Tenant, TenantState};
 use tracing::*;
+use utils::id::TenantId;
 
-pub fn start_background_loops(tenant: &Arc<Tenant>) {
-    let tenant_id = tenant.tenant_id;
-
-    let tenant_clone = Arc::clone(tenant);
+pub fn start_background_loops(tenant_id: TenantId) {
     task_mgr::spawn(
         BACKGROUND_RUNTIME.handle(),
+        TaskKind::Compaction,
+        Some(tenant_id),
+        None,
         &format!("compactor for tenant {tenant_id}"),
         false,
         async move {
-            compaction_loop(&tenant_clone)
+            compaction_loop(tenant_id)
                 .instrument(info_span!("compaction_loop", tenant_id = %tenant_id))
                 .await;
+            Ok(())
         },
     );
-    let tenant_clone = Arc::clone(tenant);
     task_mgr::spawn(
         BACKGROUND_RUNTIME.handle(),
+        TaskKind::GarbageCollector,
+        Some(tenant_id),
+        None,
         &format!("garbage collector for tenant {tenant_id}"),
         false,
         async move {
-            gc_loop(&tenant_clone)
+            gc_loop(tenant_id)
                 .instrument(info_span!("gc_loop", tenant_id = %tenant_id))
                 .await;
+            Ok(())
         },
     );
 }
@@ -41,27 +47,25 @@ pub fn start_background_loops(tenant: &Arc<Tenant>) {
 ///
 /// Compaction task's main loop
 ///
-async fn compaction_loop(tenant: &Arc<Tenant>) {
+async fn compaction_loop(tenant_id: TenantId) {
     let wait_duration = Duration::from_secs(2);
     info!("starting");
     TENANT_TASK_EVENTS.with_label_values(&["start"]).inc();
     async {
-        let top_ctx = RequestContext::new(TaskKind::Compaction, DownloadBehavior::Download);
-
-        let tenant_ctx = match tenant.get_context(&top_ctx) {
-            Ok(ctx) => ctx,
-            Err(state) => {
-                // This could happen if the tenant is detached or the pageserver is shut
-                // down immediately after loading or attaching completed and the tenant
-                // was activated. It seems unlikely enough in practice that we better print
-                // a warning, as it could also be a bug.
-                error!("Not running compaction loop, tenant is not active: {state:?}");
-                return;
-            }
-        };
         loop {
             trace!("waking up");
 
+            let tenant = tokio::select! {
+                _ = task_mgr::shutdown_watcher() => {
+                    info!("received cancellation request");
+                    return;
+                },
+                tenant_wait_result = wait_for_active_tenant(tenant_id, wait_duration) => match tenant_wait_result {
+                    ControlFlow::Break(()) => return,
+                    ControlFlow::Continue(tenant) => tenant,
+                },
+            };
+
             let mut sleep_duration = tenant.get_compaction_period();
             if sleep_duration == Duration::ZERO {
                 info!("automatic compaction is disabled");
@@ -69,7 +73,7 @@ async fn compaction_loop(tenant: &Arc<Tenant>) {
                 sleep_duration = Duration::from_secs(10);
             } else {
                 // Run compaction
-                if let Err(e) = tenant.compaction_iteration(&tenant_ctx).await {
+                if let Err(e) = tenant.compaction_iteration().await {
                     sleep_duration = wait_duration;
                     error!("Compaction failed, retrying in {:?}: {e:?}", sleep_duration);
                 }
@@ -77,9 +81,9 @@ async fn compaction_loop(tenant: &Arc<Tenant>) {
 
             // Sleep
             tokio::select! {
-                _ = tenant_ctx.cancelled() => {
+                _ = task_mgr::shutdown_watcher() => {
                     info!("received cancellation request during idling");
-                    break;
+                    break ;
                 },
                 _ = tokio::time::sleep(sleep_duration) => {},
             }
@@ -94,28 +98,25 @@ async fn compaction_loop(tenant: &Arc<Tenant>) {
 ///
 /// GC task's main loop
 ///
-async fn gc_loop(tenant: &Arc<Tenant>) {
+async fn gc_loop(tenant_id: TenantId) {
     let wait_duration = Duration::from_secs(2);
     info!("starting");
     TENANT_TASK_EVENTS.with_label_values(&["start"]).inc();
     async {
-        // GC might require downloading, to find the cutoff LSN that corresponds to the
-        // cutoff specified as time.
-        let top_ctx = RequestContext::new(TaskKind::GarbageCollector, DownloadBehavior::Download);
-        let tenant_ctx = match tenant.get_context(&top_ctx) {
-            Ok(ctx) => ctx,
-            Err(state) => {
-                // This could happen if the tenant is detached or the pageserver is shut
-                // down immediately after loading or attaching completed and the tenant
-                // was activated. It seems unlikely enough in practice that we better print
-                // a warning, as it could also be a bug.
-                error!("Not running GC loop, tenant is not active: {state:?}");
-                return;
-            }
-        };
         loop {
             trace!("waking up");
 
+            let tenant = tokio::select! {
+                _ = task_mgr::shutdown_watcher() => {
+                    info!("received cancellation request");
+                    return;
+                },
+                tenant_wait_result = wait_for_active_tenant(tenant_id, wait_duration) => match tenant_wait_result {
+                    ControlFlow::Break(()) => return,
+                    ControlFlow::Continue(tenant) => tenant,
+                },
+            };
+
             let gc_period = tenant.get_gc_period();
             let gc_horizon = tenant.get_gc_horizon();
             let mut sleep_duration = gc_period;
@@ -126,10 +127,7 @@ async fn gc_loop(tenant: &Arc<Tenant>) {
             } else {
                 // Run gc
                 if gc_horizon > 0 {
-                    // Run compaction
-                    if let Err(e) = tenant
-                        .gc_iteration(None, gc_horizon, tenant.get_pitr_interval(), &tenant_ctx)
-                        .await
+                    if let Err(e) = tenant.gc_iteration(None, gc_horizon, tenant.get_pitr_interval()).await
                     {
                         sleep_duration = wait_duration;
                         error!("Gc failed, retrying in {:?}: {e:?}", sleep_duration);
@@ -139,7 +137,7 @@ async fn gc_loop(tenant: &Arc<Tenant>) {
 
             // Sleep
             tokio::select! {
-                _ = tenant_ctx.cancelled() => {
+                _ = task_mgr::shutdown_watcher() => {
                     info!("received cancellation request during idling");
                     break;
                 },
@@ -151,3 +149,46 @@ async fn gc_loop(tenant: &Arc<Tenant>) {
     TENANT_TASK_EVENTS.with_label_values(&["stop"]).inc();
     trace!("GC loop stopped.");
 }
+
+async fn wait_for_active_tenant(
+    tenant_id: TenantId,
+    wait: Duration,
+) -> ControlFlow<(), Arc<Tenant>> {
+    let tenant = loop {
+        match mgr::get_tenant(tenant_id, false).await {
+            Ok(tenant) => break tenant,
+            Err(e) => {
+                error!("Failed to get a tenant {tenant_id}: {e:#}");
+                tokio::time::sleep(wait).await;
+            }
+        }
+    };
+
+    // if the tenant has a proper status already, no need to wait for anything
+    if tenant.current_state() == TenantState::Active {
+        ControlFlow::Continue(tenant)
+    } else {
+        let mut tenant_state_updates = tenant.subscribe_for_state_updates();
+        loop {
+            match tenant_state_updates.changed().await {
+                Ok(()) => {
+                    let new_state = *tenant_state_updates.borrow();
+                    match new_state {
+                        TenantState::Active => {
+                            debug!("Tenant state changed to active, continuing the task loop");
+                            return ControlFlow::Continue(tenant);
+                        }
+                        state => {
+                            debug!("Not running the task loop, tenant is not active: {state:?}");
+                            continue;
+                        }
+                    }
+                }
+                Err(_sender_dropped_error) => {
+                    info!("Tenant dropped the state updates sender, quitting waiting for tenant and the task loop");
+                    return ControlFlow::Break(());
+                }
+            }
+        }
+    }
+}

pageserver/src/tenant/upload_queue.rs

@@ -4,7 +4,6 @@ use super::storage_layer::LayerFileName;
 use crate::tenant::metadata::TimelineMetadata;
 use crate::tenant::remote_timeline_client::index::IndexPart;
 use crate::tenant::remote_timeline_client::index::LayerFileMetadata;
-use crate::tenant::TimelineRequestContext;
 use std::collections::{HashMap, VecDeque};
 use std::fmt::Debug;
 
@@ -74,13 +73,6 @@ pub(crate) struct UploadQueueInitialized {
     /// tasks to finish. For example, metadata upload cannot be performed before all
     /// preceding layer file uploads have completed.
     pub(crate) queued_operations: VecDeque<UploadOp>,
-
-    /// Context used for the upload tasks. Note that this is associated with the
-    /// Timeline, so this prevents the Timeline from being shut down. To ensure quick
-    /// shutdown, RemoteTimelineClient spawns a task to wait for cancellation on the
-    /// context and stop the queue. Otherwise we woudn't notice the cancellation
-    /// until next upload attempt.
-    pub(crate) upload_ctx: TimelineRequestContext,
 }
 
 pub(crate) struct UploadQueueStopped {
@@ -91,7 +83,6 @@ impl UploadQueue {
     pub(crate) fn initialize_empty_remote(
         &mut self,
         metadata: &TimelineMetadata,
-        upload_ctx: TimelineRequestContext,
     ) -> anyhow::Result<&mut UploadQueueInitialized> {
         match self {
             UploadQueue::Uninitialized => (),
@@ -117,7 +108,6 @@ impl UploadQueue {
             num_inprogress_deletions: 0,
             inprogress_tasks: HashMap::new(),
             queued_operations: VecDeque::new(),
-            upload_ctx,
         };
 
         *self = UploadQueue::Initialized(state);
@@ -127,7 +117,6 @@ impl UploadQueue {
     pub(crate) fn initialize_with_current_remote_index_part(
         &mut self,
         index_part: &IndexPart,
-        upload_ctx: TimelineRequestContext,
     ) -> anyhow::Result<&mut UploadQueueInitialized> {
         match self {
             UploadQueue::Uninitialized => (),
@@ -164,7 +153,6 @@ impl UploadQueue {
             num_inprogress_deletions: 0,
             inprogress_tasks: HashMap::new(),
             queued_operations: VecDeque::new(),
-            upload_ctx,
         };
 
         *self = UploadQueue::Initialized(state);

pageserver/src/walreceiver.rs

@@ -31,7 +31,6 @@ use once_cell::sync::OnceCell;
 use std::future::Future;
 use storage_broker::BrokerClientChannel;
 use tokio::sync::watch;
-use tokio_util::sync::CancellationToken;
 use tracing::*;
 
 pub use connection_manager::spawn_connection_manager_task;
@@ -77,7 +76,7 @@ pub fn is_broker_client_initialized() -> bool {
 
 /// A handle of an asynchronous task.
 /// The task has a channel that it can use to communicate its lifecycle events in a certain form, see [`TaskEvent`]
-/// and a cancellation token that it can listen to for earlier interrupts.
+/// and a cancellation channel that it can listen to for earlier interrupts.
 ///
 /// Note that the communication happens via the `watch` channel, that does not accumulate the events, replacing the old one with the never one on submission.
 /// That may lead to certain events not being observed by the listener.
@@ -85,7 +84,7 @@ pub fn is_broker_client_initialized() -> bool {
 pub struct TaskHandle<E> {
     join_handle: Option<tokio::task::JoinHandle<anyhow::Result<()>>>,
     events_receiver: watch::Receiver<TaskStateUpdate<E>>,
-    cancellation: CancellationToken,
+    cancellation: watch::Sender<()>,
 }
 
 pub enum TaskEvent<E> {
@@ -103,18 +102,20 @@ pub enum TaskStateUpdate<E> {
 impl<E: Clone> TaskHandle<E> {
     /// Initializes the task, starting it immediately after the creation.
     pub fn spawn<Fut>(
-        task: impl FnOnce(watch::Sender<TaskStateUpdate<E>>) -> Fut + Send + 'static,
-        cancellation: CancellationToken,
+        task: impl FnOnce(watch::Sender<TaskStateUpdate<E>>, watch::Receiver<()>) -> Fut
+            + Send
+            + 'static,
     ) -> Self
     where
         Fut: Future<Output = anyhow::Result<()>> + Send,
         E: Send + Sync + 'static,
     {
+        let (cancellation, cancellation_receiver) = watch::channel(());
         let (events_sender, events_receiver) = watch::channel(TaskStateUpdate::Started);
 
         let join_handle = WALRECEIVER_RUNTIME.spawn(async move {
             events_sender.send(TaskStateUpdate::Started).ok();
-            task(events_sender).await
+            task(events_sender, cancellation_receiver).await
         });
 
         TaskHandle {
@@ -156,7 +157,7 @@ impl<E: Clone> TaskHandle<E> {
     /// Aborts current task, waiting for it to finish.
     pub async fn shutdown(self) {
         if let Some(jh) = self.join_handle {
-            self.cancellation.cancel();
+            self.cancellation.send(()).ok();
             match jh.await {
                 Ok(Ok(())) => debug!("Shutdown success"),
                 Ok(Err(e)) => error!("Shutdown task error: {e:?}"),

pageserver/src/walreceiver/connection_manager.rs

@@ -11,9 +11,9 @@
 
 use std::{collections::HashMap, num::NonZeroU64, ops::ControlFlow, sync::Arc, time::Duration};
 
-use crate::context::{DownloadBehavior, RequestContext, TaskKind};
+use crate::task_mgr::TaskKind;
 use crate::task_mgr::WALRECEIVER_RUNTIME;
-use crate::tenant::{Timeline, TimelineRequestContext};
+use crate::tenant::Timeline;
 use crate::{task_mgr, walreceiver::TaskStateUpdate};
 use anyhow::Context;
 use chrono::{NaiveDateTime, Utc};
@@ -46,7 +46,6 @@ pub fn spawn_connection_manager_task(
     lagging_wal_timeout: Duration,
     max_lsn_wal_lag: NonZeroU64,
     auth_token: Option<Arc<String>>,
-    ctx: TimelineRequestContext,
 ) {
     let mut broker_client = get_broker_client().clone();
 
@@ -55,6 +54,9 @@ pub fn spawn_connection_manager_task(
 
     task_mgr::spawn(
         WALRECEIVER_RUNTIME.handle(),
+        TaskKind::WalReceiverManager,
+        Some(tenant_id),
+        Some(timeline_id),
         &format!("walreceiver for timeline {tenant_id}/{timeline_id}"),
         false,
         async move {
@@ -68,21 +70,20 @@ pub fn spawn_connection_manager_task(
             );
             loop {
                 select! {
-                    _ = ctx.cancelled() => {
+                    _ = task_mgr::shutdown_watcher() => {
                         info!("WAL receiver shutdown requested, shutting down");
                         walreceiver_state.shutdown().await;
-                        return;
+                        return Ok(());
                     },
                     loop_step_result = connection_manager_loop_step(
                         &mut broker_client,
                         &mut walreceiver_state,
-                        &ctx,
                     ) => match loop_step_result {
                         ControlFlow::Continue(()) => continue,
                         ControlFlow::Break(()) => {
                             info!("Connection manager loop ended, shutting down");
                             walreceiver_state.shutdown().await;
-                            return;
+                            return Ok(());
                         }
                     },
                 }
@@ -100,7 +101,6 @@ pub fn spawn_connection_manager_task(
 async fn connection_manager_loop_step(
     broker_client: &mut BrokerClientChannel,
     walreceiver_state: &mut WalreceiverState,
-    ctx: &TimelineRequestContext,
 ) -> ControlFlow<(), ()> {
     let mut timeline_state_updates = walreceiver_state.timeline.subscribe_for_state_updates();
 
@@ -226,7 +226,6 @@ async fn connection_manager_loop_step(
                 .change_connection(
                     new_candidate.safekeeper_id,
                     new_candidate.wal_source_connconf,
-                    ctx,
                 )
                 .await
         }
@@ -390,38 +389,26 @@ impl WalreceiverState {
         &mut self,
         new_sk_id: NodeId,
         new_wal_source_connconf: PgConnectionConfig,
-        ctx: &TimelineRequestContext,
     ) {
         self.drop_old_connection(true).await;
 
         let id = self.id;
         let connect_timeout = self.wal_connect_timeout;
         let timeline = Arc::clone(&self.timeline);
-
-        let child_ctx = ctx.register_another(RequestContext::with_parent(
-            TaskKind::WalReceiverConnection,
-            DownloadBehavior::Download,
-            ctx,
-        ));
-        let cancellation_token = child_ctx.cancellation_token().clone();
-
-        let connection_handle = TaskHandle::spawn(
-            move |events_sender| {
-                async move {
-                    super::walreceiver_connection::handle_walreceiver_connection(
-                        timeline,
-                        new_wal_source_connconf,
-                        events_sender,
-                        connect_timeout,
-                        child_ctx,
-                    )
-                    .await
-                    .context("walreceiver connection handling failure")
-                }
-                .instrument(info_span!("walreceiver_connection", id = %id, node_id = %new_sk_id))
-            },
-            cancellation_token,
-        );
+        let connection_handle = TaskHandle::spawn(move |events_sender, cancellation| {
+            async move {
+                super::walreceiver_connection::handle_walreceiver_connection(
+                    timeline,
+                    new_wal_source_connconf,
+                    events_sender,
+                    cancellation,
+                    connect_timeout,
+                )
+                .await
+                .context("walreceiver connection handling failure")
+            }
+            .instrument(info_span!("walreceiver_connection", id = %id, node_id = %new_sk_id))
+        });
 
         let now = Utc::now().naive_utc();
         self.wal_connection = Some(WalConnection {
@@ -833,7 +820,6 @@ fn wal_stream_connection_config(
 mod tests {
     use super::*;
     use crate::tenant::harness::{TenantHarness, TIMELINE_ID};
-    use tokio_util::sync::CancellationToken;
     use url::Host;
 
     fn dummy_broker_sk_timeline(
@@ -914,15 +900,12 @@ mod tests {
             started_at: now,
             sk_id: connected_sk_id,
             status: connection_status,
-            connection_task: TaskHandle::spawn(
-                move |sender| async move {
-                    sender
-                        .send(TaskStateUpdate::Progress(connection_status))
-                        .ok();
-                    Ok(())
-                },
-                CancellationToken::new(),
-            ),
+            connection_task: TaskHandle::spawn(move |sender, _| async move {
+                sender
+                    .send(TaskStateUpdate::Progress(connection_status))
+                    .ok();
+                Ok(())
+            }),
             discovered_new_wal: None,
         });
         state.wal_stream_candidates = HashMap::from([
@@ -1079,15 +1062,12 @@ mod tests {
             started_at: now,
             sk_id: connected_sk_id,
             status: connection_status,
-            connection_task: TaskHandle::spawn(
-                move |sender| async move {
-                    sender
-                        .send(TaskStateUpdate::Progress(connection_status))
-                        .ok();
-                    Ok(())
-                },
-                CancellationToken::new(),
-            ),
+            connection_task: TaskHandle::spawn(move |sender, _| async move {
+                sender
+                    .send(TaskStateUpdate::Progress(connection_status))
+                    .ok();
+                Ok(())
+            }),
             discovered_new_wal: None,
         });
         state.wal_stream_candidates = HashMap::from([
@@ -1147,15 +1127,12 @@ mod tests {
             started_at: now,
             sk_id: NodeId(1),
             status: connection_status,
-            connection_task: TaskHandle::spawn(
-                move |sender| async move {
-                    sender
-                        .send(TaskStateUpdate::Progress(connection_status))
-                        .ok();
-                    Ok(())
-                },
-                CancellationToken::new(),
-            ),
+            connection_task: TaskHandle::spawn(move |sender, _| async move {
+                sender
+                    .send(TaskStateUpdate::Progress(connection_status))
+                    .ok();
+                Ok(())
+            }),
             discovered_new_wal: None,
         });
         state.wal_stream_candidates = HashMap::from([(
@@ -1212,10 +1189,7 @@ mod tests {
             started_at: now,
             sk_id: NodeId(1),
             status: connection_status,
-            connection_task: TaskHandle::spawn(
-                move |_| async move { Ok(()) },
-                CancellationToken::new(),
-            ),
+            connection_task: TaskHandle::spawn(move |_, _| async move { Ok(()) }),
             discovered_new_wal: Some(NewCommittedWAL {
                 discovered_at: time_over_threshold,
                 lsn: new_lsn,
@@ -1259,18 +1233,18 @@ mod tests {
     const DUMMY_SAFEKEEPER_HOST: &str = "safekeeper_connstr";
 
     async fn dummy_state(harness: &TenantHarness<'_>) -> WalreceiverState {
-        let (tenant, tenant_ctx) = harness.load().await;
-        let (timeline, timeline_ctx) = tenant
-            .create_empty_timeline(TIMELINE_ID, Lsn(0), crate::DEFAULT_PG_VERSION, &tenant_ctx)
-            .expect("Failed to create an empty timeline for dummy wal connection manager");
-        let timeline = timeline.initialize(&timeline_ctx).unwrap();
-
         WalreceiverState {
             id: TenantTimelineId {
                 tenant_id: harness.tenant_id,
                 timeline_id: TIMELINE_ID,
             },
-            timeline,
+            timeline: harness
+                .load()
+                .await
+                .create_empty_timeline(TIMELINE_ID, Lsn(0), crate::DEFAULT_PG_VERSION)
+                .expect("Failed to create an empty timeline for dummy wal connection manager")
+                .initialize()
+                .unwrap(),
             wal_connect_timeout: Duration::from_secs(1),
             lagging_wal_timeout: Duration::from_secs(1),
             max_lsn_wal_lag: NonZeroU64::new(1024 * 1024).unwrap(),

pageserver/src/walreceiver/walreceiver_connection.rs

@@ -24,8 +24,9 @@ use tracing::{debug, error, info, trace, warn};
 use crate::{metrics::LIVE_CONNECTIONS_COUNT, walreceiver::TaskStateUpdate};
 use crate::{
     task_mgr,
+    task_mgr::TaskKind,
     task_mgr::WALRECEIVER_RUNTIME,
-    tenant::{Timeline, TimelineRequestContext, WalReceiverInfo},
+    tenant::{Timeline, WalReceiverInfo},
     walingest::WalIngest,
     walrecord::DecodedWALRecord,
 };
@@ -58,8 +59,8 @@ pub async fn handle_walreceiver_connection(
     timeline: Arc<Timeline>,
     wal_source_connconf: PgConnectionConfig,
     events_sender: watch::Sender<TaskStateUpdate<WalConnectionStatus>>,
+    mut cancellation: watch::Receiver<()>,
     connect_timeout: Duration,
-    ctx: TimelineRequestContext,
 ) -> anyhow::Result<()> {
     // Connect to the database in replication mode.
     info!("connecting to {wal_source_connconf:?}");
@@ -97,9 +98,12 @@ pub async fn handle_walreceiver_connection(
 
     // The connection object performs the actual communication with the database,
     // so spawn it off to run on its own.
-    let cancellation_token = ctx.cancellation_token().clone();
+    let mut connection_cancellation = cancellation.clone();
     task_mgr::spawn(
         WALRECEIVER_RUNTIME.handle(),
+        TaskKind::WalReceiverConnection,
+        Some(timeline.tenant_id),
+        Some(timeline.timeline_id),
         "walreceiver connection",
         false,
         async move {
@@ -113,8 +117,9 @@ pub async fn handle_walreceiver_connection(
                     }
                 },
 
-                _ = cancellation_token.cancelled() => info!("Connection cancelled"),
+                _ = connection_cancellation.changed() => info!("Connection cancelled"),
             }
+            Ok(())
         },
     );
 
@@ -174,13 +179,11 @@ pub async fn handle_walreceiver_connection(
 
     let mut waldecoder = WalStreamDecoder::new(startpoint, timeline.pg_version);
 
-    let mut walingest = WalIngest::new(timeline.as_ref(), startpoint, &ctx).await?;
-
-    let cancellation = ctx.cancellation_token().clone();
+    let mut walingest = WalIngest::new(timeline.as_ref(), startpoint).await?;
 
     while let Some(replication_message) = {
         select! {
-            _ = cancellation.cancelled() => {
+            _ = cancellation.changed() => {
                 info!("walreceiver interrupted");
                 None
             }
@@ -247,7 +250,7 @@ pub async fn handle_walreceiver_connection(
                         ensure!(lsn.is_aligned());
 
                         walingest
-                            .ingest_record(recdata, lsn, &mut modification, &mut decoded, &ctx)
+                            .ingest_record(recdata.clone(), lsn, &mut modification, &mut decoded)
                             .await
                             .with_context(|| format!("could not ingest record at {lsn}"))?;
 
@@ -325,7 +328,7 @@ pub async fn handle_walreceiver_connection(
             // Send the replication feedback message.
             // Regular standby_status_update fields are put into this message.
             let (timeline_logical_size, _) = timeline
-                .get_current_logical_size(&ctx)
+                .get_current_logical_size()
                 .context("Status update creation failed to get current logical size")?;
             let status_update = ReplicationFeedback {
                 current_timeline_size: timeline_logical_size,

pgxn/neon/pagestore_smgr.c

@@ -52,7 +52,6 @@
 #include "access/xlogdefs.h"
 #include "catalog/pg_class.h"
 #include "common/hashfn.h"
-#include "executor/instrument.h"
 #include "pagestore_client.h"
 #include "postmaster/interrupt.h"
 #include "postmaster/autovacuum.h"
@@ -251,6 +250,11 @@ PrefetchState *MyPState;
 	) \
 )
 
+int			n_prefetch_hits = 0;
+int			n_prefetch_misses = 0;
+int			n_prefetch_missed_caches = 0;
+int			n_prefetch_dupes = 0;
+
 XLogRecPtr	prefetch_lsn = 0;
 
 static bool compact_prefetch_buffers(void);
@@ -766,7 +770,7 @@ prefetch_register_buffer(BufferTag tag, bool *force_latest, XLogRecPtr *force_ls
 		else
 		{
 			/* The buffered request is good enough, return that index */
-			pgBufferUsage.prefetch.duplicates++;
+			n_prefetch_dupes++;
 			return ring_index;
 		}
 	}
@@ -1841,7 +1845,7 @@ neon_read_at_lsn(RelFileNode rnode, ForkNumber forkNum, BlockNumber blkno,
 		if (slot->effective_request_lsn >= request_lsn)
 		{
 			ring_index = slot->my_ring_index;
-			pgBufferUsage.prefetch.hits += 1;
+			n_prefetch_hits += 1;
 		}
 		else /* the current prefetch LSN is not large enough, so drop the prefetch */
 		{
@@ -1856,7 +1860,7 @@ neon_read_at_lsn(RelFileNode rnode, ForkNumber forkNum, BlockNumber blkno,
 			}
 			/* drop caches */
 			prefetch_set_unused(slot->my_ring_index);
-			pgBufferUsage.prefetch.expired += 1;
+			n_prefetch_missed_caches += 1;
 			/* make it look like a prefetch cache miss */
 			entry = NULL;
 		}
@@ -1866,7 +1870,7 @@ neon_read_at_lsn(RelFileNode rnode, ForkNumber forkNum, BlockNumber blkno,
 	{
 		if (entry == NULL)
 		{
-			pgBufferUsage.prefetch.misses += 1;
+			n_prefetch_misses += 1;
 
 			ring_index = prefetch_register_buffer(buftag, &request_latest,
 												  &request_lsn);

rust-toolchain.toml

@@ -1,5 +1,10 @@
 [toolchain]
-channel = "1.66.1"
+# We try to stick to a toolchain version that is widely available on popular distributions, so that most people
+# can use the toolchain that comes with their operating system. But if there's a feature we miss badly from a later
+# 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"
 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/src/timelines_global_map.rs

@@ -279,9 +279,7 @@ impl GlobalTimelines {
                 let (dir_existed, was_active) = timeline.delete_from_disk(&mut shared_state)?;
 
                 // Remove timeline from the map.
-                // FIXME: re-enable it once we fix the issue with recreation of deleted timelines
-                // https://github.com/neondatabase/neon/issues/3146
-                // TIMELINES_STATE.lock().unwrap().timelines.remove(ttid);
+                TIMELINES_STATE.lock().unwrap().timelines.remove(ttid);
 
                 Ok(TimelineDeleteForceResult {
                     dir_existed,
@@ -348,16 +346,15 @@ impl GlobalTimelines {
                 .tenant_dir(tenant_id),
         )?;
 
-        // FIXME: we temporarily disabled removing timelines from the map, see `delete_force`
-        // let tlis_after_delete = Self::get_all_for_tenant(*tenant_id);
-        // if !tlis_after_delete.is_empty() {
-        //     // Some timelines were created while we were deleting them, returning error
-        //     // to the caller, so it can retry later.
-        //     bail!(
-        //         "failed to delete all timelines for tenant {}: some timelines were created while we were deleting them",
-        //         tenant_id
-        //     );
-        // }
+        let tlis_after_delete = Self::get_all_for_tenant(*tenant_id);
+        if !tlis_after_delete.is_empty() {
+            // Some timelines were created while we were deleting them, returning error
+            // to the caller, so it can retry later.
+            bail!(
+                "failed to delete all timelines for tenant {}: some timelines were created while we were deleting them",
+                tenant_id
+            );
+        }
 
         Ok(deleted)
     }

test_runner/fixtures/neon_fixtures.py

@@ -1914,9 +1914,9 @@ class NeonPageserver(PgProtocol):
             ".*Shutdown task error: walreceiver connection handling failure.*",
             ".*wal_connection_manager.*tcp connect error: Connection refused.*",
             ".*query handler for .* failed: Socket IO error: Connection reset by peer.*",
-            ".*connection handler exited with error: Postgres connection error.*",
-            ".*connection handler exited with error: Connection reset by peer.*",
-            ".*connection handler exited with error: Postgres query error.*",
+            ".*serving compute connection task.*exited with error: Postgres connection error.*",
+            ".*serving compute connection task.*exited with error: Connection reset by peer.*",
+            ".*serving compute connection task.*exited with error: Postgres query error.*",
             ".*Connection aborted: connection error: error communicating with the server: Broken pipe.*",
             ".*Connection aborted: connection error: error communicating with the server: Transport endpoint is not connected.*",
             ".*Connection aborted: connection error: error communicating with the server: Connection reset by peer.*",
@@ -1945,7 +1945,6 @@ class NeonPageserver(PgProtocol):
             # Tenant::delete_timeline() can cause any of the four following errors.
             # FIXME: we shouldn't be considering it an error: https://github.com/neondatabase/neon/issues/2946
             ".*could not flush frozen layer.*queue is in state Stopped",  # when schedule layer upload fails because queued got closed before compaction got killed
-            ".*Compaction failed, .* queue is in state Stopped.*",  # similarly for compaction
             ".*wait for layer upload ops to complete.*",  # .*Caused by:.*wait_completion aborted because upload queue was stopped
             ".*gc_loop.*Gc failed, retrying in.*timeline is Stopping",  # When gc checks timeline state after acquiring layer_removal_cs
             ".*compaction_loop.*Compaction failed, retrying in.*timeline is Stopping",  # When compaction checks timeline state after acquiring layer_removal_cs

test_runner/performance/test_perf_pgbench.py

@@ -172,8 +172,10 @@ def get_scales_matrix(default: int = 10) -> List[int]:
 @pytest.mark.parametrize("duration", get_durations_matrix())
 def test_pgbench(neon_with_baseline: PgCompare, scale: int, duration: int):
     run_test_pgbench(neon_with_baseline, scale, duration, PgBenchLoadType.INIT)
-    run_test_pgbench(neon_with_baseline, scale, duration, PgBenchLoadType.SIMPLE_UPDATE)
-    run_test_pgbench(neon_with_baseline, scale, duration, PgBenchLoadType.SELECT_ONLY)
+    run_test_pgbench(neon_with_baseline, scale, duration, PgBenchLoadType.INIT)
+    run_test_pgbench(neon_with_baseline, scale, duration, PgBenchLoadType.INIT)
+    # run_test_pgbench(neon_with_baseline, scale, duration, PgBenchLoadType.SIMPLE_UPDATE)
+    # run_test_pgbench(neon_with_baseline, scale, duration, PgBenchLoadType.SELECT_ONLY)
 
 
 # The following 3 tests run on an existing database as it was set up by previous tests,

test_runner/regress/test_broken_timeline.py

@@ -15,7 +15,7 @@ def test_broken_timeline(neon_env_builder: NeonEnvBuilder):
 
     env.pageserver.allowed_errors.extend(
         [
-            ".*Failed to load delta layer.*",
+            ".*Failed to reconstruct the page.*",
             ".*could not find data for key.*",
             ".*is not active. Current state: Broken.*",
             ".*will not become active. Current state: Broken.*",
@@ -99,7 +99,7 @@ def test_broken_timeline(neon_env_builder: NeonEnvBuilder):
     # Third timeline will also fail during basebackup, because the layer file is corrupt.
     # It will fail when we try to read (and reconstruct) a page from it, ergo the error message.
     # (We don't check layer file contents on startup, when loading the timeline)
-    with pytest.raises(Exception, match="Failed to load delta layer") as err:
+    with pytest.raises(Exception, match="Failed to reconstruct the page") as err:
         pg3.start()
     log.info(
         f"As expected, compute startup failed for timeline {tenant3}/{timeline3} with corrupt layers: {err}"

test_runner/regress/test_ondemand_download.py

@@ -87,22 +87,8 @@ def test_ondemand_download_large_rel(
     # run checkpoint manually to be sure that data landed in remote storage
     client.timeline_checkpoint(tenant_id, timeline_id)
 
-    # wait until pageserver successfully uploaded all the data, and the result
-    # of the compaction, to remote storage
+    # wait until pageserver successfully uploaded a checkpoint to remote storage
     wait_for_upload(client, tenant_id, timeline_id, current_lsn)
-    def assert_queue_empty(file_kind, op_kind):
-        val = client.get_remote_timeline_client_metric(
-            "pageserver_remote_timeline_client_calls_unfinished",
-            tenant_id,
-            timeline_id,
-            file_kind,
-            op_kind,
-        )
-        assert val is not None, "expecting metric to be present"
-        assert int(val) == 0
-    wait_until(60, 0.5, lambda: assert_queue_empty(file_kind="layer", op_kind="upload"))
-    wait_until(60, 0.5, lambda: assert_queue_empty(file_kind="index", op_kind="upload"))
-    wait_until(60, 0.5, lambda: assert_queue_empty(file_kind="layer", op_kind="delete"))
     log.info("uploads have finished")
 
     ##### Stop the first pageserver instance, erase all its data
@@ -126,7 +112,7 @@ def test_ondemand_download_large_rel(
     # relation, so they are likely already downloaded. But the middle of the
     # table should not have been needed by anything yet.
     with pg.cursor() as cur:
-        assert query_scalar(cur, f"select count(*) from tbl where id = {num_rows}/2") == 1
+        assert query_scalar(cur, "select count(*) from tbl where id = 500000") == 1
 
     after_downloads = get_num_downloaded_layers(client, tenant_id, timeline_id)
     log.info(f"layers downloaded before {before_downloads} and after {after_downloads}")
@@ -385,7 +371,7 @@ def test_download_remote_layers_api(
     env.pageserver.allowed_errors.extend(
         [
             f".*download_all_remote_layers.*{tenant_id}.*{timeline_id}.*layer download failed.*remote-storage-download-pre-rename failpoint",
-            f".*initial size calculation for {tenant_id}/{timeline_id} failed.*remote-storage-download-pre-rename failpoint",
+            f".*initial size calculation.*{tenant_id}.*{timeline_id}.*Failed to calculate logical size",
         ]
     )
 

test_runner/regress/test_remote_storage.py

@@ -543,7 +543,6 @@ def test_timeline_deletion_with_files_stuck_in_upload_queue(
     )
 
     env = neon_env_builder.init_start()
-    env.pageserver.allowed_errors.append(".*request cancelled while waiting on uploads to finish.*")
 
     # create tenant with config that will determinstically allow
     # compaction and gc

test_runner/regress/test_tenant_detach.py

@@ -209,7 +209,6 @@ def test_tenant_reattach_while_busy(
     env.pageserver.allowed_errors.append(
         ".*Tenant .* will not become active\\. Current state: Stopping.*"
     )
-    env.pageserver.allowed_errors.append(".*Tenant .* is not active, state: Stopping.*")
 
     pageserver_http = env.pageserver.http_client()
 

test_runner/regress/test_wal_acceptor.py

@@ -567,8 +567,6 @@ def test_s3_wal_replay(neon_env_builder: NeonEnvBuilder, remote_storage_kind: Re
     for sk in env.safekeepers:
         cli = sk.http_client()
         cli.timeline_delete_force(tenant_id, timeline_id)
-        # restart safekeeper to clear its in-memory state
-        sk.stop().start()
         cli.timeline_create(tenant_id, timeline_id, pg_version, last_lsn)
         f_partial_path = (
             Path(sk.data_dir()) / str(tenant_id) / str(timeline_id) / f_partial_saved.name
@@ -1107,7 +1105,6 @@ def test_delete_force(neon_env_builder: NeonEnvBuilder, auth_enabled: bool):
     env.pageserver.allowed_errors.extend(
         [
             ".*Failed to process query for timeline .*: Timeline .* was not found in global map.*",
-            ".*Failed to process query for timeline .*: Timeline .* was cancelled and cannot be used anymore.*",
         ]
     )
 
@@ -1208,7 +1205,7 @@ def test_delete_force(neon_env_builder: NeonEnvBuilder, auth_enabled: bool):
 
     # Remove initial tenant again.
     response = sk_http.tenant_delete_force(tenant_id)
-    # assert response == {}
+    assert response == {}
     assert not (sk_data_dir / str(tenant_id)).exists()
     assert (sk_data_dir / str(tenant_id_other) / str(timeline_id_other)).is_dir()
 

workspace_hack/Cargo.toml

@@ -25,6 +25,7 @@ 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"] }
@@ -53,6 +54,7 @@ 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"] }