libsql-server: cap individual snapshot files at SQLD_MAX_SNAPSHOT_SIZE

Catalog-init for large shops downloads the whole .snap in a single Snapshot
RPC because the compactor merges accumulated .snap files into one giant
blob whenever their cumulative size exceeds 2x the live database
(SNAPHOT_SPACE_AMPLIFICATION_FACTOR). Combined with a 200 MB max_log_size,
the result is a single 200-500 MB streaming RPC that routinely runs over
the 100 s SFE cap on degraded mobile links.

Add a new SQLD_MAX_SNAPSHOT_SIZE env var (in MB, mirroring
SQLD_MAX_LOG_SIZE). When set:

  - should_compact triggers a merge when the cumulative size of accumulated
    snapshots reaches the cap, instead of using the 2x amplification rule.
  - merge_snapshots groups input snapshots greedily into contiguous batches
    whose summed frame count fits under the cap, producing one output file
    per batch instead of a single combined blob.
  - A pre-existing snapshot whose own size already exceeds the cap is left
    in place as a singleton batch, so the merger never produces a file
    larger than the configured cap.

When unset, behavior is unchanged: legacy 2x-db-page-count amplification,
single combined merge output. The env var is read once via OnceLock to keep
the change local to snapshot.rs (no DbConfig/PrimaryConfig plumbing) and
small for a fork patch.

Pairs with lowering SQLD_MAX_LOG_SIZE in the production env so each .snap
written by the compactor is small (e.g., 20 MB), making the existing
client-side replicator loop stream a 500 MB catalog as ~25 small Snapshot
RPCs instead of one giant one.

Refs: ae-task 259, GSD 48518

Author: RaVbaker

libsql-server/src/replication/snapshot.rs

@@ -3,7 +3,7 @@ use std::io::SeekFrom;
 use std::mem::size_of;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
 use std::time::Instant;
 
 use anyhow::{bail, Context};
@@ -26,10 +26,101 @@ use super::FrameNo;
 
 /// This is the ratio of the space required to store snapshot vs size of the actual database.
 /// When this ratio is exceeded, compaction is triggered.
+///
+/// Used only when `SQLD_MAX_SNAPSHOT_SIZE` is unset (legacy behavior).
 const SNAPHOT_SPACE_AMPLIFICATION_FACTOR: u64 = 2;
-/// The maximum amount of snapshot allowed before a compaction is required
+/// The default maximum number of snapshot files allowed before a compaction is required.
+/// Overridable via `SQLD_MAX_SNAPSHOT_COUNT`.
 const MAX_SNAPSHOT_NUMBER: usize = 32;
 
