Fix prewarming terminatoin condition

libs/remote_storage/src/config.rs

@@ -1,9 +1,4 @@
-use std::{
-    fmt::Debug,
-    num::{NonZeroU32, NonZeroUsize},
-    str::FromStr,
-    time::Duration,
-};
+use std::{fmt::Debug, num::NonZeroUsize, str::FromStr, time::Duration};
 
 use aws_sdk_s3::types::StorageClass;
 use camino::Utf8PathBuf;
@@ -45,10 +40,7 @@ fn is_default_timeout(d: &Duration) -> bool {
 pub enum RemoteStorageKind {
     /// Storage based on local file system.
     /// Specify a root folder to place all stored files into.
-    LocalFs {
-        local_path: Utf8PathBuf,
-        max_keys_per_list_response: Option<NonZeroU32>,
-    },
+    LocalFs { local_path: Utf8PathBuf },
     /// AWS S3 based storage, storing all files in the S3 bucket
     /// specified by the config
     AwsS3(S3Config),
@@ -208,8 +200,7 @@ timeout = '5s'";
             config,
             RemoteStorageConfig {
                 storage: RemoteStorageKind::LocalFs {
-                    local_path: Utf8PathBuf::from("."),
-                    max_keys_per_list_response: None,
+                    local_path: Utf8PathBuf::from(".")
                 },
                 timeout: Duration::from_secs(5)
             }

libs/remote_storage/src/lib.rs

@@ -21,7 +21,7 @@ mod support;
 use std::{
     collections::HashMap,
     fmt::Debug,
-    num::{NonZeroU32, NonZeroUsize},
+    num::NonZeroU32,
     ops::Bound,
     pin::{pin, Pin},
     sync::Arc,
@@ -33,10 +33,10 @@ use camino::{Utf8Path, Utf8PathBuf};
 
 use bytes::Bytes;
 use futures::{stream::Stream, StreamExt};
-use itertools::Itertools;
+use itertools::Itertools as _;
 use serde::{Deserialize, Serialize};
 use tokio::sync::Semaphore;
-use tokio_util::{either, sync::CancellationToken};
+use tokio_util::sync::CancellationToken;
 use tracing::info;
 
 pub use self::{
@@ -162,7 +162,7 @@ pub struct ListingObject {
     pub size: u64,
 }
 
-#[derive(Debug, Default)]
+#[derive(Default)]
 pub struct Listing {
     pub prefixes: Vec<RemotePath>,
     pub keys: Vec<ListingObject>,
@@ -585,16 +585,9 @@ impl GenericRemoteStorage {
     pub async fn from_config(storage_config: &RemoteStorageConfig) -> anyhow::Result<Self> {
         let timeout = storage_config.timeout;
         Ok(match &storage_config.storage {
-            RemoteStorageKind::LocalFs {
-                local_path: path,
-                max_keys_per_list_response: max_keys_in_list_response,
-            } => {
+            RemoteStorageKind::LocalFs { local_path: path } => {
                 info!("Using fs root '{path}' as a remote storage");
-                Self::LocalFs(LocalFs::new(
-                    path.clone(),
-                    timeout,
-                    *max_keys_in_list_response,
-                )?)
+                Self::LocalFs(LocalFs::new(path.clone(), timeout)?)
             }
             RemoteStorageKind::AwsS3(s3_config) => {
                 // The profile and access key id are only printed here for debugging purposes,
@@ -706,60 +699,6 @@ impl ConcurrencyLimiter {
     }
 }
 
-impl Listing {
-    /// For testing, slice up [`Self`] into pages with `prefixes.len() + keys.len() <= max_keys`.
-    ///
-    /// Also, when paginating, anecdotally, S3 does {prefixes \u keys} | sort | paginate.
-    /// So, do the same here
-    pub fn paginate_like_s3(self, max_keys: Option<NonZeroU32>) -> impl Iterator<Item = Self> {
-        let Some(max_keys) = max_keys else {
-            return itertools::Either::Left(std::iter::once(self));
-        };
-        enum PrefixOrKey {
-            Prefix(RemotePath),
-            Key(ListingObject),
-        }
-        impl PrefixOrKey {
-            fn key(&self) -> impl Ord {
-                match self {
-                    Self::Prefix(prefix) => prefix.0.clone(),
-                    Self::Key(o) => o.key.0.clone(),
-                }
-            }
-        }
-
-        let mut combined: Vec<_> = self
-            .prefixes
-            .into_iter()
-            .map(PrefixOrKey::Prefix)
-            .chain(self.keys.into_iter().map(PrefixOrKey::Key))
-            .collect();
-        combined.sort_by_cached_key(|combined| combined.key());
-        itertools::Either::Right(
-            combined
-                .into_iter()
-                .chunks(
-                    usize::try_from(max_keys.get())
-                        .expect("true on the architectures we care about"),
-                )
-                .into_iter()
-                .map(|chunk| {
-                    let mut prefixes = Vec::new();
-                    let mut keys = Vec::new();
-                    for item in chunk {
-                        match item {
-                            PrefixOrKey::Prefix(prefix) => prefixes.push(prefix),
-                            PrefixOrKey::Key(key) => keys.push(key),
-                        }
-                    }
-                    Listing { prefixes, keys }
-                })
-                .collect::<Vec<_>>()
-                .into_iter(),
-        )
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

libs/remote_storage/src/local_fs.rs

@@ -36,17 +36,12 @@ const LOCAL_FS_TEMP_FILE_SUFFIX: &str = "___temp";
 pub struct LocalFs {
     storage_root: Utf8PathBuf,
     timeout: Duration,
-    max_keys_per_list_response: Option<NonZeroU32>,
 }
 
 impl LocalFs {
     /// Attempts to create local FS storage, along with its root directory.
     /// Storage root will be created (if does not exist) and transformed into an absolute path (if passed as relative).
-    pub fn new(
-        mut storage_root: Utf8PathBuf,
-        timeout: Duration,
-        max_keys_per_list_response: Option<NonZeroU32>,
-    ) -> anyhow::Result<Self> {
+    pub fn new(mut storage_root: Utf8PathBuf, timeout: Duration) -> anyhow::Result<Self> {
         if !storage_root.exists() {
             std::fs::create_dir_all(&storage_root).with_context(|| {
                 format!("Failed to create all directories in the given root path {storage_root:?}")
@@ -61,7 +56,6 @@ impl LocalFs {
         Ok(Self {
             storage_root,
             timeout,
-            max_keys_per_list_response,
         })
     }
 
@@ -154,12 +148,10 @@ impl LocalFs {
     // recursively lists all files in a directory,
     // mirroring the `list_files` for `s3_bucket`
     async fn list_recursive(&self, folder: Option<&RemotePath>) -> anyhow::Result<Vec<RemotePath>> {
-        dbg!(&folder);
         let full_path = match folder {
             Some(folder) => folder.with_base(&self.storage_root),
             None => self.storage_root.clone(),
         };
-        dbg!(&full_path);
 
         // If we were given a directory, we may use it as our starting point.
         // Otherwise, we must go up to the first ancestor dir that exists.  This is because
@@ -213,7 +205,6 @@ impl LocalFs {
                 let full_file_name = cur_folder.join(file_name);
                 if full_file_name.as_str().starts_with(prefix) {
                     let file_remote_path = self.local_file_to_relative_path(full_file_name.clone());
-                    dbg!(&file_remote_path);
                     files.push(file_remote_path);
                     if full_file_name.is_dir() {
                         directory_queue.push(full_file_name);
@@ -337,13 +328,6 @@ impl LocalFs {
 
         Ok(())
     }
-
-    fn effective_max_keys(&self, max_keys: Option<NonZeroU32>) -> Option<NonZeroU32> {
-        self.max_keys_per_list_response
-            .into_iter()
-            .chain(max_keys.into_iter())
-            .min()
-    }
 }
 
 impl RemoteStorage for LocalFs {
@@ -354,12 +338,8 @@ impl RemoteStorage for LocalFs {
         max_keys: Option<NonZeroU32>,
         cancel: &CancellationToken,
     ) -> impl Stream<Item = Result<Listing, DownloadError>> {
-        async_stream::try_stream! {
-            let listing = self.list(prefix, mode, max_keys, cancel).await?;
-            for page in listing.paginate_like_s3(self.effective_max_keys(max_keys)) {
-                yield dbg!(page);
-            }
-        }
+        let listing = self.list(prefix, mode, max_keys, cancel);
+        futures::stream::once(listing)
     }
 
     async fn list(
@@ -377,34 +357,35 @@ impl RemoteStorage for LocalFs {
                 .list_recursive(prefix)
                 .await
                 .map_err(DownloadError::Other)?;
-            let mut objects = Vec::with_capacity(keys.len());
-            for key in keys {
-                let path = key.with_base(&self.storage_root);
-                let metadata = file_metadata(&path).await?;
-                if metadata.is_dir() {
-                    continue;
-                }
-                objects.push(ListingObject {
-                    key: key.clone(),
-                    last_modified: metadata.modified()?,
-                    size: metadata.len(),
-                });
-            }
-            let objects = objects;
+            let objects = keys
+                .into_iter()
+                .filter_map(|k| {
+                    let path = k.with_base(&self.storage_root);
+                    if path.is_dir() {
+                        None
+                    } else {
+                        Some(ListingObject {
+                            key: k.clone(),
+                            // LocalFs is just for testing, so just specify a dummy time
+                            last_modified: SystemTime::now(),
+                            size: 0,
+                        })
+                    }
+                })
+                .collect();
 
             if let ListingMode::NoDelimiter = mode {
                 result.keys = objects;
             } else {
                 let mut prefixes = HashSet::new();
                 for object in objects {
-                    let key = dbg!(object.key);
+                    let key = object.key;
                     // If the part after the prefix includes a "/", take only the first part and put it in `prefixes`.
                     let relative_key = if let Some(prefix) = prefix {
                         let mut prefix = prefix.clone();
                         // We only strip the dirname of the prefix, so that when we strip it from the start of keys we
                         // end up with full file/dir names.
                         let prefix_full_local_path = prefix.with_base(&self.storage_root);
-                        dbg!(&prefix_full_local_path);
                         let has_slash = prefix.0.to_string().ends_with('/');
                         let strip_prefix = if prefix_full_local_path.is_dir() && has_slash {
                             prefix
@@ -412,7 +393,6 @@ impl RemoteStorage for LocalFs {
                             prefix.0.pop();
                             prefix
                         };
-                        dbg!(&strip_prefix);
 
                         RemotePath::new(key.strip_prefix(&strip_prefix).unwrap()).unwrap()
                     } else {
@@ -420,25 +400,19 @@ impl RemoteStorage for LocalFs {
                     };
 
                     let relative_key = format!("{}", relative_key);
-                    dbg!(&relative_key);
                     if relative_key.contains(REMOTE_STORAGE_PREFIX_SEPARATOR) {
                         let first_part = relative_key
                             .split(REMOTE_STORAGE_PREFIX_SEPARATOR)
                             .next()
                             .unwrap()
                             .to_owned();
-                        dbg!(&first_part);
-                        dbg!(&prefix);
-                        prefixes.insert(if let Some(prefix) = prefix {
-                            prefix.join(first_part).to_string()
-                        } else {
-                            first_part
-                        });
+                        prefixes.insert(first_part);
                     } else {
                         result.keys.push(ListingObject {
                             key: RemotePath::from_string(&relative_key).unwrap(),
-                            last_modified: object.last_modified,
-                            size: object.size,
+                            // LocalFs is just for testing
+                            last_modified: SystemTime::now(),
+                            size: 0,
                         });
                     }
                 }
@@ -526,9 +500,6 @@ impl RemoteStorage for LocalFs {
         let target_path = from.with_base(&self.storage_root);
 
         let file_metadata = file_metadata(&target_path).await?;
-        if let Some(etag) = &opts.etag {
-
-        }
         let etag = mock_etag(&file_metadata);
 
         if opts.etag.as_ref() == Some(&etag) {
@@ -740,10 +711,9 @@ mod fs_tests {
         Ok(())
     }
 
-    pub(crate) fn create_storage() -> anyhow::Result<(LocalFs, CancellationToken)> {
+    fn create_storage() -> anyhow::Result<(LocalFs, CancellationToken)> {
         let storage_root = tempdir()?.path().to_path_buf();
-        LocalFs::new(storage_root, Duration::from_secs(120), None)
-            .map(|s| (s, CancellationToken::new()))
+        LocalFs::new(storage_root, Duration::from_secs(120)).map(|s| (s, CancellationToken::new()))
     }
 
     #[tokio::test]
@@ -991,7 +961,7 @@ mod fs_tests {
         );
         assert_eq!(
             listing.prefixes,
-            [RemotePath::from_string("timelines/some_timeline/grandparent/parent").unwrap()].to_vec()
+            [RemotePath::from_string("parent").unwrap()].to_vec()
         );
 
         // Delimiter and prefix without a trailing slash
@@ -1006,7 +976,7 @@ mod fs_tests {
         assert_eq!(listing.keys, vec![]);
         assert_eq!(
             listing.prefixes,
-            [RemotePath::from_string("timelines/some_timeline/grandparent").unwrap()].to_vec()
+            [RemotePath::from_string("grandparent").unwrap()].to_vec()
         );
 
         // Delimiter and prefix that's partway through a path component
@@ -1021,7 +991,7 @@ mod fs_tests {
         assert_eq!(listing.keys, vec![]);
         assert_eq!(
             listing.prefixes,
-            [RemotePath::from_string("timelines/some_timeline/grandparent").unwrap()].to_vec()
+            [RemotePath::from_string("grandparent").unwrap()].to_vec()
         );
 
         Ok(())

libs/remote_storage/tests/common/mod.rs

@@ -8,7 +8,7 @@ use bytes::Bytes;
 use camino::Utf8Path;
 use futures::stream::Stream;
 use once_cell::sync::OnceCell;
-use remote_storage::{Download, GenericRemoteStorage, ListingObject, RemotePath};
+use remote_storage::{Download, GenericRemoteStorage, RemotePath};
 use tokio::task::JoinSet;
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, info};
@@ -52,26 +52,11 @@ pub(crate) async fn download_to_vec(dl: Download) -> anyhow::Result<Vec<u8>> {
     Ok(buf)
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-pub(crate) struct RemoteBlobInfo {
-    pub(crate) path: RemotePath,
-    pub(crate) size: u64,
-}
-
-impl From<ListingObject> for RemoteBlobInfo {
-    fn from(listing: ListingObject) -> Self {
-        RemoteBlobInfo {
-            path: listing.key,
-            size: listing.size,
-        }
-    }
-}
-
 // Uploads files `folder{j}/blob{i}.txt`. See test description for more details.
 pub(crate) async fn upload_simple_remote_data(
     client: &Arc<GenericRemoteStorage>,
     upload_tasks_count: usize,
-) -> ControlFlow<HashSet<RemoteBlobInfo>, HashSet<RemoteBlobInfo>> {
+) -> ControlFlow<HashSet<RemotePath>, HashSet<RemotePath>> {
     info!("Creating {upload_tasks_count} remote files");
     let mut upload_tasks = JoinSet::new();
     let cancel = CancellationToken::new();
@@ -88,16 +73,12 @@ pub(crate) async fn upload_simple_remote_data(
             .with_context(|| format!("{blob_path:?} to RemotePath conversion"))?;
             debug!("Creating remote item {i} at path {blob_path:?}");
 
-            let data = format!("remote blob data {i}").into_bytes();
-            let (data, len) = upload_stream(data.into());
+            let (data, len) = upload_stream(format!("remote blob data {i}").into_bytes().into());
             task_client
                 .upload(data, len, &blob_path, None, &cancel)
                 .await?;
 
-            Ok::<_, anyhow::Error>(RemoteBlobInfo {
-                path: blob_path,
-                size: len as u64,
-            })
+            Ok::<_, anyhow::Error>(blob_path)
         });
     }
 
@@ -108,8 +89,8 @@ pub(crate) async fn upload_simple_remote_data(
             .context("task join failed")
             .and_then(|task_result| task_result.context("upload task failed"))
         {
-            Ok(remote_blob_info) => {
-                uploaded_blobs.insert(remote_blob_info);
+            Ok(upload_path) => {
+                uploaded_blobs.insert(upload_path);
             }
             Err(e) => {
                 error!("Upload task failed: {e:?}");
@@ -127,7 +108,7 @@ pub(crate) async fn upload_simple_remote_data(
 
 pub(crate) async fn cleanup(
     client: &Arc<GenericRemoteStorage>,
-    objects_to_delete: HashSet<RemoteBlobInfo>,
+    objects_to_delete: HashSet<RemotePath>,
 ) {
     info!(
         "Removing {} objects from the remote storage during cleanup",
@@ -135,11 +116,7 @@ pub(crate) async fn cleanup(
     );
     let cancel = CancellationToken::new();
     let mut delete_tasks = JoinSet::new();
-    for RemoteBlobInfo {
-        path: object_to_delete,
-        ..
-    } in objects_to_delete
-    {
+    for object_to_delete in objects_to_delete {
         let task_client = Arc::clone(client);
         let cancel = cancel.clone();
         delete_tasks.spawn(async move {
@@ -163,7 +140,7 @@ pub(crate) async fn cleanup(
 }
 pub(crate) struct Uploads {
     pub(crate) prefixes: HashSet<RemotePath>,
-    pub(crate) blobs: HashSet<RemoteBlobInfo>,
+    pub(crate) blobs: HashSet<RemotePath>,
 }
 
 pub(crate) async fn upload_remote_data(
@@ -192,13 +169,7 @@ pub(crate) async fn upload_remote_data(
                 .upload(data, data_len, &blob_path, None, &cancel)
                 .await?;
 
-            anyhow::Ok((
-                blob_prefix,
-                RemoteBlobInfo {
-                    path: blob_path,
-                    size: data_len as u64,
-                },
-            ))
+            Ok::<_, anyhow::Error>((blob_prefix, blob_path))
         });
     }
 

libs/remote_storage/tests/common/tests.rs

@@ -9,7 +9,7 @@ use test_context::test_context;
 use tokio_util::sync::CancellationToken;
 use tracing::debug;
 
-use crate::common::{download_to_vec, upload_stream, wrap_stream, RemoteBlobInfo};
+use crate::common::{download_to_vec, upload_stream, wrap_stream};
 
 use super::{
     MaybeEnabledStorage, MaybeEnabledStorageWithSimpleTestBlobs, MaybeEnabledStorageWithTestBlobs,
@@ -85,7 +85,7 @@ async fn pagination_should_work(ctx: &mut MaybeEnabledStorageWithTestBlobs) -> a
         .collect::<HashSet<_>>();
     assert_eq!(
         remote_only_prefixes.len() + missing_uploaded_prefixes.len(), 0,
-        "remote storage nested prefixes list mismatches with the uploads.\n\nRemote only prefixes: {remote_only_prefixes:?}\n\nmissing uploaded prefixes: {missing_uploaded_prefixes:?}",
+        "remote storage nested prefixes list mismatches with the uploads. Remote only prefixes: {remote_only_prefixes:?}, missing uploaded prefixes: {missing_uploaded_prefixes:?}",
     );
 
     // list_streaming
@@ -102,7 +102,6 @@ async fn pagination_should_work(ctx: &mut MaybeEnabledStorageWithTestBlobs) -> a
     let mut segment_max_size = 0;
     while let Some(st) = nested_remote_prefixes_st.next().await {
         let st = st?;
-        dbg!(&st.prefixes);
         segment_max_size = segment_max_size.max(st.prefixes.len());
         nested_remote_prefixes_combined.extend(st.prefixes.into_iter());
         segments += 1;
@@ -157,7 +156,7 @@ async fn list_no_delimiter_works(
         .context("client list root files failure")?
         .keys
         .into_iter()
-        .map(RemoteBlobInfo::from)
+        .map(|o| o.key)
         .collect::<HashSet<_>>();
     assert_eq!(
         root_files,
@@ -184,14 +183,14 @@ async fn list_no_delimiter_works(
         .context("client list nested files failure")?
         .keys
         .into_iter()
-        .map(RemoteBlobInfo::from)
+        .map(|o| o.key)
         .collect::<HashSet<_>>();
     let trim_remote_blobs: HashSet<_> = ctx
         .remote_blobs
         .iter()
-        .filter(|x| x.path.get_path().starts_with("folder1"))
-        .cloned()
-        .map(RemoteBlobInfo::from)
+        .map(|x| x.get_path())
+        .filter(|x| x.starts_with("folder1"))
+        .map(|x| RemotePath::new(x).expect("must be valid path"))
         .collect();
     assert_eq!(
         nested_remote_files, trim_remote_blobs,
@@ -229,7 +228,7 @@ async fn list_partial_prefix(
         .await?
         .keys
         .into_iter()
-        .map(RemoteBlobInfo::from)
+        .map(|o| o.key)
         .collect();
     assert_eq!(&objects, &ctx.remote_blobs);
 
@@ -259,7 +258,7 @@ async fn list_partial_prefix(
         .await?
         .keys
         .into_iter()
-        .map(RemoteBlobInfo::from)
+        .map(|o| o.key)
         .collect();
     assert_eq!(&objects, &ctx.remote_blobs);
 
@@ -304,14 +303,13 @@ async fn list_partial_prefix(
         .await?
         .keys
         .into_iter()
-        .map(RemoteBlobInfo::from)
+        .map(|o| o.key)
         .collect();
     let expect: HashSet<_> = ctx
         .remote_blobs
         .iter()
-        .filter(|o| o.path.get_path().starts_with("folder2"))
+        .filter(|o| o.get_path().starts_with("folder2"))
         .cloned()
-        .map(RemoteBlobInfo::from)
         .collect();
     assert_eq!(&objects, &expect);
 
@@ -423,7 +421,7 @@ async fn delete_prefix(ctx: &mut MaybeEnabledStorageWithSimpleTestBlobs) -> anyh
                 .await?
                 .keys
                 .into_iter()
-                .map(RemoteBlobInfo::from)
+                .map(|o| o.key)
                 .collect();
             assert_eq!($expect, listing);
         }};
@@ -446,19 +444,13 @@ async fn delete_prefix(ctx: &mut MaybeEnabledStorageWithSimpleTestBlobs) -> anyh
 
     // Deleting a path which overlaps with an existing object should do nothing. We pick the first
     // path in the set as our common prefix.
-    let path = expect
-        .iter()
-        .next()
-        .expect("empty set")
-        .path
-        .clone()
-        .join("xyz");
+    let path = expect.iter().next().expect("empty set").clone().join("xyz");
     test_client.delete_prefix(&path, &cancel).await?;
     assert_list!(expect);
 
     // Deleting an exact path should work. We pick the first path in the set.
     let path = expect.iter().next().expect("empty set").clone();
-    test_client.delete_prefix(&path.path, &cancel).await?;
+    test_client.delete_prefix(&path, &cancel).await?;
     expect.remove(&path);
     assert_list!(expect);
 
@@ -466,7 +458,7 @@ async fn delete_prefix(ctx: &mut MaybeEnabledStorageWithSimpleTestBlobs) -> anyh
     test_client
         .delete_prefix(&RemotePath::from_string("folder0/blob_")?, &cancel)
         .await?;
-    expect.retain(|RemoteBlobInfo { path: p, .. }| !p.get_path().as_str().starts_with("folder0/"));
+    expect.retain(|p| !p.get_path().as_str().starts_with("folder0/"));
     assert_list!(expect);
 
     // Deleting a common prefix should delete all objects.

libs/remote_storage/tests/test_localfs.rs

@@ -1,167 +0,0 @@
-use std::num::NonZeroU32;
-use std::ops::ControlFlow;
-use std::sync::Arc;
-use std::{collections::HashSet, time::Duration};
-
-use anyhow::Context;
-use remote_storage::{
-    GenericRemoteStorage, LocalFs, RemotePath, RemoteStorageConfig, RemoteStorageKind,
-};
-use test_context::AsyncTestContext;
-use tokio_util::sync::CancellationToken;
-use tracing::info;
-
-mod common;
-
-#[path = "common/tests.rs"]
-mod tests_localfs;
-
-use common::{
-    cleanup, ensure_logging_ready, upload_remote_data, upload_simple_remote_data, RemoteBlobInfo,
-};
-
-const BASE_PREFIX: &str = "test";
-
-struct EnabledLocalFs {
-    client: Arc<GenericRemoteStorage>,
-    base_prefix: &'static str,
-}
-
-impl EnabledLocalFs {
-    async fn setup(max_keys_per_list_response: Option<u32>) -> Self {
-        let storage_root = camino_tempfile::tempdir()
-            .expect("create tempdir")
-            .path()
-            .to_path_buf();
-
-        let remote_storage_config = RemoteStorageConfig {
-            storage: RemoteStorageKind::LocalFs {
-                local_path: storage_root,
-                max_keys_per_list_response: max_keys_per_list_response.and_then(NonZeroU32::new),
-            },
-            timeout: Duration::from_secs(120),
-        };
-        let client = Arc::new(
-            GenericRemoteStorage::from_config(&remote_storage_config)
-                .await
-                .context("remote storage init")
-                .unwrap(),
-        );
-
-        EnabledLocalFs {
-            client,
-            base_prefix: &BASE_PREFIX,
-        }
-    }
-}
-
-enum MaybeEnabledStorage {
-    Enabled(EnabledLocalFs),
-    Disabled,
-}
-
-impl AsyncTestContext for MaybeEnabledStorage {
-    async fn setup() -> Self {
-        ensure_logging_ready();
-        Self::Enabled(EnabledLocalFs::setup(None).await)
-    }
-}
-
-enum MaybeEnabledStorageWithTestBlobs {
-    Enabled(LocalFsWithTestBlobs),
-    Disabled,
-    UploadsFailed(anyhow::Error, LocalFsWithTestBlobs),
-}
-
-struct LocalFsWithTestBlobs {
-    enabled: EnabledLocalFs,
-    remote_prefixes: HashSet<RemotePath>,
-    remote_blobs: HashSet<RemoteBlobInfo>,
-}
-
-impl AsyncTestContext for MaybeEnabledStorageWithTestBlobs {
-    async fn setup() -> Self {
-        ensure_logging_ready();
-
-        let max_keys_in_list_response = 10;
-        let upload_tasks_count = 1 + (2 * usize::try_from(max_keys_in_list_response).unwrap());
-
-        let enabled = EnabledLocalFs::setup(Some(max_keys_in_list_response)).await;
-
-        match upload_remote_data(&enabled.client, enabled.base_prefix, upload_tasks_count).await {
-            ControlFlow::Continue(uploads) => {
-                info!("Remote objects created successfully");
-
-                Self::Enabled(LocalFsWithTestBlobs {
-                    enabled,
-                    remote_prefixes: uploads.prefixes,
-                    remote_blobs: uploads.blobs,
-                })
-            }
-            ControlFlow::Break(uploads) => Self::UploadsFailed(
-                anyhow::anyhow!("One or multiple blobs failed to upload to Azure"),
-                LocalFsWithTestBlobs {
-                    enabled,
-                    remote_prefixes: uploads.prefixes,
-                    remote_blobs: uploads.blobs,
-                },
-            ),
-        }
-    }
-
-    async fn teardown(self) {
-        match self {
-            Self::Disabled => {}
-            Self::Enabled(ctx) | Self::UploadsFailed(_, ctx) => {
-                cleanup(&ctx.enabled.client, ctx.remote_blobs).await;
-            }
-        }
-    }
-}
-
-enum MaybeEnabledStorageWithSimpleTestBlobs {
-    Enabled(LocalFsWithSimpleTestBlobs),
-    Disabled,
-    UploadsFailed(anyhow::Error, LocalFsWithSimpleTestBlobs),
-}
-struct LocalFsWithSimpleTestBlobs {
-    enabled: EnabledLocalFs,
-    remote_blobs: HashSet<RemoteBlobInfo>,
-}
-
-impl AsyncTestContext for MaybeEnabledStorageWithSimpleTestBlobs {
-    async fn setup() -> Self {
-        ensure_logging_ready();
-        let max_keys_in_list_response = 10;
-        let upload_tasks_count = 1 + (2 * usize::try_from(max_keys_in_list_response).unwrap());
-
-        let enabled = EnabledLocalFs::setup(Some(max_keys_in_list_response)).await;
-
-        match upload_simple_remote_data(&enabled.client, upload_tasks_count).await {
-            ControlFlow::Continue(uploads) => {
-                info!("Remote objects created successfully");
-
-                Self::Enabled(LocalFsWithSimpleTestBlobs {
-                    enabled,
-                    remote_blobs: uploads,
-                })
-            }
-            ControlFlow::Break(uploads) => Self::UploadsFailed(
-                anyhow::anyhow!("One or multiple blobs failed to upload to Azure"),
-                LocalFsWithSimpleTestBlobs {
-                    enabled,
-                    remote_blobs: uploads,
-                },
-            ),
-        }
-    }
-
-    async fn teardown(self) {
-        match self {
-            Self::Disabled => {}
-            Self::Enabled(ctx) | Self::UploadsFailed(_, ctx) => {
-                cleanup(&ctx.enabled.client, ctx.remote_blobs).await;
-            }
-        }
-    }
-}

libs/remote_storage/tests/test_real_azure.rs

@@ -17,9 +17,7 @@ mod common;
 #[path = "common/tests.rs"]
 mod tests_azure;
 
-use common::{
-    cleanup, ensure_logging_ready, upload_remote_data, upload_simple_remote_data, RemoteBlobInfo,
-};
+use common::{cleanup, ensure_logging_ready, upload_remote_data, upload_simple_remote_data};
 
 const ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME: &str = "ENABLE_REAL_AZURE_REMOTE_STORAGE";
 
@@ -85,7 +83,7 @@ enum MaybeEnabledStorageWithTestBlobs {
 struct AzureWithTestBlobs {
     enabled: EnabledAzure,
     remote_prefixes: HashSet<RemotePath>,
-    remote_blobs: HashSet<RemoteBlobInfo>,
+    remote_blobs: HashSet<RemotePath>,
 }
 
 impl AsyncTestContext for MaybeEnabledStorageWithTestBlobs {
@@ -142,7 +140,7 @@ enum MaybeEnabledStorageWithSimpleTestBlobs {
 }
 struct AzureWithSimpleTestBlobs {
     enabled: EnabledAzure,
-    remote_blobs: HashSet<RemoteBlobInfo>,
+    remote_blobs: HashSet<RemotePath>,
 }
 
 impl AsyncTestContext for MaybeEnabledStorageWithSimpleTestBlobs {

libs/remote_storage/tests/test_real_s3.rs

@@ -26,9 +26,7 @@ mod common;
 #[path = "common/tests.rs"]
 mod tests_s3;
 
-use common::{
-    cleanup, ensure_logging_ready, upload_remote_data, upload_simple_remote_data, RemoteBlobInfo,
-};
+use common::{cleanup, ensure_logging_ready, upload_remote_data, upload_simple_remote_data};
 use utils::backoff;
 
 const ENABLE_REAL_S3_REMOTE_STORAGE_ENV_VAR_NAME: &str = "ENABLE_REAL_S3_REMOTE_STORAGE";
@@ -262,7 +260,7 @@ enum MaybeEnabledStorageWithTestBlobs {
 struct S3WithTestBlobs {
     enabled: EnabledS3,
     remote_prefixes: HashSet<RemotePath>,
-    remote_blobs: HashSet<RemoteBlobInfo>,
+    remote_blobs: HashSet<RemotePath>,
 }
 
 impl AsyncTestContext for MaybeEnabledStorageWithTestBlobs {
@@ -319,7 +317,7 @@ enum MaybeEnabledStorageWithSimpleTestBlobs {
 }
 struct S3WithSimpleTestBlobs {
     enabled: EnabledS3,
-    remote_blobs: HashSet<RemoteBlobInfo>,
+    remote_blobs: HashSet<RemotePath>,
 }
 
 impl AsyncTestContext for MaybeEnabledStorageWithSimpleTestBlobs {

pageserver/ctl/src/key.rs

@@ -345,6 +345,7 @@ impl AuxFileV2 {
                 AuxFileV2::Recognized("pg_logical/replorigin_checkpoint", hash)
             }
             (2, 1) => AuxFileV2::Recognized("pg_replslot/", hash),
+            (4, 1) => AuxFileV2::Recognized("lfc.state", hash),
             (1, 0xff) => AuxFileV2::OtherWithPrefix("pg_logical/", hash),
             (0xff, 0xff) => AuxFileV2::Other(hash),
             _ => return None,

pageserver/src/aux_file.rs

@@ -39,6 +39,7 @@ fn aux_hash_to_metadata_key(dir_level1: u8, dir_level2: u8, data: &[u8]) -> Key
 
 const AUX_DIR_PG_LOGICAL: u8 = 0x01;
 const AUX_DIR_PG_REPLSLOT: u8 = 0x02;
+const AUX_DIR_LFC_STATE: u8 = 0x04;
 const AUX_DIR_PG_UNKNOWN: u8 = 0xFF;
 
 /// Encode the aux file into a fixed-size key.
@@ -75,6 +76,8 @@ pub fn encode_aux_file_key(path: &str) -> Key {
         aux_hash_to_metadata_key(AUX_DIR_PG_LOGICAL, 0xFF, fname.as_bytes())
     } else if let Some(fname) = path.strip_prefix("pg_replslot/") {
         aux_hash_to_metadata_key(AUX_DIR_PG_REPLSLOT, 0x01, fname.as_bytes())
+    } else if let Some(fname) = path.strip_prefix("lfc.state") {
+        aux_hash_to_metadata_key(AUX_DIR_LFC_STATE, 0x01, fname.as_bytes())
     } else {
         if cfg!(debug_assertions) {
             warn!(

pageserver/src/deletion_queue.rs

@@ -841,7 +841,6 @@ mod test {
         let storage_config = RemoteStorageConfig {
             storage: RemoteStorageKind::LocalFs {
                 local_path: remote_fs_dir.clone(),
-                max_keys_per_list_response: None,
             },
             timeout: RemoteStorageConfig::DEFAULT_TIMEOUT,
         };

pageserver/src/tenant.rs

@@ -4642,7 +4642,6 @@ pub(crate) mod harness {
             let config = RemoteStorageConfig {
                 storage: RemoteStorageKind::LocalFs {
                     local_path: remote_fs_dir.clone(),
-                    max_keys_per_list_response: None,
                 },
                 timeout: RemoteStorageConfig::DEFAULT_TIMEOUT,
             };

pgxn/neon/Makefile

@@ -32,6 +32,8 @@ DATA = \
 	neon--1.2--1.3.sql \
 	neon--1.3--1.4.sql \
 	neon--1.4--1.5.sql \
+	neon--1.5--1.6.sql \
+	neon--1.6--1.5.sql \
 	neon--1.5--1.4.sql \
 	neon--1.4--1.3.sql \
 	neon--1.3--1.2.sql \

pgxn/neon/file_cache.c

@@ -22,6 +22,7 @@
 #include "neon_pgversioncompat.h"
 
 #include "access/parallel.h"
+#include "access/xlog.h"
 #include "funcapi.h"
 #include "miscadmin.h"
 #include "pagestore_client.h"
@@ -30,22 +31,28 @@
 #include "port/pg_iovec.h"
 #include "postmaster/bgworker.h"
 #include RELFILEINFO_HDR
+#include "replication/message.h"
 #include "storage/buf_internals.h"
 #include "storage/fd.h"
 #include "storage/ipc.h"
 #include "storage/latch.h"
 #include "storage/lwlock.h"
 #include "storage/pg_shmem.h"
+#include "tcop/tcopprot.h"
 #include "utils/builtins.h"
 #include "utils/dynahash.h"
 #include "utils/guc.h"
 
+#if PG_VERSION_NUM >= 150000
+#include "access/xlogrecovery.h"
+#endif
+
 #include "hll.h"
 #include "bitmap.h"
 #include "neon.h"
 #include "neon_perf_counters.h"
 
-#define CriticalAssert(cond) do if (!(cond)) elog(PANIC, "Assertion %s failed at %s:%d: ", #cond, __FILE__, __LINE__); while (0)
+#define CriticalAssert(cond) do if (!(cond)) elog(PANIC, "LFC: assertion %s failed at %s:%d: ", #cond, __FILE__, __LINE__); while (0)
 
 /*
  * Local file cache is used to temporary store relations pages in local file system.
@@ -100,7 +107,9 @@ typedef struct FileCacheEntry
 	BufferTag	key;
 	uint32		hash;
 	uint32		offset;
-	uint32		access_count;
+	uint32		access_count : 30;
+	uint32      prewarm_requested : 1; /* entry should be filled by prewarm */
+	uint32      prewarm_started : 1;   /* chunk is written by lfc_prewarm */
 	uint32		bitmap[CHUNK_BITMAP_SIZE];
 	dlist_node	list_node;		/* LRU/holes list node */
 } FileCacheEntry;
@@ -118,26 +127,57 @@ typedef struct FileCacheControl
 	uint64		writes;			/* number of writes issued */
 	uint64		time_read;		/* time spent reading (us) */
 	uint64		time_write;		/* time spent writing (us) */
+	uint32		prewarm_total_chunks;
+	uint32		prewarm_curr_chunk;
+	uint32		prewarmed_pages;
+	uint32		skipped_pages;
 	dlist_head	lru;			/* double linked list for LRU replacement
 								 * algorithm */
 	dlist_head  holes;          /* double linked list of punched holes */
 	HyperLogLogState wss_estimation; /* estimation of working set size */
 } FileCacheControl;
 
+typedef struct FileCacheStateEntry
+{
+	BufferTag	key;
+	uint32		bitmap[CHUNK_BITMAP_SIZE];
+} FileCacheStateEntry;
+
 static HTAB *lfc_hash;
 static int	lfc_desc = 0;
 static LWLockId lfc_lock;
 static int	lfc_max_size;
 static int	lfc_size_limit;
+static int	lfc_prewarm_limit;
+static int	lfc_prewarm_batch;
 static char *lfc_path;
 static FileCacheControl *lfc_ctl;
 static shmem_startup_hook_type prev_shmem_startup_hook;
 #if PG_VERSION_NUM>=150000
 static shmem_request_hook_type prev_shmem_request_hook;
 #endif
+static CustomCheckpointHookType PrevCheckpointHook;
+
 
 #define LFC_ENABLED() (lfc_ctl->limit != 0)
 
+PGDLLEXPORT void LfcPrewarmMain(Datum main_arg);
+
+static void
+LfcCheckpointHook(int flags)
+{
+	if (flags & CHECKPOINT_IS_SHUTDOWN)
+	{
+		lfc_save_state();
+	}
+
+	if (PrevCheckpointHook)
+	{
+		PrevCheckpointHook(flags);
+	}
+}
+
+
 /*
  * Local file cache is optional and Neon can work without it.
  * In case of any any errors with this cache, we should disable it but to not throw error.
@@ -149,7 +189,7 @@ lfc_disable(char const *op)
 {
 	int			fd;
 
-	elog(WARNING, "Failed to %s local file cache at %s: %m, disabling local file cache", op, lfc_path);
+	elog(WARNING, "LFC: failed to %s local file cache at %s: %m, disabling local file cache", op, lfc_path);
 
 	/* Invalidate hash */
 	LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
@@ -184,7 +224,7 @@ lfc_disable(char const *op)
 			pgstat_report_wait_end();
 
 			if (rc < 0)
-				elog(WARNING, "Failed to truncate local file cache %s: %m", lfc_path);
+				elog(WARNING, "LFC: failed to truncate local file cache %s: %m", lfc_path);
 		}
 	}
 
@@ -196,7 +236,7 @@ lfc_disable(char const *op)
 
 	fd = BasicOpenFile(lfc_path, O_RDWR | O_CREAT | O_TRUNC);
 	if (fd < 0)
-		elog(WARNING, "Failed to recreate local file cache %s: %m", lfc_path);
+		elog(WARNING, "LFC: failed to recreate local file cache %s: %m", lfc_path);
 	else
 		close(fd);
 
@@ -236,6 +276,17 @@ lfc_ensure_opened(void)
 	return enabled;
 }
 
+PGDLLEXPORT void
+LfcPrewarmMain(Datum main_arg)
+{
+	pqsignal(SIGTERM, die);
+
+	BackgroundWorkerUnblockSignals();
+
+	lfc_load_pages();
+}
+
+
 static void
 lfc_shmem_startup(void)
 {
@@ -267,14 +318,7 @@ lfc_shmem_startup(void)
 								 n_chunks + 1, n_chunks + 1,
 								 &info,
 								 HASH_ELEM | HASH_BLOBS);
-		lfc_ctl->generation = 0;
-		lfc_ctl->size = 0;
-		lfc_ctl->used = 0;
-		lfc_ctl->hits = 0;
-		lfc_ctl->misses = 0;
-		lfc_ctl->writes = 0;
-		lfc_ctl->time_read = 0;
-		lfc_ctl->time_write = 0;
+		memset(lfc_ctl, 0, sizeof *lfc_ctl);
 		dlist_init(&lfc_ctl->lru);
 		dlist_init(&lfc_ctl->holes);
 
@@ -285,7 +329,7 @@ lfc_shmem_startup(void)
 		fd = BasicOpenFile(lfc_path, O_RDWR | O_CREAT | O_TRUNC);
 		if (fd < 0)
 		{
-			elog(WARNING, "Failed to create local file cache %s: %m", lfc_path);
+			elog(WARNING, "LFC: failed to create local file cache %s: %m", lfc_path);
 			lfc_ctl->limit = 0;
 		}
 		else
@@ -295,6 +339,9 @@ lfc_shmem_startup(void)
 		}
 	}
 	LWLockRelease(AddinShmemInitLock);
+
+	PrevCheckpointHook = CustomCheckpointHook;
+	CustomCheckpointHook = LfcCheckpointHook;
 }
 
 static void
@@ -327,7 +374,7 @@ lfc_check_limit_hook(int *newval, void **extra, GucSource source)
 {
 	if (*newval > lfc_max_size)
 	{
-		elog(ERROR, "neon.file_cache_size_limit can not be larger than neon.max_file_cache_size");
+		elog(ERROR, "LFC: neon.file_cache_size_limit can not be larger than neon.max_file_cache_size");
 		return false;
 	}
 	return true;
@@ -436,6 +483,32 @@ lfc_init(void)
 							   NULL,
 							   NULL);
 
+	DefineCustomIntVariable("neon.file_cache_prewarm_limit",
+							"Maximal number of prewarmed pages",
+							NULL,
+							&lfc_prewarm_limit,
+							0,	/* disabled by default */
+							0,
+							INT_MAX,
+							PGC_SIGHUP,
+							0,
+							NULL,
+							NULL,
+							NULL);
+
+	DefineCustomIntVariable("neon.file_cache_prewarm_batch",
+							"Number of pages retrivied by prewarm from page server",
+							NULL,
+							&lfc_prewarm_batch,
+							64,
+							1,
+							INT_MAX,
+							PGC_SIGHUP,
+							0,
+							NULL,
+							NULL,
+							NULL);
+
 	if (lfc_max_size == 0)
 		return;
 
@@ -447,8 +520,326 @@ lfc_init(void)
 #else
 	lfc_shmem_request();
 #endif
+
+	if (lfc_prewarm_limit != 0)
+	{
+		BackgroundWorker bgw;
+		memset(&bgw, 0, sizeof(bgw));
+		bgw.bgw_flags = BGWORKER_SHMEM_ACCESS;
+
+		bgw.bgw_start_time = BgWorkerStart_ConsistentState;
+		snprintf(bgw.bgw_library_name, BGW_MAXLEN, "neon");
+		snprintf(bgw.bgw_function_name, BGW_MAXLEN, "LfcPrewarmMain");
+		snprintf(bgw.bgw_name, BGW_MAXLEN, "LFC prewarm");
+		snprintf(bgw.bgw_type, BGW_MAXLEN, "LFC prewarm");
+
+		RegisterBackgroundWorker(&bgw);
+	}
 }
 
+static FileCacheStateEntry*
+lfc_get_state(size_t* n_entries)
+{
+	size_t max_entries = *n_entries;
+	size_t i = 0;
+	FileCacheStateEntry* fs = (FileCacheStateEntry*)palloc(sizeof(FileCacheStateEntry) * max_entries);
+
+	LWLockAcquire(lfc_lock, LW_SHARED);
+
+	if (LFC_ENABLED())
+	{
+		dlist_iter	iter;
+		dlist_reverse_foreach(iter, &lfc_ctl->lru)
+		{
+			FileCacheEntry *entry = dlist_container(FileCacheEntry, list_node, iter.cur);
+			memcpy(&fs[i].key, &entry->key, sizeof entry->key);
+			memcpy(fs[i].bitmap, entry->bitmap, sizeof entry->bitmap);
+			if (++i == max_entries)
+				break;
+		}
+		elog(LOG, "LFC: save state of %ld chunks", (long)i);
+	}
+
+	LWLockRelease(lfc_lock);
+
+	*n_entries = i;
+	return fs;
+}
+
+/*
+ * Save state of local file cache as AUX file. Size of saved state is limited by lfc_prewarm_limit.
+ * This function saves first mostrecently used pages.
+ * It is expected to be called at shutdown checkpoint by checkpointer.
+ */
+void
+lfc_save_state(void)
+{
+	size_t n_entries = lfc_prewarm_limit;
+	FileCacheStateEntry* fs;
+
+	if (n_entries == 0)
+		return;
+
+	fs = lfc_get_state(&n_entries);
+	if (n_entries != 0)
+	{
+#if PG_MAJORVERSION_NUM < 17
+		XLogFlush(LogLogicalMessage("neon-file:lfc.state", (char const*)fs, sizeof(FileCacheStateEntry) * n_entries, false));
+#else
+		LogLogicalMessage("neon-file:lfc.state", (char const*)fs, sizeof(FileCacheStateEntry) * n_entries, false, true);
+#endif
+	}
+	pfree(fs);
+}
+
+/*
+ * Prewarm LFC cache to the specified state.
+ *
+ * Prewarming can interfere with accesses to the pages by other backends. Usually access to LFC is protected by shared buffers: when Postgres
+ * is reading page, it pins shared buffer and enforces that only one backend is reading it, while other are waiting for read completion.
+ *
+ * But it is not true for prewarming: backend can fetch page itself, modify and then write it to LFC. At the
+ * same time `lfc_prewarm` tries to write deteriorated image of this page in LFC. To increase concurrency, access to LFC files (both read and write)
+ * is performed without holding locks. So it can happen that two or more processes write different content to the same location in the LFC file.
+ * Certainly we can not rely on disk content in this case.
+ *
+ * To solve this problem we use two flags in LFC entry: `prewarm_requested` and `prewarm_started`. First is set before prewarm is actually started.
+ * `lfc_prewarm` writes to LFC file only if this flag is set. This flag is cleared if any other backend performs write to this LFC chunk.
+ * In this case data loaded by `lfc_prewarm` is considered to be deteriorated and should be just ignored.
+ *
+ * But as far as write to LFC is performed without holding lock, there is no guarantee that no such write is in progress.
+ * This is why second flag is used: `prewarm_started`. It is set by `lfc_prewarm` when is starts writing page and cleared when write is completed.
+ * Any other backend writing to LFC should abandon it's write to LFC file (just not mark page as loaded in bitmap) if this flag is set.
+ * So neither `lfc_prewarm`, neither backend are saving page in LFC in this case - it is just skipped.
+ */
+
+static void
+lfc_prewarm(FileCacheStateEntry* fs, size_t n_entries)
+{
+	ssize_t rc;
+	size_t snd_idx = 0, rcv_idx = 0;
+	size_t n_sent = 0, n_received = 0;
+	FileCacheEntry *entry;
+	uint64 generation;
+	uint32 entry_offset;
+	uint32 hash;
+	size_t i;
+	bool   found;
+	int    shard_no;
+
+	if (!lfc_ensure_opened())
+		return;
+
+	if (n_entries == 0 || fs == NULL)
+	{
+		elog(LOG, "LFC: prewarm is disabled");
+		return;
+	}
+
+	LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
+
+	/* Do not prewarm more entries than LFC limit */
+	if (lfc_ctl->limit <= lfc_ctl->size)
+	{
+		LWLockRelease(lfc_lock);
+		return;
+	}
+	if (n_entries > lfc_ctl->limit - lfc_ctl->size)
+	{
+		n_entries = lfc_ctl->limit - lfc_ctl->size;
+	}
+
+	/* Initialize fields used to track prewarming progress */
+	lfc_ctl->prewarm_total_chunks = n_entries;
+	lfc_ctl->prewarm_curr_chunk = 0;
+
+    /*
+	 * Load LFC state and add entries in hash table.
+	 * It is needed to track modification of prewarmed pages.
+	 * All such entries have `prewarm_requested` flag set. When entry is updated (some backed reads or writes
+	 * some pages from this chunk), then `prewarm_requested` flag is cleared, prohibiting prewarm of this chunk.
+	 * It prevents overwritting page updated or loaded by backend with older one, loaded by prewarm.
+	 */
+	for (i = 0; i < n_entries; i++)
+	{
+		hash = get_hash_value(lfc_hash, &fs[i].key);
+		entry = hash_search_with_hash_value(lfc_hash, &fs[i].key, hash, HASH_ENTER, &found);
+		/* Do not prewarm chunks which are already present in LFC */
+		if (!found)
+		{
+			entry->offset = lfc_ctl->size++;
+			entry->hash = hash;
+			entry->access_count = 0;
+			entry->prewarm_requested = true;
+			entry->prewarm_started = false;
+			memset(entry->bitmap, 0, sizeof entry->bitmap);
+			/* Most recently visted pages are stored first */
+			dlist_push_head(&lfc_ctl->lru, &entry->list_node);
+			lfc_ctl->used += 1;
+		}
+	}
+	LWLockRelease(lfc_lock);
+
+	elog(LOG, "LFC: start loading %ld chunks", (long)n_entries);
+
+	while (true)
+	{
+		size_t chunk_no = snd_idx / BLOCKS_PER_CHUNK;
+		size_t offs_in_chunk = snd_idx % BLOCKS_PER_CHUNK;
+		if (chunk_no < n_entries)
+		{
+			if (fs[chunk_no].bitmap[offs_in_chunk >> 5] & (1 << (offs_in_chunk & 31)))
+			{
+				/*
+				 * In case of prewarming replica we should be careful not to load too new version
+				 * of the page - with LSN larger than current replay LSN.
+				 * At primary we are always loading latest version.
+				 */
+				XLogRecPtr req_lsn = RecoveryInProgress() ? GetXLogReplayRecPtr(NULL) : UINT64_MAX;
+
+				NeonGetPageRequest request = {
+					.req.tag = T_NeonGetPageRequest,
+					/* lsn and not_modified_since are filled in below */
+					.rinfo = BufTagGetNRelFileInfo(fs[chunk_no].key),
+					.forknum = fs[chunk_no].key.forkNum,
+					.blkno = fs[chunk_no].key.blockNum + offs_in_chunk,
+					.req.lsn = req_lsn,
+					.req.not_modified_since = 0
+				};
+				shard_no = get_shard_number(&fs[chunk_no].key);
+				while (!page_server->send(shard_no, (NeonRequest *) &request)
+					   || !page_server->flush(shard_no))
+				{
+					/* do nothing */
+				}
+				n_sent += 1;
+			}
+			snd_idx += 1;
+		}
+		if (n_sent >= n_received + lfc_prewarm_batch || chunk_no == n_entries)
+		{
+			NeonResponse * resp;
+			do
+			{
+				chunk_no = rcv_idx / BLOCKS_PER_CHUNK;
+				offs_in_chunk = rcv_idx % BLOCKS_PER_CHUNK;
+				rcv_idx += 1;
+			} while (!(fs[chunk_no].bitmap[offs_in_chunk >> 5] & (1 << (offs_in_chunk & 31))));
+
+			shard_no = get_shard_number(&fs[chunk_no].key);
+			resp = page_server->receive(shard_no);
+			lfc_ctl->prewarm_curr_chunk = chunk_no;
+
+			if (resp->tag != T_NeonGetPageResponse)
+			{
+				elog(LOG, "LFC: unexpected response type: %d", resp->tag);
+				return;
+			}
+
+			hash = get_hash_value(lfc_hash, &fs[chunk_no].key);
+
+			LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
+			entry = hash_search_with_hash_value(lfc_hash, &fs[chunk_no].key, hash, HASH_FIND, NULL);
+			if (entry != NULL && entry->prewarm_requested)
+			{
+				/* Unlink entry from LRU list to pin it for the duration of IO operation */
+				if (entry->access_count++ == 0)
+					dlist_delete(&entry->list_node);
+
+				generation = lfc_ctl->generation;
+				entry_offset = entry->offset;
+				Assert(!entry->prewarm_started);
+				entry->prewarm_started = true;
+
+				LWLockRelease(lfc_lock);
+
+				rc = pwrite(lfc_desc, ((NeonGetPageResponse*)resp)->page, BLCKSZ, ((off_t) entry_offset * BLOCKS_PER_CHUNK + offs_in_chunk) * BLCKSZ);
+				if (rc != BLCKSZ)
+				{
+					lfc_disable("write");
+					break;
+				}
+				else
+				{
+					LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
+
+					if (lfc_ctl->generation == generation)
+					{
+						CriticalAssert(LFC_ENABLED());
+						if (--entry->access_count == 0)
+							dlist_push_tail(&lfc_ctl->lru, &entry->list_node);
+						if (entry->prewarm_requested)
+						{
+							lfc_ctl->used_pages += 1 - ((entry->bitmap[offs_in_chunk >> 5] >> (offs_in_chunk & 31)) & 1);
+							entry->bitmap[offs_in_chunk >> 5] |= 1 << (offs_in_chunk & 31);
+							lfc_ctl->prewarmed_pages += 1;
+						}
+						else
+						{
+							lfc_ctl->skipped_pages += 1;
+						}
+						Assert(entry->prewarm_started);
+						entry->prewarm_started = false;
+					}
+
+					LWLockRelease(lfc_lock);
+				}
+			}
+			else
+			{
+				Assert(!entry || !entry->prewarm_started);
+				lfc_ctl->skipped_pages += 1;
+				LWLockRelease(lfc_lock);
+			}
+
+			if (++n_received == n_sent && snd_idx >= n_entries * BLOCKS_PER_CHUNK)
+			{
+				break;
+			}
+		}
+	}
+	Assert(n_sent == n_received);
+	lfc_ctl->prewarm_curr_chunk = n_entries;
+	elog(LOG, "LFC: complete prewarming: loaded %ld pages", (long)n_received);
+}
+
+
+/*
+ * Load pages from LFC state saved in AUX file.
+ */
+void
+lfc_load_pages(void)
+{
+	int fd;
+	FileCacheStateEntry *fs;
+	ssize_t rc;
+	size_t max_entries = lfc_prewarm_limit;
+
+	fd = OpenTransientFile("lfc.state", O_RDONLY | PG_BINARY);
+	if (fd < 0)
+	{
+		elog(LOG, "LFC: state file is missing");
+		return;
+	}
+
+	fs = (FileCacheStateEntry*)palloc(sizeof(FileCacheStateEntry) * max_entries);
+	rc = read(fd, fs, sizeof(FileCacheStateEntry) * max_entries);
+	if (rc <= 0)
+	{
+		elog(LOG, "LFC: Failed to read state file: %m");
+		CloseTransientFile(fd);
+	}
+	else
+	{
+		CloseTransientFile(fd);
+		elog(LOG, "LFC: read state with %lu entries", (long)(rc / sizeof(FileCacheStateEntry)));
+
+		lfc_prewarm(fs, rc / sizeof(FileCacheStateEntry));
+	}
+	pfree(fs);
+}
+
+
 /*
  * Check if page is present in the cache.
  * Returns true if page is found in local cache.
@@ -616,6 +1007,7 @@ lfc_evict(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno)
 
 	/* remove the page from the cache */
 	entry->bitmap[chunk_offs >> 5] &= ~(1 << (chunk_offs & (32 - 1)));
+	entry->prewarm_requested = false; /* prohibit prewarm of this LFC entry */
 
 	if (entry->access_count == 0)
 	{
@@ -861,7 +1253,7 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 
 	CriticalAssert(BufTagGetRelNumber(&tag) != InvalidRelFileNumber);
 
-	/* 
+	/*
 	 * For every chunk that has blocks we're interested in, we
 	 * 1. get the chunk header
 	 * 2. Check if the chunk actually has the blocks we're interested in
@@ -899,6 +1291,17 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 
 		if (found)
 		{
+			if (entry->prewarm_started)
+			{
+				/*
+				 * Some page of this chunk is currently written by `lfc_prewarm`.
+				 * We should give-up not to interfere with it.
+				 * But clearing `prewarm_requested` flag also will not allow `lfc_prewarm` to fix it result.
+				 */
+				entry->prewarm_requested = false;
+				LWLockRelease(lfc_lock);
+				return;
+			}
 			/*
 			 * Unlink entry from LRU list to pin it for the duration of IO
 			 * operation
@@ -928,7 +1331,7 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 			{
 				/* Cache overflow: evict least recently used chunk */
 				FileCacheEntry *victim = dlist_container(FileCacheEntry, list_node, dlist_pop_head_node(&lfc_ctl->lru));
-	
+
 				for (int i = 0; i < BLOCKS_PER_CHUNK; i++)
 				{
 					lfc_ctl->used_pages -= (victim->bitmap[i >> 5] >> (i & 31)) & 1;
@@ -944,10 +1347,10 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 				FileCacheEntry *hole = dlist_container(FileCacheEntry, list_node, dlist_pop_head_node(&lfc_ctl->holes));
 				uint32		offset = hole->offset;
 				bool		hole_found;
-	
+
 				hash_search_with_hash_value(lfc_hash, &hole->key, hole->hash, HASH_REMOVE, &hole_found);
 				CriticalAssert(hole_found);
-	
+
 				lfc_ctl->used += 1;
 				entry->offset = offset;	/* reuse the hole */
 			}
@@ -959,9 +1362,11 @@ lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 			}
 			entry->access_count = 1;
 			entry->hash = hash;
+			entry->prewarm_started = false;
 			memset(entry->bitmap, 0, sizeof entry->bitmap);
 		}
 
+		entry->prewarm_requested = false; /* prohibit prewarm if LFC entry is updated by some backend */
 		generation = lfc_ctl->generation;
 		entry_offset = entry->offset;
 		LWLockRelease(lfc_lock);
@@ -1334,3 +1739,74 @@ approximate_working_set_size(PG_FUNCTION_ARGS)
 	}
 	PG_RETURN_NULL();
 }
+
+PG_FUNCTION_INFO_V1(save_local_cache_state);
+
+Datum
+save_local_cache_state(PG_FUNCTION_ARGS)
+{
+	lfc_save_state();
+	PG_RETURN_NULL();
+}
+
+PG_FUNCTION_INFO_V1(get_local_cache_state);
+
+Datum
+get_local_cache_state(PG_FUNCTION_ARGS)
+{
+	size_t n_entries = PG_ARGISNULL(0) ? lfc_prewarm_limit : PG_GETARG_INT32(0);
+	FileCacheStateEntry* fs = lfc_get_state(&n_entries);
+	size_t size_in_bytes = sizeof(FileCacheStateEntry) * n_entries;
+	bytea* res = (bytea*)palloc(VARHDRSZ + size_in_bytes);
+
+	SET_VARSIZE(res, VARHDRSZ + size_in_bytes);
+	memcpy(VARDATA(res), fs, size_in_bytes);
+	pfree(fs);
+
+	PG_RETURN_BYTEA_P(res);
+}
+
+PG_FUNCTION_INFO_V1(prewarm_local_cache);
+
+Datum
+prewarm_local_cache(PG_FUNCTION_ARGS)
+{
+	bytea* state = PG_GETARG_BYTEA_PP(0);
+	uint32 n_entries = VARSIZE_ANY_EXHDR(state);
+	FileCacheStateEntry* fs = (FileCacheStateEntry*)VARDATA_ANY(state);
+
+	lfc_prewarm(fs, n_entries);
+
+	PG_RETURN_NULL();
+}
+
+PG_FUNCTION_INFO_V1(get_prewarm_info);
+
+Datum
+get_prewarm_info(PG_FUNCTION_ARGS)
+{
+	Datum		values[4];
+	bool		nulls[4];
+	TupleDesc	tupdesc;
+
+	if (lfc_size_limit == 0)
+		PG_RETURN_NULL();
+
+	tupdesc = CreateTemplateTupleDesc(4);
+	TupleDescInitEntry(tupdesc, (AttrNumber) 1, "total_chunks", INT4OID, -1, 0);
+	TupleDescInitEntry(tupdesc, (AttrNumber) 2, "curr_chunk", INT4OID, -1, 0);
+	TupleDescInitEntry(tupdesc, (AttrNumber) 3, "prewarmed_pages", INT4OID, -1, 0);
+	TupleDescInitEntry(tupdesc, (AttrNumber) 4, "skipped_pages", INT4OID, -1, 0);
+	tupdesc = BlessTupleDesc(tupdesc);
+
+	MemSet(nulls, 0, sizeof(nulls));
+	LWLockAcquire(lfc_lock, LW_SHARED);
+	values[0] = Int32GetDatum(lfc_ctl->prewarm_total_chunks);
+	values[1] = Int32GetDatum(lfc_ctl->prewarm_curr_chunk);
+	values[2] = Int32GetDatum(lfc_ctl->prewarmed_pages);
+	values[3] = Int32GetDatum(lfc_ctl->skipped_pages);
+	LWLockRelease(lfc_lock);
+
+	PG_RETURN_DATUM(HeapTupleGetDatum(heap_form_tuple(tupdesc, values, nulls)));
+}
+

pgxn/neon/neon--1.5--1.6.sql

@@ -0,0 +1,28 @@
+\echo Use "ALTER EXTENSION neon UPDATE TO '1.6'" to load this file. \quit
+
+CREATE FUNCTION save_local_cache_state()
+RETURNS void
+AS 'MODULE_PATHNAME', 'save_local_cache_state'
+LANGUAGE C STRICT
+PARALLEL UNSAFE;
+
+CREATE FUNCTION get_prewarm_info(out total_chunks integer, out curr_chunk integer, out prewarmed_pages integer, out skipped_pages integer)
+RETURNS record
+AS 'MODULE_PATHNAME', 'get_prewarm_info'
+LANGUAGE C STRICT
+PARALLEL SAFE;
+
+CREATE FUNCTION get_local_cache_state(max_chunks integer default null)
+RETURNS bytea
+AS 'MODULE_PATHNAME', 'get_local_cache_state'
+LANGUAGE C
+PARALLEL UNSAFE;
+
+CREATE FUNCTION prewarm_local_cache(state bytea)
+RETURNS void
+AS 'MODULE_PATHNAME', 'prewarm_local_cache'
+LANGUAGE C STRICT
+PARALLEL UNSAFE;
+
+
+

pgxn/neon/neon--1.6--1.5.sql

@@ -0,0 +1,9 @@
+DROP FUNCTION IF EXISTS save_local_cache_state();
+
+DROP FUNCTION IF EXISTS get_prewarm_info(out total_chunks integer, out curr_chunk integer, out prewarmed_pages integer, out skipped_pages integer);
+
+DROP FUNCTION IF EXISTS get_local_cache_state(max_chunks integer);
+
+DROP FUNCTION IF EXISTS prewarm_local_cache(state bytea);
+
+

pgxn/neon/pagestore_client.h

@@ -276,6 +276,8 @@ extern int lfc_cache_containsv(NRelFileInfo rinfo, ForkNumber forkNum,
 							   BlockNumber blkno, int nblocks, bits8 *bitmap);
 extern void lfc_evict(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno);
 extern void lfc_init(void);
+extern void lfc_save_state(void);
+extern void lfc_load_pages(void);
 
 static inline bool
 lfc_read(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,

proxy/src/context/parquet.rs

@@ -543,7 +543,6 @@ mod tests {
         let remote_storage_config = RemoteStorageConfig {
             storage: RemoteStorageKind::LocalFs {
                 local_path: tmpdir.to_path_buf(),
-                max_keys_per_list_response: None,
             },
             timeout: std::time::Duration::from_secs(120),
         };

test_runner/regress/test_lfc_prewarm.py

@@ -0,0 +1,52 @@
+import time
+
+from fixtures.log_helper import log
+from fixtures.neon_fixtures import NeonEnv
+
+
+def test_lfc_prewarm(neon_simple_env: NeonEnv):
+    env = neon_simple_env
+    n_records = 1000000
+
+    endpoint = env.endpoints.create_start(
+        branch_name="main",
+        config_lines=[
+            "autovacuum = off",
+            "shared_buffers=1MB",
+            "neon.max_file_cache_size=1GB",
+            "neon.file_cache_size_limit=1GB",
+            "neon.file_cache_prewarm_limit=1000",
+        ],
+    )
+    conn = endpoint.connect()
+    cur = conn.cursor()
+    cur.execute("create table t(pk integer primary key, payload text default repeat('?', 128))")
+    cur.execute(f"insert into t (pk) values (generate_series(1,{n_records}))")
+
+    endpoint.stop()
+    endpoint.start()
+
+    conn = endpoint.connect()
+    cur = conn.cursor()
+    cur.execute("create extension neon version '1.6'")
+
+    for _ in range(60):
+        time.sleep(1)  # give prewarm BGW some time to proceed
+        cur.execute("select lfc_value from neon_lfc_stats where lfc_key='file_cache_used_pages'")
+        lfc_used_pages = cur.fetchall()[0][0]
+        log.info(f"Used LFC size: {lfc_used_pages}")
+        cur.execute("select * from get_prewarm_info()")
+        prewarm_info = cur.fetchall()[0]
+        log.info(f"Prewarm info: {prewarm_info}")
+        if prewarm_info[0] > 0:
+            log.info(f"Prewarm progress: {prewarm_info[1]*100//prewarm_info[0]}%")
+            if prewarm_info[0] == prewarm_info[1]:
+                break
+
+    assert lfc_used_pages > 10000
+    assert prewarm_info[0] > 0 and prewarm_info[0] == prewarm_info[1]
+
+    cur.execute("select sum(pk) from t")
+    assert cur.fetchall()[0][0] == n_records * (n_records + 1) / 2
+
+    assert prewarm_info[1] > 0

vendor/revisions.json

@@ -1,18 +1,18 @@
 {
   "v17": [
     "17.0",
-    "68b5038f27e493bde6ae552fe066f10cbdfe6a14"
+    "37d5ead146b028dd9a5c07e7a37068ec0df9f465"
   ],
   "v16": [
     "16.4",
-    "e131a9c027b202ce92bd7b9cf2569d48a6f9948e"
+    "cc36e03bd0c927022cf3b3563e291e42d75366a1"
   ],
   "v15": [
     "15.8",
-    "22e580fe9ffcea7e02592110b1c9bf426d83cada"
+    "a4830163a65811578824ce4022c1cd3daef33d4e"
   ],
   "v14": [
     "14.13",
-    "2199b83fb72680001ce0f43bf6187a21dfb8f45d"
+    "ecb1020ff71927e9dd59c526254bb8846bb73ee1"
   ]
 }