Allow capping data to N blocks with API retention NET-896

crates/hotblocks/src/cli.rs

@@ -98,7 +98,7 @@ impl CLI {
 
         let api_controlled_datasets = datasets
             .iter()
-            .filter_map(|(id, cfg)| (cfg.retention_strategy == RetentionConfig::Api).then_some(*id))
+            .filter_map(|(id, cfg)| matches!(cfg.retention_strategy, RetentionConfig::Api { .. }).then_some(*id))
             .collect();
 
         let data_service = DataService::start(db.clone(), datasets).await.map(Arc::new)?;

crates/hotblocks/src/data_service.rs

@@ -46,16 +46,21 @@ impl DataService {
                     .map(|url| ReqwestDataClient::new(http_client.clone(), url))
                     .collect();
 
-                let retention = match cfg.retention_strategy {
-                    RetentionConfig::FromBlock { number, parent_hash } => {
-                        RetentionStrategy::FromBlock { number, parent_hash }
-                    }
-                    RetentionConfig::Head(n) => RetentionStrategy::Head(n),
-                    RetentionConfig::Api | RetentionConfig::None => RetentionStrategy::None
+                let (retention, max_blocks) = match &cfg.retention_strategy {
+                    RetentionConfig::FromBlock { number, parent_hash } => (
+                        RetentionStrategy::FromBlock {
+                            number: *number,
+                            parent_hash: parent_hash.clone()
+                        },
+                        None
+                    ),
+                    RetentionConfig::Head(n) => (RetentionStrategy::Head(*n), None),
+                    RetentionConfig::Api { max_blocks } => (RetentionStrategy::None, *max_blocks),
+                    RetentionConfig::None => (RetentionStrategy::None, None)
                 };
 
                 tokio::task::spawn_blocking(move || {
-                    DatasetController::new(db, dataset_id, cfg.kind, retention, data_sources).map(|c| {
+                    DatasetController::new(db, dataset_id, cfg.kind, retention, max_blocks, data_sources).map(|c| {
                         c.enable_compaction(!cfg.disable_compaction);
                         Arc::new(c)
                     })

crates/hotblocks/src/dataset_config.rs

@@ -1,13 +1,16 @@
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, fmt};
 
-use serde::{Deserialize, Serialize};
+use serde::{
+    Deserialize, Deserializer, Serialize,
+    de::{self, IgnoredAny, MapAccess, Visitor}
+};
 use sqd_query::BlockNumber;
 use sqd_storage::db::DatasetId;
 use url::Url;
 
 use crate::types::DatasetKind;
 
-#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Eq, PartialEq, Serialize)]
 pub enum RetentionConfig {
     // Fixed, starting from the block number
     FromBlock {
@@ -16,11 +19,93 @@ pub enum RetentionConfig {
     },
     // Moving window that keeps up to N blocks
     Head(u64),
-    // Retention is set dynamically from the portal
-    Api,
+    // Retention is set dynamically from the portal. `max_blocks` is a soft cap
+    // applied when the portal stops advancing the floor, see `Ctl::max_blocks`.
+    Api {
+        max_blocks: Option<u64>
+    },
     None
 }
 
+const RETENTION_VARIANTS: &[&str] = &["FromBlock", "Head", "Api", "None"];
+
+// Hand-written to accept both the bare `Api` string and `Api: { max_blocks: N }`.
+// The config is read through `singleton_map_recursive`, which encodes unit
+// variants as strings and struct variants as maps, so a derived impl can accept
+// only one of the two forms. Can be derived again once no config uses bare `Api`.
+impl<'de> Deserialize<'de> for RetentionConfig {
+    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
+        #[derive(Deserialize)]
+        #[serde(deny_unknown_fields)]
+        struct FromBlockCfg {
+            number: BlockNumber,
+            #[serde(default)]
+            parent_hash: Option<String>
+        }
+
+        #[derive(Deserialize)]
+        #[serde(deny_unknown_fields)]
+        struct ApiCfg {
+            #[serde(default)]
+            max_blocks: Option<u64>
+        }
+
+        struct RetentionVisitor;
+
+        impl<'de> Visitor<'de> for RetentionVisitor {
+            type Value = RetentionConfig;
+
+            fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
+                f.write_str("a retention strategy")
+            }
+
+            fn visit_str<E: de::Error>(self, value: &str) -> Result<Self::Value, E> {
+                match value {
+                    "Api" => Ok(RetentionConfig::Api { max_blocks: None }),
+                    "None" => Ok(RetentionConfig::None),
+                    other => Err(E::unknown_variant(other, RETENTION_VARIANTS))
+                }
+            }
+
+            fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<Self::Value, A::Error> {
+                let tag: String = map
+                    .next_key()?
+                    .ok_or_else(|| de::Error::custom("expected a retention strategy"))?;
+
+                let strategy = match tag.as_str() {
+                    "FromBlock" => {
+                        let cfg: FromBlockCfg = map.next_value()?;
+                        RetentionConfig::FromBlock {
+                            number: cfg.number,
+                            parent_hash: cfg.parent_hash
+                        }
+                    }
+                    "Head" => RetentionConfig::Head(map.next_value()?),
+                    "Api" => {
+                        let cfg: ApiCfg = map.next_value()?;
+                        RetentionConfig::Api {
+                            max_blocks: cfg.max_blocks
+                        }
+                    }
+                    "None" => {
+                        map.next_value::<IgnoredAny>()?;
+                        RetentionConfig::None
+                    }
+                    other => return Err(de::Error::unknown_variant(other, RETENTION_VARIANTS))
+                };
+
+                if map.next_key::<IgnoredAny>()?.is_some() {
+                    return Err(de::Error::custom("retention strategy must have a single key"));
+                }
+
+                Ok(strategy)
+            }
+        }
+
+        deserializer.deserialize_any(RetentionVisitor)
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(deny_unknown_fields)]
 pub struct DatasetConfig {
@@ -39,3 +124,41 @@ impl DatasetConfig {
         Ok(config)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // Mirror how the field is read inside `read_config_file`.
+    fn parse(yaml: &str) -> Result<RetentionConfig, serde_yaml::Error> {
+        let deser = serde_yaml::Deserializer::from_str(yaml);
+        serde_yaml::with::singleton_map_recursive::deserialize(deser)
+    }
+
+    #[test]
+    fn api_accepts_both_forms() {
+        assert_eq!(parse("Api").unwrap(), RetentionConfig::Api { max_blocks: None });
+        assert_eq!(parse("Api: {}").unwrap(), RetentionConfig::Api { max_blocks: None });
+        assert_eq!(
+            parse("Api:\n  max_blocks: 100000").unwrap(),
+            RetentionConfig::Api {
+                max_blocks: Some(100000)
+            }
+        );
+        assert!(parse("Api:\n  max_blcks: 5").is_err());
+    }
+
+    #[test]
+    fn other_strategies_still_parse() {
+        assert_eq!(parse("None").unwrap(), RetentionConfig::None);
+        assert_eq!(parse("Head: 2000").unwrap(), RetentionConfig::Head(2000));
+        assert_eq!(
+            parse("FromBlock:\n  number: 10\n  parent_hash: '0xabc'").unwrap(),
+            RetentionConfig::FromBlock {
+                number: 10,
+                parent_hash: Some("0xabc".to_owned())
+            }
+        );
+        assert!(parse("Bogus").is_err());
+    }
+}

crates/hotblocks/src/dataset_controller/dataset_controller.rs

@@ -42,6 +42,7 @@ impl DatasetController {
         dataset_id: DatasetId,
         dataset_kind: DatasetKind,
         retention: RetentionStrategy,
+        max_blocks: Option<u64>,
         data_sources: Vec<ReqwestDataClient>
     ) -> anyhow::Result<Self> {
         let mut write = WriteController::new(db.clone(), dataset_id, dataset_kind)?;
@@ -63,6 +64,7 @@ impl DatasetController {
             dataset_id,
             dataset_kind,
             data_sources,
+            max_blocks,
             retention_recv,
             head_sender,
             finalized_head_sender
@@ -162,12 +164,9 @@ impl WriteCtx {
                 self.notify_head();
                 self.notify_finalized_head();
                 if let Some(n) = head {
-                    if self
-                        .write
-                        .first_chunk_head()
-                        .map_or(false, |h| self.write.next_block() - h.number > n)
-                    {
-                        self.retain(self.write.next_block() - n, None)?;
+                    let first_chunk_head = self.write.first_chunk_head().map(|h| h.number);
+                    if let Some(floor) = trim_floor(first_chunk_head, self.write.next_block(), n) {
+                        self.retain(floor, None)?;
                     }
                 }
             }
@@ -248,6 +247,16 @@ fn send_if_new<T: Eq>(sender: &tokio::sync::watch::Sender<T>, value: T) {
     });
 }
 
+// Returns the new floor when the tail has to be trimmed to keep
+// `max_blocks` behind the tip, or `None` when the window still fits.
+//
+// `max_blocks` is a soft limit. Since `retain()` only drops whole chunks, trimming
+// may keep a part of the first chunk.
+fn trim_floor(first_chunk_head: Option<BlockNumber>, next_block: BlockNumber, max_blocks: u64) -> Option<BlockNumber> {
+    let first_chunk_head = first_chunk_head?;
+    (next_block - first_chunk_head > max_blocks).then(|| next_block - max_blocks)
+}
+
 enum State {
     Idle,
     Init {
@@ -283,6 +292,10 @@ struct Ctl {
     dataset_id: DatasetId,
     dataset_kind: DatasetKind,
     data_sources: Vec<ReqwestDataClient>,
+    // Safety cap on retained blocks for Api-controlled datasets: even if the portal
+    // stops advancing the floor, the tail is trimmed to keep roughly this many blocks
+    // behind the tip. `None` means grow indefinitely.
+    max_blocks: Option<u64>,
     retention_recv: tokio::sync::watch::Receiver<RetentionStrategy>,
     head_sender: tokio::sync::watch::Sender<Option<BlockRef>>,
     finalized_head_sender: tokio::sync::watch::Sender<Option<BlockRef>>
@@ -343,17 +356,18 @@ impl Ctl {
         let retention = self.retention_recv.borrow_and_update().clone();
         let mut state = match retention {
             RetentionStrategy::FromBlock { number, parent_hash } => {
+                let (number, parent_hash) = self.clamp_floor(&write, number, parent_hash);
                 if !write.starts_at(number, &parent_hash) {
                     blocking! {
                         write.retain(number, parent_hash)
                     }?;
                 }
-                State::Init { head: None }
+                State::Init { head: self.max_blocks }
             }
             RetentionStrategy::Head(n) => State::Init { head: Some(n) },
             RetentionStrategy::None => {
                 if write.write.head().is_some() {
-                    State::Init { head: None }
+                    State::Init { head: self.max_blocks }
                 } else {
                     State::Idle
                 }
@@ -456,17 +470,32 @@ impl Ctl {
         }
     }
 
+    // Don't allow auto-moving the floor backwards because it causes a full resync.
+    fn clamp_floor(
+        &self,
+        write: &WriteCtx,
+        number: BlockNumber,
+        parent_hash: Option<String>
+    ) -> (BlockNumber, Option<String>) {
+        if self.max_blocks.is_some() && number < write.write.start_block() {
+            (write.write.start_block(), None)
+        } else {
+            (number, parent_hash)
+        }
+    }
+
     async fn handle_retention_change(&mut self, state: &mut State, mut write: WriteCtx) -> anyhow::Result<WriteCtx> {
         // need this variable to please the compiler
         let retention = self.retention_recv.borrow_and_update().clone();
         match retention {
             RetentionStrategy::FromBlock { number, parent_hash } => {
+                let (number, parent_hash) = self.clamp_floor(&write, number, parent_hash);
                 let will_erase_head = write.write.head().map_or(false, |h| h.number < number) || // FromBlock is greater than current head, so everything is cleared
                     write.write.start_block() > number; // FromBlock is less than current front, dropping everything by design
                 blocking_write!(write, write.retain(number, parent_hash))?;
                 match state {
                     State::Ingest { .. } if !will_erase_head => {} // Keep ingesting, head is valid
-                    _ => *state = State::Init { head: None }       // New ingest needed
+                    _ => *state = State::Init { head: self.max_blocks } // New ingest needed
                 }
             }
             RetentionStrategy::Head(n) => match state {
@@ -650,3 +679,26 @@ async fn compaction_loop(db: DBRef, dataset_id: DatasetId, mut enabled: tokio::s
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::trim_floor;
+
+    #[test]
+    fn nothing_is_trimmed_while_the_window_fits() {
+        assert_eq!(trim_floor(None, 500, 100), None);
+        // The whole dataset is one chunk [0..50], well inside the cap.
+        assert_eq!(trim_floor(Some(50), 51, 100), None);
+        // Exactly at the cap: the first chunk still has a block in the window.
+        assert_eq!(trim_floor(Some(0), 100, 100), None);
+    }
+
+    #[test]
+    fn the_tail_is_trimmed_once_the_first_chunk_leaves_the_window() {
+        // First chunk ends at 0, so trimming starts one block past the cap.
+        assert_eq!(trim_floor(Some(0), 101, 100), Some(1));
+        // The soft-limit overshoot: [0..150K] under a 100K cap survives until 250K.
+        assert_eq!(trim_floor(Some(150_000), 250_000, 100_000), None);
+        assert_eq!(trim_floor(Some(150_000), 250_001, 100_000), Some(150_001));
+    }
+}