+/// Returns the configured upper bound, in frames, on the cumulative size of accumulated
+/// snapshot files before merging is triggered, and on the size of any single merged
+/// snapshot file produced by the merger. Sourced once from the `SQLD_MAX_SNAPSHOT_SIZE`
+/// environment variable (in MB, mirroring `SQLD_MAX_LOG_SIZE`).
+///
+/// When unset, the legacy `SNAPHOT_SPACE_AMPLIFICATION_FACTOR * db_page_count` rule is
+/// used and the merger combines all accumulated snapshots into a single file regardless
+/// of size.
+/// Group an ordered list of accumulated snapshots into contiguous batches for merging,
+/// ensuring no batch's cumulative frame count exceeds `max_frames`.
+///
+/// Returns ranges into `snapshots`. With `max_frames = None`, returns a single range
+/// covering the whole input (legacy single-file merge). An input snapshot whose own
+/// frame count already exceeds `max_frames` is placed in a singleton batch so the merger
+/// can leave it as-is rather than produce an even larger file.
+fn group_snapshots_for_merge(
+    snapshots: &[(String, u64)],
+    max_frames: Option<u64>,
+) -> Vec<std::ops::Range<usize>> {
+    if snapshots.is_empty() {
+        return Vec::new();
+    }
+    let Some(max) = max_frames else {
+        return vec![0..snapshots.len()];
+    };
+
+    let mut batches: Vec<std::ops::Range<usize>> = Vec::new();
+    let mut start = 0usize;
+    let mut acc: u64 = 0;
+    for (i, (_, count)) in snapshots.iter().enumerate() {
+        // If this single file already exceeds the cap, flush any pending batch and emit
+        // it as its own singleton.
+        if *count >= max {
+            if i > start {
+                batches.push(start..i);
+            }
+            batches.push(i..i + 1);
+            start = i + 1;
+            acc = 0;
+            continue;
+        }
+        if acc + count > max && i > start {
+            batches.push(start..i);
+            start = i;
+            acc = 0;
+        }
+        acc += count;
+    }
+    if start < snapshots.len() {
+        batches.push(start..snapshots.len());
+    }
+    batches
+}
+
+fn max_snapshot_frames() -> Option<u64> {
+    static CACHED: OnceLock<Option<u64>> = OnceLock::new();
+    *CACHED.get_or_init(|| {
+        let mb = std::env::var("SQLD_MAX_SNAPSHOT_SIZE")
+            .ok()?
+            .parse::<u64>()
+            .ok()?;
+        if mb == 0 {
+            return None;
+        }
+        Some(mb * 1_000_000 / LogFile::FRAME_SIZE as u64)
+    })
+}
+
+/// Maximum number of accumulated snapshot files before the merger is forced to compact,
+/// independently of total size. Sourced from `SQLD_MAX_SNAPSHOT_COUNT`; falls back to
+/// `MAX_SNAPSHOT_NUMBER` (32) when unset.
+///
+/// When pairing a low `SQLD_MAX_SNAPSHOT_SIZE` with a low `SQLD_MAX_LOG_SIZE` (so that
+/// each `.snap` file stays small on disk), this count needs to be raised in lockstep,
+/// otherwise the count-trigger fires and undoes the chunking by merging many small files
+/// into one large one. Rule of thumb: set it to at least
+/// `expected_total_snapshot_bytes / max_log_size_bytes`.
+fn max_snapshot_count() -> usize {
+    static CACHED: OnceLock<usize> = OnceLock::new();
+    *CACHED.get_or_init(|| {
+        std::env::var("SQLD_MAX_SNAPSHOT_COUNT")
+            .ok()
+            .and_then(|s| s.parse::<usize>().ok())
+            .filter(|n| *n > 0)
+            .unwrap_or(MAX_SNAPSHOT_NUMBER)
+    })
+}
+
 /// returns (db_id, start_frame_no, end_frame_no) for the given snapshot name
 fn parse_snapshot_name(name: &str) -> Option<(Uuid, u64, u64)> {
     let (db_id_str, remaining) = name.split_at(36);
@@ -271,8 +362,11 @@ impl SnapshotMerger {
 
     fn should_compact(snapshots: &[(String, u64)], db_page_count: u32) -> bool {
         let snapshots_size: u64 = snapshots.iter().map(|(_, s)| *s).sum();
-        snapshots_size >= SNAPHOT_SPACE_AMPLIFICATION_FACTOR * db_page_count as u64
-            || snapshots.len() > MAX_SNAPSHOT_NUMBER
+        let size_trigger = match max_snapshot_frames() {
+            Some(max) => snapshots_size >= max,
+            None => snapshots_size >= SNAPHOT_SPACE_AMPLIFICATION_FACTOR * db_page_count as u64,
+        };
+        size_trigger || snapshots.len() > max_snapshot_count()
     }
 
     async fn run_snapshot_merger_loop(
@@ -310,8 +404,13 @@ impl SnapshotMerger {
                     working = false;
                     job.set(std::future::pending());
                     let ret = ret??;
-                    // the new merged snapshot is prepended to the snapshot list
-                    snapshots.insert(0, ret);
+                    // The merged snapshot(s) cover the oldest frames, so they are prepended
+                    // to the snapshot list in chronological (oldest-first) order. With
+                    // SQLD_MAX_SNAPSHOT_SIZE set, merge_snapshots may return more than one
+                    // surviving snapshot.
+                    for (i, item) in ret.into_iter().enumerate() {
+                        snapshots.insert(i, item);
+                    }
                 }
                 else => return Ok(())
             }
@@ -359,41 +458,69 @@ impl SnapshotMerger {
         log_id: Uuid,
         scripted_backup: Option<ScriptBackupManager>,
         namespace: NamespaceName,
-    ) -> anyhow::Result<(String, u64)> {
-        let mut builder = SnapshotBuilder::new(db_path, log_id, scripted_backup, namespace).await?;
+    ) -> anyhow::Result<Vec<(String, u64)>> {
+        // When SQLD_MAX_SNAPSHOT_SIZE is set, group input snapshots greedily into batches
+        // whose summed frame count fits under the cap. Each batch becomes one output
+        // `.snap` file, ensuring no merger-produced file ever exceeds the cap. A batch of
+        // size 1 is left as-is (skip the merge copy). When unset, fall back to the legacy
+        // behavior: collapse everything into a single file.
+        let max_frames = max_snapshot_frames();
+        let batches = group_snapshots_for_merge(&snapshots, max_frames);
+        tracing::debug!(
+            "merging {} snapshots for {log_id} into {} batch(es) (max_frames={:?})",
+            snapshots.len(),
+            batches.len(),
+            max_frames,
+        );
+
         let snapshot_dir_path = snapshot_dir_path(db_path);
-        let mut size_after = None;
-        tracing::debug!("merging {} snashots for {log_id}", snapshots.len());
-        for (name, _) in snapshots.iter().rev() {
-            // NOTICE: no encryptor passed in order to read frames as is, still encrypted
-            let snapshot = SnapshotFile::open(&snapshot_dir_path.join(name), None).await?;
-            // The size after the merged snapshot is the size after the first snapshot to be merged
-            if size_after.is_none() {
-                size_after.replace(snapshot.header().size_after);
+        let mut out: Vec<(String, u64)> = Vec::with_capacity(batches.len());
+        for batch in batches {
+            let batch = &snapshots[batch];
+            if batch.len() == 1 {
+                // Single oversized or stand-alone snapshot — leave it in place untouched
+                // so that we never produce a merged file larger than the configured cap.
+                out.push(batch[0].clone());
+                continue;
             }
-            builder
-                .append_frames(snapshot.into_stream_mut().map_err(|e| anyhow::anyhow!(e)))
-                .await?;
-        }
 
-        let (_, start_frame_no, _) = parse_snapshot_name(&snapshots[0].0).unwrap();
-        let (_, _, end_frame_no) = parse_snapshot_name(&snapshots.last().unwrap().0).unwrap();
+            let mut builder =
+                SnapshotBuilder::new(db_path, log_id, scripted_backup.clone(), namespace.clone())
+                    .await?;
+            let mut size_after = None;
+            for (name, _) in batch.iter().rev() {
+                // NOTICE: no encryptor passed in order to read frames as is, still encrypted
+                let snapshot = SnapshotFile::open(&snapshot_dir_path.join(name), None).await?;
+                // The size after the merged snapshot is the size after the first snapshot to be merged
+                if size_after.is_none() {
+                    size_after.replace(snapshot.header().size_after);
+                }
+                builder
+                    .append_frames(snapshot.into_stream_mut().map_err(|e| anyhow::anyhow!(e)))
+                    .await?;
+            }
 
-        tracing::debug!(
-            "created merged snapshot for {log_id} from frame {start_frame_no} to {end_frame_no}"
-        );
+            let (_, start_frame_no, _) = parse_snapshot_name(&batch[0].0).unwrap();
+            let (_, _, end_frame_no) = parse_snapshot_name(&batch.last().unwrap().0).unwrap();
 
-        builder.header.start_frame_no = start_frame_no.into();
-        builder.header.end_frame_no = end_frame_no.into();
-        builder.header.size_after = size_after.unwrap();
+            tracing::debug!(
+                "created merged snapshot for {log_id} from frame {start_frame_no} to {end_frame_no}"
+            );
 
-        let meta = builder.finish().await?;
+            builder.header.start_frame_no = start_frame_no.into();
+            builder.header.end_frame_no = end_frame_no.into();
+            builder.header.size_after = size_after.unwrap();
 
-        for (name, _) in snapshots.iter() {
-            tokio::fs::remove_file(&snapshot_dir_path.join(name)).await?;
+            let meta = builder.finish().await?;
+
+            for (name, _) in batch.iter() {
+                tokio::fs::remove_file(&snapshot_dir_path.join(name)).await?;
+            }
+
+            out.push((meta.0, meta.1));
         }
 
-        Ok((meta.0, meta.1))
+        Ok(out)
     }
 
     async fn register_snapshot(
@@ -582,6 +709,48 @@ mod test {
     use crate::replication::snapshot::SnapshotFile;
     use crate::LIBSQL_PAGE_SIZE;
 
+    fn s(name: &str, count: u64) -> (String, u64) {
+        (name.to_string(), count)
+    }
+
+    #[test]
+    fn group_snapshots_legacy_no_max_returns_single_batch() {
+        let snaps = vec![s("a", 10), s("b", 20), s("c", 5)];
+        let batches = super::group_snapshots_for_merge(&snaps, None);
+        assert_eq!(batches, vec![0..3]);
+    }
+
+    #[test]
+    fn group_snapshots_empty_input() {
+        let batches = super::group_snapshots_for_merge(&[], Some(100));
+        assert!(batches.is_empty());
+    }
+
+    #[test]
+    fn group_snapshots_packs_under_max() {
+        // max = 25, sequence 10,10,10,10 → (10+10) (10+10)
+        let snaps = vec![s("a", 10), s("b", 10), s("c", 10), s("d", 10)];
+        let batches = super::group_snapshots_for_merge(&snaps, Some(25));
+        assert_eq!(batches, vec![0..2, 2..4]);
+    }
+
+    #[test]
+    fn group_snapshots_oversized_input_is_singleton() {
+        // Pre-existing big file (50) exceeds cap (25). It must be left alone in its own
+        // batch so the merger does not produce an even bigger merged file.
+        let snaps = vec![s("a", 10), s("big", 50), s("c", 10), s("d", 10)];
+        let batches = super::group_snapshots_for_merge(&snaps, Some(25));
+        assert_eq!(batches, vec![0..1, 1..2, 2..4]);
+    }
+
+    #[test]
+    fn group_snapshots_each_under_cap_but_sum_over() {
+        // max=15, sequence 10,10,10 → each fits alone (10<15), pairs do not (20>15).
+        let snaps = vec![s("a", 10), s("b", 10), s("c", 10)];
+        let batches = super::group_snapshots_for_merge(&snaps, Some(15));
+        assert_eq!(batches, vec![0..1, 1..2, 2..3]);
+    }
+
     use super::*;
 
     async fn dir_is_empty(p: &Path) -> bool {