Improve memory

server-rs/src/main.rs

@@ -19,6 +19,25 @@ mod routes;
 mod state;
 pub mod utils;
 
+// Use jemalloc as the global allocator. glibc malloc keeps freed memory in many
+// per-CPU arenas and rarely returns it to the OS, so the transient buffers from
+// the rayon/Polars-parallel data load stay resident and inflate steady-state RSS.
+// jemalloc's decay-based purging (configured below) hands those pages back.
+#[cfg(not(target_env = "msvc"))]
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
+// Return dirty/muzzy pages to the OS ~1s after they go idle, and run a background
+// thread to do so proactively (so RSS drops after the startup load peak without
+// waiting for the next allocation). Read by jemalloc at startup via the
+// `malloc_conf` symbol (unprefixed on Linux). Can be overridden by the
+// `_RJEM_MALLOC_CONF` / `MALLOC_CONF` env var.
+#[cfg(not(target_env = "msvc"))]
+#[allow(non_upper_case_globals)]
+#[used]
+#[export_name = "malloc_conf"]
+pub static malloc_conf: &[u8] = b"background_thread:true,dirty_decay_ms:1000,muzzy_decay_ms:1000\0";
+
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::time::Duration;

server-rs/src/routes/tiles.rs

@@ -1,13 +1,64 @@
+use std::os::unix::fs::FileExt;
 use std::sync::Arc;
 
 use axum::extract::{Path, Query, State};
 use axum::http::{header, StatusCode};
 use axum::response::{IntoResponse, Response};
-use pmtiles::{AsyncPmTilesReader, MmapBackend, TileCoord};
+use bytes::Bytes;
+use pmtiles::{
+    AsyncBackend, AsyncPmTilesReader, BackendResponse, HashMapCache, PmtError, PmtResult, TileCoord,
+};
 use serde::Deserialize;
 use tracing::warn;
 
-pub type TileReader = AsyncPmTilesReader<MmapBackend>;
+/// PMTiles archives are read straight from disk with positional `pread` calls
+/// instead of being memory-mapped. With an mmap backend every touched tile page
+/// is attributed to the process RSS (~21 GB across all tilesets); with `pread`
+/// the file contents stay in the (reclaimable) kernel page cache and only the
+/// small per-request buffer lives in the process. A [`HashMapCache`] keeps the
+/// archive's directory entries in memory so tile lookups don't re-read and
+/// re-decompress directory pages from disk on every request.
+pub type TileReader = AsyncPmTilesReader<FileBackend, HashMapCache>;
+
+/// An [`AsyncBackend`] that serves byte ranges from a local file via `pread`.
+pub struct FileBackend {
+    file: Arc<std::fs::File>,
+    len: u64,
+}
+
+impl FileBackend {
+    pub fn open(path: &std::path::Path) -> std::io::Result<Self> {
+        let file = std::fs::File::open(path)?;
+        let len = file.metadata()?.len();
+        Ok(Self {
+            file: Arc::new(file),
+            len,
+        })
+    }
+}
+
+impl AsyncBackend for FileBackend {
+    async fn read(&self, offset: usize, length: usize) -> PmtResult<BackendResponse> {
+        let available = (self.len as usize).saturating_sub(offset);
+        let read_len = length.min(available);
+        if read_len == 0 {
+            return Ok(BackendResponse::new(Bytes::new()));
+        }
+
+        let file = Arc::clone(&self.file);
+        // pread is blocking; keep it off the async runtime's worker threads.
+        let buf = tokio::task::spawn_blocking(move || {
+            let mut buf = vec![0u8; read_len];
+            file.read_exact_at(&mut buf, offset as u64)?;
+            std::io::Result::Ok(buf)
+        })
+        .await
+        .map_err(|err| PmtError::Reading(std::io::Error::other(err)))?
+        .map_err(PmtError::Reading)?;
+
+        Ok(BackendResponse::new(Bytes::from(buf)))
+    }
+}
 
 pub async fn get_tile(
     State(reader): State<Arc<TileReader>>,
@@ -260,7 +311,7 @@ fn build_style(is_dark: bool, layers: &[serde_json::Value], tile_url: &str) -> s
 }
 
 pub async fn init_tile_reader(path: &std::path::Path) -> anyhow::Result<TileReader> {
-    let backend = MmapBackend::try_from(path).await?;
-    let reader = AsyncPmTilesReader::try_from_source(backend).await?;
+    let backend = FileBackend::open(path)?;
+    let reader = AsyncPmTilesReader::try_from_cached_source(backend, HashMapCache::default()).await?;
     Ok(reader)
 }