MTC: Cache landmark bundle and optimize computation (#163)

* MTC: Implement caching for hot and cold calls to /get-landmark-bundle

* Started moving landmark bundle comptuation to cronjob

* Fix type signature of CheckpointCallbacker to let the landmark bundle creator access the checkpoint bytes

* Save the landmark bundle in R2

* Make landmark subtrees endpoint fetch from R2

* Return get-landmark-bundle directly from R2

* Implement cached bucket backend

* Add comments

---------

Co-authored-by: Michael Rosenberg <mrosenberg@cloudflare.com>

Author: rozbb

crates/generic_log_worker/src/lib.rs

@@ -25,11 +25,13 @@ use serde::{Deserialize, Serialize};
 use serde_bytes::ByteBuf;
 use sha2::{Digest, Sha256};
 use std::cell::RefCell;
-use std::collections::{HashMap, VecDeque};
+use std::collections::btree_map::Entry;
+use std::collections::{BTreeMap, HashMap, VecDeque};
 use std::io::Write;
 use std::str::FromStr;
 use std::sync::Once;
 use tlog_tiles::{LookupKey, PendingLogEntry, SequenceMetadata};
+use tokio::sync::Mutex;
 use util::now_millis;
 use worker::{
     js_sys, kv, kv::KvStore, wasm_bindgen, Bucket, Env, Error, HttpMetadata, Result, State,
@@ -628,3 +630,46 @@ impl ObjectBackend for ObjectBucket {
         res
     }
 }
+
+/// A read-only ObjectBucket that caches every fetch no matter how big
+///
+/// **NOTE:** The cache here has no size limit. If you use a `CachedRoObjectBucket` for too many
+/// fetches, you will run out of memory.
+pub struct CachedRoObjectBucket {
+    bucket: ObjectBucket,
+    cache: Mutex<BTreeMap<String, Option<Vec<u8>>>>,
+}
+
+impl CachedRoObjectBucket {
+    pub fn new(bucket: ObjectBucket) -> Self {
+        CachedRoObjectBucket {
+            bucket,
+            cache: Mutex::new(BTreeMap::new()),
+        }
+    }
+}
+
+impl ObjectBackend for CachedRoObjectBucket {
+    async fn upload<S: AsRef<str>, D: Into<Vec<u8>>>(
+        &self,
+        _key: S,
+        _data: D,
+        _opts: &UploadOptions,
+    ) -> Result<()> {
+        unimplemented!("CachedRoObjectBucket does not implement ObjectBackend::upload")
+    }
+
+    async fn fetch<S: AsRef<str>>(&self, key: S) -> Result<Option<Vec<u8>>> {
+        // See if the key is in the cache
+        match self.cache.blocking_lock().entry(key.as_ref().to_string()) {
+            // If so, return the value
+            Entry::Occupied(oentry) => Ok(oentry.get().clone()),
+            // Otherwise, fetch the value, cache it, and return it
+            Entry::Vacant(ventry) => {
+                let val = self.bucket.fetch(key).await?;
+                ventry.insert(val.clone());
+                Ok(val)
+            }
+        }
+    }
+}

crates/generic_log_worker/src/log_ops.rs

@@ -1061,7 +1061,7 @@ async fn sequence_entries<L: LogEntry>(
 
     // Call the checkpoint callback. This is a no-op for CT, but is used to
     // update landmark checkpoints for MTC.
-    if let Err(e) = (config.checkpoint_callback)(n, old_time, timestamp).await {
+    if let Err(e) = (config.checkpoint_callback)(old_time, timestamp, new.checkpoint()).await {
         warn!("{name}: Checkpoint callback failed: {e}");
     }
 

crates/generic_log_worker/src/sequencer_do.rs

@@ -307,14 +307,14 @@ impl<L: LogEntry> GenericSequencer<L> {
 /// so that it can have side-effects.
 ///
 /// The parameters are as follows:
-/// - `tree_size: u64`: The tree size of the latest checkpoint.
 /// - `old_time: UnixTimestamp`: The timestamp of the previous checkpoint.
 /// - `new_time: UnixTimestamp`: The timestamp of the latest checkpoint.
+/// - `new_checkpoint: &[u8]`: The latest checkpoint bytes. This is a signed note.
 pub type CheckpointCallbacker = Box<
     dyn Fn(
-            u64,
             UnixTimestamp,
             UnixTimestamp,
+            &[u8],
         ) -> Pin<Box<dyn Future<Output = Result<(), WorkerError>> + 'static>>
         + 'static,
 >;
@@ -323,9 +323,8 @@ pub type CheckpointCallbacker = Box<
 /// don't need to perform any action after the checkpoint is updated.
 pub fn empty_checkpoint_callback() -> CheckpointCallbacker {
     Box::new(
-        move |_tree_size: u64, _old_time: UnixTimestamp, _new_time: UnixTimestamp| {
+        move |_old_time: UnixTimestamp, _new_time: UnixTimestamp, _new_checkpoint: &[u8]| {
             Box::pin(async move { Ok(()) })
-                as Pin<Box<dyn Future<Output = Result<(), WorkerError>>>>
         },
     )
 }

crates/mtc_api/src/landmark.rs

@@ -9,8 +9,11 @@ pub struct LandmarkSequence {
     pub landmarks: VecDeque<u64>,
 }
 
-/// The location in object storage for the landmark bundle.
-pub static LANDMARK_KEY: &str = "landmark";
+/// The location in object storage for the landmark sequence
+pub const LANDMARK_KEY: &str = "landmark";
+
+/// The location in object storage for the landmark bundle. Its serialized form is JSON
+pub const LANDMARK_BUNDLE_KEY: &str = "landmark-bundle";
 
 impl LandmarkSequence {
     /// Create a new landmark sequence with the given `max_landmarks` and an

crates/mtc_worker/src/frontend_worker.rs

@@ -11,21 +11,19 @@ use der::{
 use generic_log_worker::{
     batcher_id_from_lookup_key, deserialize, get_durable_object_stub, init_logging,
     load_public_bucket,
-    log_ops::{
-        prove_subtree_consistency, prove_subtree_inclusion, read_leaf, ProofError, CHECKPOINT_KEY,
-    },
+    log_ops::{prove_subtree_inclusion, read_leaf, ProofError, CHECKPOINT_KEY},
     serialize,
     util::now_millis,
     ObjectBackend, ObjectBucket, ENTRY_ENDPOINT, METRICS_ENDPOINT,
 };
 use mtc_api::{
     serialize_signatureless_cert, AddEntryRequest, AddEntryResponse, BootstrapMtcLogEntry,
-    GetRootsResponse, LandmarkSequence, ID_RDNA_TRUSTANCHOR_ID, LANDMARK_KEY,
+    GetRootsResponse, LandmarkSequence, ID_RDNA_TRUSTANCHOR_ID, LANDMARK_BUNDLE_KEY, LANDMARK_KEY,
 };
 use serde::{Deserialize, Serialize};
 use serde_with::{base64::Base64, serde_as};
 use signed_note::VerifierList;
-use std::{collections::VecDeque, time::Duration};
+use std::time::Duration;
 use tlog_tiles::{
     open_checkpoint, CheckpointSigner, CheckpointText, LeafIndex, PendingLogEntry,
     PendingLogEntryBlob,
@@ -72,24 +70,6 @@ pub struct GetCertificateResponse {
     pub landmark_id: usize,
 }
 
-/// GET response structure for the `/get-landmark-bundle` endpoint
-#[serde_as]
-#[derive(Serialize)]
-pub struct GetLandmarkBundleResponse<'a> {
-    pub checkpoint: &'a str,
-    pub subtrees: Vec<SubtreeWithConsistencyProof>,
-    pub landmarks: &'a VecDeque<u64>,
-}
-
-#[serde_as]
-#[derive(Serialize, Deserialize)]
-pub struct SubtreeWithConsistencyProof {
-    #[serde_as(as = "Base64")]
-    pub hash: [u8; 32],
-    #[serde_as(as = "Vec<Base64>")]
-    pub consistency_proof: Vec<[u8; 32]>,
-}
-
 /// Start is the first code run when the Wasm module is loaded.
 #[event(start)]
 fn start() {
@@ -412,41 +392,14 @@ async fn add_entry(mut req: Request, env: &Env, name: &str) -> Result<Response>
 async fn get_landmark_bundle(env: &Env, name: &str) -> Result<Response> {
     let object_backend = ObjectBucket::new(load_public_bucket(env, name)?);
 
-    // Fetch the current checkpoint.
-    let (checkpoint, checkpoint_bytes) = get_current_checkpoint(env, name, &object_backend).await?;
-
-    // Fetch the current landmark sequence.
-    let landmark_sequence = get_landmark_sequence(name, &object_backend).await?;
-
-    // Compute the sequence of landmark subtrees and, for each subtree, a proof of consistency with
-    // the checkpoint. Each signatureless MTC includes an inclusion proof in one of these subtrees.
-    let mut subtrees = Vec::new();
-    for landmark_subtree in landmark_sequence.subtrees() {
-        let (consistency_proof, landmark_subtree_hash) = match prove_subtree_consistency(
-            *checkpoint.hash(),
-            checkpoint.size(),
-            landmark_subtree.lo(),
-            landmark_subtree.hi(),
-            &object_backend,
-        )
-        .await
-        {
-            Ok(p) => p,
-            Err(ProofError::Tlog(s)) => return Response::error(s.to_string(), 422),
-            Err(ProofError::Other(e)) => return Err(e.to_string().into()),
-        };
-
-        subtrees.push(SubtreeWithConsistencyProof {
-            hash: landmark_subtree_hash.0,
-            consistency_proof: consistency_proof.iter().map(|h| h.0).collect(),
-        });
-    }
+    // Fetch the current landmark bundle from R2 (already encoded in JSON) and return it
+    let Some(landmark_bundle_bytes) = object_backend.fetch(LANDMARK_BUNDLE_KEY).await? else {
+        return Err("failed to get landmark bundle".into());
+    };
 
-    Response::from_json(&GetLandmarkBundleResponse {
-        checkpoint: std::str::from_utf8(&checkpoint_bytes).unwrap(),
-        subtrees,
-        landmarks: &landmark_sequence.landmarks,
-    })
+    Ok(ResponseBuilder::new()
+        .with_header("content-type", "application/json")?
+        .body(ResponseBody::Body(landmark_bundle_bytes)))
 }
 
 async fn get_current_checkpoint(

crates/mtc_worker/src/sequencer_do.rs

@@ -3,15 +3,20 @@
 
 //! Sequencer is the 'brain' of the CT log, responsible for sequencing entries and maintaining log state.
 
-use std::{future::Future, pin::Pin, time::Duration};
+use std::{collections::VecDeque, time::Duration};
 
 use crate::{load_checkpoint_cosigner, load_origin, CONFIG};
 use generic_log_worker::{
-    get_durable_object_name, load_public_bucket, CheckpointCallbacker, GenericSequencer,
-    SequencerConfig, SEQUENCER_BINDING,
+    get_durable_object_name, load_public_bucket,
+    log_ops::{prove_subtree_consistency, ProofError},
+    CachedRoObjectBucket, CheckpointCallbacker, GenericSequencer, ObjectBucket, SequencerConfig,
+    SEQUENCER_BINDING,
 };
-use mtc_api::{BootstrapMtcLogEntry, LandmarkSequence, LANDMARK_KEY};
-use tlog_tiles::UnixTimestamp;
+use mtc_api::{BootstrapMtcLogEntry, LandmarkSequence, LANDMARK_BUNDLE_KEY, LANDMARK_KEY};
+use serde::{Deserialize, Serialize};
+use serde_with::{base64::Base64, serde_as};
+use signed_note::Note;
+use tlog_tiles::{CheckpointText, Hash, UnixTimestamp};
 #[allow(clippy::wildcard_imports)]
 use worker::*;
 
@@ -53,14 +58,47 @@ impl DurableObject for Sequencer {
     }
 }
 
+#[serde_as]
+#[derive(Serialize, Deserialize)]
+pub struct SubtreeWithConsistencyProof {
+    #[serde_as(as = "Base64")]
+    pub hash: [u8; 32],
+    #[serde_as(as = "Vec<Base64>")]
+    pub consistency_proof: Vec<[u8; 32]>,
+}
+
+/// GET response structure for the `/get-landmark-bundle` endpoint
+#[derive(Serialize, Deserialize)]
+pub struct LandmarkBundle {
+    pub checkpoint: String,
+    pub subtrees: Vec<SubtreeWithConsistencyProof>,
+    pub landmarks: VecDeque<u64>,
+}
+
 /// Return a callback function that gets passed into the generic sequencer and
 /// called each time a new checkpoint is created. For MTC, this is used to
 /// periodically update the landmark checkpoint sequence.
 fn checkpoint_callback(env: &Env, name: &str) -> CheckpointCallbacker {
     let params = &CONFIG.logs[name];
     let bucket = load_public_bucket(env, name).unwrap();
     Box::new(
-        move |tree_size: u64, old_time: UnixTimestamp, new_time: UnixTimestamp| {
+        move |old_time: UnixTimestamp, new_time: UnixTimestamp, new_checkpoint_bytes: &[u8]| {
+            let new_checkpoint = {
+                // TODO: Make more efficient. There are two unnecessary allocations here.
+
+                // We can unwrap because the checkpoint provided is the checkpoint that the
+                // sequencer just created, so it must be well formed.
+                let note = Note::from_bytes(new_checkpoint_bytes)
+                    .expect("freshly created checkpoint is not a note");
+                CheckpointText::from_bytes(note.text())
+                    .expect("freshly created checkpoint is not a checkpoint")
+            };
+            let tree_size = new_checkpoint.size();
+            let root_hash = *new_checkpoint.hash();
+            // We can unwrap here for the same reason as above
+            let new_checkpoint_str = String::from_utf8(new_checkpoint_bytes.to_vec())
+                .expect("freshly created checkpoint is not UTF-8");
+
             Box::pin({
                 // We have to clone each time since the bucket gets moved into
                 // the async function.
@@ -100,9 +138,63 @@ fn checkpoint_callback(env: &Env, name: &str) -> CheckpointCallbacker {
                             .execute()
                             .await?;
                     }
+
+                    // Compute the landmark bundle and save it
+                    let subtrees =
+                        get_landmark_subtrees(&seq, root_hash, tree_size, bucket_clone.clone())
+                            .await?;
+                    let bundle = LandmarkBundle {
+                        checkpoint: new_checkpoint_str,
+                        subtrees,
+                        landmarks: seq.landmarks,
+                    };
+                    // TODO: the put operation here should be done with the put operation above.
+                    // Otherwise an error here might put us in a state where the landmark bundle is
+                    // out of sync with the landmark sequence. We need an all-or-nothing multi-put
+                    // operation. Tracking issue here https://github.com/cloudflare/workers-rs/issues/876
+                    bucket_clone
+                        // Can unwrap here because we use the autoderived Serialize impl for LandmarkBundle
+                        .put(LANDMARK_BUNDLE_KEY, serde_json::to_vec(&bundle).unwrap())
+                        .execute()
+                        .await?;
+
                     Ok(())
                 }
-            }) as Pin<Box<dyn Future<Output = Result<()>>>>
+            })
         },
     )
 }
+
+// Computes the sequence of landmark subtrees and, for each subtree, a proof of consistency with the
+// checkpoint. Each signatureless MTC includes an inclusion proof in one of these subtrees.
+async fn get_landmark_subtrees(
+    landmark_sequence: &LandmarkSequence,
+    checkpoint_hash: Hash,
+    checkpoint_size: u64,
+    bucket: Bucket,
+) -> Result<Vec<SubtreeWithConsistencyProof>> {
+    let cached_object_backend = CachedRoObjectBucket::new(ObjectBucket::new(bucket));
+    let mut subtrees = Vec::new();
+    for landmark_subtree in landmark_sequence.subtrees() {
+        let (consistency_proof, landmark_subtree_hash) = match prove_subtree_consistency(
+            checkpoint_hash,
+            checkpoint_size,
+            landmark_subtree.lo(),
+            landmark_subtree.hi(),
+            &cached_object_backend,
+        )
+        .await
+        {
+            Ok(p) => p,
+            Err(ProofError::Tlog(s)) => return Err(s.to_string().into()),
+            Err(ProofError::Other(e)) => return Err(e.to_string().into()),
+        };
+
+        subtrees.push(SubtreeWithConsistencyProof {
+            hash: landmark_subtree_hash.0,
+            consistency_proof: consistency_proof.iter().map(|h| h.0).collect(),
+        });
+    }
+
+    Ok(subtrees)
+}