Implement autoprewarming in Neon uig pg_prewarm extension

pageserver/src/basebackup.rs

@@ -17,6 +17,7 @@ use itertools::Itertools;
 use std::fmt::Write as FmtWrite;
 use std::io;
 use std::io::Write;
+use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::SystemTime;
 use tar::{Builder, EntryType, Header};
@@ -25,8 +26,10 @@ use tracing::*;
 use crate::tenant::Timeline;
 use pageserver_api::reltag::{RelTag, SlruKind};
 
+use postgres_ffi::pg_constants::{
+    AUTOPREWARM_FILE_NAME, PGDATA_SPECIAL_FILES, PGDATA_SUBDIRS, PG_HBA,
+};
 use postgres_ffi::pg_constants::{DEFAULTTABLESPACE_OID, GLOBALTABLESPACE_OID};
-use postgres_ffi::pg_constants::{PGDATA_SPECIAL_FILES, PGDATA_SUBDIRS, PG_HBA};
 use postgres_ffi::TransactionId;
 use postgres_ffi::XLogFileName;
 use postgres_ffi::PG_TLI;
@@ -145,6 +148,7 @@ where
                 self.ar.append(&header, &mut io::empty())?;
             }
         }
+        self.add_prewarm_file()?;
 
         // Gather non-relational files from object storage pages.
         for kind in [
@@ -218,6 +222,21 @@ where
         Ok(())
     }
 
+    //
+    // Include "autoprewarm-bin.blocks" in archive (if exists)
+    //
+    fn add_prewarm_file(&mut self) -> anyhow::Result<()> {
+        let path = self
+            .timeline
+            .conf
+            .timeline_path(&self.timeline.timeline_id, &self.timeline.tenant_id)
+            .join(AUTOPREWARM_FILE_NAME);
+        if PathBuf::from(&path).exists() {
+            self.ar.append_path_with_name(path, AUTOPREWARM_FILE_NAME)?;
+        }
+        Ok(())
+    }
+
     //
     // Generate SLRU segment files from repository.
     //

pageserver/src/page_service.rs

@@ -15,10 +15,11 @@ use futures::{Stream, StreamExt};
 use pageserver_api::models::{
     PagestreamBeMessage, PagestreamDbSizeRequest, PagestreamDbSizeResponse,
     PagestreamErrorResponse, PagestreamExistsRequest, PagestreamExistsResponse,
-    PagestreamFeMessage, PagestreamGetPageRequest, PagestreamGetPageResponse,
-    PagestreamNblocksRequest, PagestreamNblocksResponse,
+    PagestreamFcntlRequest, PagestreamFeMessage, PagestreamGetPageRequest,
+    PagestreamGetPageResponse, PagestreamNblocksRequest, PagestreamNblocksResponse,
 };
 use std::io;
+use std::io::Write;
 use std::net::TcpListener;
 use std::str;
 use std::str::FromStr;
@@ -45,9 +46,12 @@ use crate::task_mgr;
 use crate::task_mgr::TaskKind;
 use crate::tenant::Timeline;
 use crate::tenant_mgr;
+use crate::virtual_file::VirtualFile;
 use crate::CheckpointConfig;
-
-use postgres_ffi::pg_constants::DEFAULTTABLESPACE_OID;
+use crate::TEMP_FILE_SUFFIX;
+use postgres_ffi::pg_constants::{
+    AUTOPREWARM_FILE_NAME, DEFAULTTABLESPACE_OID, SMGR_FCNTL_CACHE_SNAPSHOT,
+};
 use postgres_ffi::BLCKSZ;
 
 fn copyin_stream(pgb: &mut PostgresBackend) -> impl Stream<Item = io::Result<Bytes>> + '_ {
@@ -317,6 +321,10 @@ impl PageServerHandler {
                     let _timer = metrics.get_db_size.start_timer();
                     self.handle_db_size_request(&timeline, &req).await
                 }
+                PagestreamFeMessage::Fcntl(req) => {
+                    self.handle_fcntl_request(&timeline, &req).await?;
+                    continue;
+                }
             };
 
             let response = response.unwrap_or_else(|e| {
@@ -586,6 +594,33 @@ impl PageServerHandler {
         }))
     }
 
+    async fn handle_fcntl_request(
+        &self,
+        timeline: &Timeline,
+        req: &PagestreamFcntlRequest,
+    ) -> Result<()> {
+        if req.cmd == SMGR_FCNTL_CACHE_SNAPSHOT {
+            let temp_path = self
+                .conf
+                .timeline_path(&timeline.timeline_id, &timeline.tenant_id)
+                .join(format!("{AUTOPREWARM_FILE_NAME}.{TEMP_FILE_SUFFIX}"));
+            let mut file = VirtualFile::open_with_options(
+                &temp_path,
+                std::fs::OpenOptions::new().write(true).create_new(true),
+            )?;
+            file.write_all(&req.data)?;
+            drop(file);
+            let final_path = self
+                .conf
+                .timeline_path(&timeline.timeline_id, &timeline.tenant_id)
+                .join(AUTOPREWARM_FILE_NAME);
+            std::fs::rename(temp_path, &final_path)?;
+        } else {
+            warn!("Fcntl request {} is not supported", req.cmd);
+        }
+        Ok(())
+    }
+
     #[instrument(skip(self, pgb))]
     async fn handle_basebackup_request(
         &self,

pageserver/src/tenant/timeline.rs

@@ -62,8 +62,8 @@ use crate::{
 };
 
 pub struct Timeline {
-    conf: &'static PageServerConf,
-    tenant_conf: Arc<RwLock<TenantConfOpt>>,
+    pub conf: &'static PageServerConf,
+    pub tenant_conf: Arc<RwLock<TenantConfOpt>>,
 
     pub tenant_id: TenantId,
     pub timeline_id: TimelineId,