TODO split into two commits

- Add PendingLogEntry struct to avoid overloading LogEntry.
  PendingLogEntry contains all of the fields of LogEntry except for the
  leaf index and timestamp.
- Instead of using a Sender per pool of entries, use one Sender per
  pending entry. This will allow us to split up the set of pending
  entries more easily, for example if we want to prefer creating full
  tiles instead of partial tiles (e.g., pick entries so that the log
  size is a multiple of 256, the full tile width).

Author: lukevalenta

crates/ct_worker/src/batcher_do.rs

@@ -9,7 +9,7 @@
 use crate::{get_stub, load_cache_kv, LookupKey, QueryParams, SequenceMetadata};
 use base64::prelude::*;
 use futures_util::future::{join_all, select, Either};
-use static_ct_api::LogEntry;
+use static_ct_api::PendingLogEntry;
 use std::{
     collections::{HashMap, HashSet},
     time::Duration,
@@ -37,7 +37,7 @@ struct Batcher {
 
 // A batch of entries to be submitted to the Sequencer together.
 struct Batch {
-    pending_leaves: Vec<LogEntry>,
+    pending_leaves: Vec<PendingLogEntry>,
     by_hash: HashSet<LookupKey>,
     done: Sender<HashMap<LookupKey, SequenceMetadata>>,
 }
@@ -68,7 +68,7 @@ impl DurableObject for Batcher {
         match req.path().as_str() {
             "/add_leaf" => {
                 let name = &req.query::<QueryParams>()?.name;
-                let entry: LogEntry = req.json().await?;
+                let entry: PendingLogEntry = req.json().await?;
                 let key = entry.lookup_key();
 
                 if self.in_flight >= MAX_IN_FLIGHT {

crates/ct_worker/src/ctlog.rs

@@ -30,7 +30,7 @@ use log::{debug, error, info, trace, warn};
 use p256::ecdsa::SigningKey as EcdsaSigningKey;
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
-use static_ct_api::{LogEntry, TileIterator, TreeWithTimestamp};
+use static_ct_api::{LogEntry, PendingLogEntry, TileIterator, TreeWithTimestamp};
 use std::collections::HashMap;
 use std::time::Duration;
 use std::{
@@ -39,7 +39,7 @@ use std::{
 };
 use thiserror::Error;
 use tlog_tiles::{Hash, HashReader, PathElem, Tile, TlogError, TlogTile, HASH_SIZE};
-use tokio::sync::watch::{Receiver, Sender};
+use tokio::sync::watch::{channel, Receiver, Sender};
 
 /// The maximum tile level is 63 (<c2sp.org/static-ct-api>), so safe to use [`u8::MAX`] as
 /// the special level for data tiles. The Go implementation uses -1.
@@ -75,65 +75,60 @@ pub(crate) struct LogConfig {
 /// <https://blog.cloudflare.com/durable-objects-easy-fast-correct-choose-three/#background-durable-objects-are-single-threaded>
 #[derive(Default, Debug)]
 pub(crate) struct PoolState {
-    // Entries that are ready to be sequenced.
-    pending_entries: Vec<LogEntry>,
+    // Entries that are ready to be sequenced, along with the Sender that will
+    // be updated when the entry is sequenced.
+    pending_entries: Vec<PendingLogEntry>,
 
-    // Deduplication cache for entries currently pending sequencing.
-    pending: HashMap<LookupKey, (u64, Receiver<SequenceMetadata>)>,
+    // Callbacks for the pending entries.
+    pending_senders: Vec<Sender<SequenceMetadata>>,
 
-    // Channel that will be updated when the current pool of pending entries is
-    // sequenced.
-    pending_done: Sender<SequenceMetadata>,
+    // Deduplication cache for entries currently pending sequencing.
+    pending: HashMap<LookupKey, Receiver<SequenceMetadata>>,
 
     // Deduplication cache for entries currently being sequenced.
-    in_sequencing: HashMap<LookupKey, (u64, Receiver<SequenceMetadata>)>,
+    in_sequencing: HashMap<LookupKey, Receiver<SequenceMetadata>>,
 }
 
 impl PoolState {
     // Check if the key is already in the pool. If so, return the index of the
     // entry in the pool, and a Receiver from which to read the entry metadata
     // when it is sequenced.
     fn check(&self, key: &LookupKey) -> Option<AddLeafResult> {
-        if let Some((index, rx)) = self.in_sequencing.get(key) {
+        if let Some(rx) = self.in_sequencing.get(key) {
             // Entry is being sequenced.
             Some(AddLeafResult::Pending {
-                pool_index: *index,
                 rx: rx.clone(),
                 source: PendingSource::InSequencing,
             })
         } else {
-            self.pending
-                .get(key)
-                .map(|(index, rx)| AddLeafResult::Pending {
-                    pool_index: *index,
-                    rx: rx.clone(),
-                    source: PendingSource::Pool,
-                })
+            self.pending.get(key).map(|rx| AddLeafResult::Pending {
+                rx: rx.clone(),
+                source: PendingSource::Pool,
+            })
         }
     }
     // Add a new entry to the pool.
-    fn add(&mut self, key: LookupKey, leaf: &LogEntry) -> AddLeafResult {
+    fn add(&mut self, key: LookupKey, entry: &PendingLogEntry) -> AddLeafResult {
         if self.pending_entries.len() >= MAX_POOL_SIZE {
             return AddLeafResult::RateLimited;
         }
-        self.pending_entries.push(leaf.clone());
-        let pool_index = (self.pending_entries.len() as u64) - 1;
-        let rx = self.pending_done.subscribe();
-        self.pending.insert(key, (pool_index, rx.clone()));
+        let (tx, rx) = channel((0, 0));
+        self.pending_entries.push(entry.clone());
+        self.pending_senders.push(tx);
+        self.pending.insert(key, rx.clone());
 
         AddLeafResult::Pending {
-            pool_index,
             rx,
             source: PendingSource::Sequencer,
         }
     }
     // Take the entries from the pool that are ready to be sequenced and the
     // corresponding Senders to update when the entries have been sequenced.
-    fn take(&mut self) -> (Vec<LogEntry>, Sender<SequenceMetadata>) {
+    fn take(&mut self) -> (Vec<PendingLogEntry>, Vec<Sender<SequenceMetadata>>) {
         self.in_sequencing = std::mem::take(&mut self.pending);
         (
             std::mem::take(&mut self.pending_entries),
-            std::mem::take(&mut self.pending_done),
+            std::mem::take(&mut self.pending_senders),
         )
     }
     // Reset the map of in-sequencing entries. This should be called after
@@ -348,7 +343,6 @@ impl SequenceState {
 pub(crate) enum AddLeafResult {
     Cached(SequenceMetadata),
     Pending {
-        pool_index: u64,
         rx: Receiver<SequenceMetadata>,
         source: PendingSource,
     },
@@ -361,15 +355,10 @@ impl AddLeafResult {
     pub(crate) async fn resolve(self) -> Option<SequenceMetadata> {
         match self {
             AddLeafResult::Cached(entry) => Some(entry),
-            AddLeafResult::Pending {
-                pool_index,
-                mut rx,
-                source: _,
-            } => {
+            AddLeafResult::Pending { mut rx, source: _ } => {
                 // Wait until sequencing completes for this entry's pool.
                 if rx.changed().await.is_ok() {
-                    let (first_index, timestamp) = *rx.borrow();
-                    Some((first_index + pool_index, timestamp))
+                    Some(*rx.borrow())
                 } else {
                     warn!("sender dropped");
                     None
@@ -383,11 +372,7 @@ impl AddLeafResult {
         match self {
             AddLeafResult::Cached(_) => "cache",
             AddLeafResult::RateLimited => "ratelimit",
-            AddLeafResult::Pending {
-                pool_index: _,
-                rx: _,
-                source,
-            } => match source {
+            AddLeafResult::Pending { rx: _, source } => match source {
                 PendingSource::InSequencing => "sequencing",
                 PendingSource::Pool => "pool",
                 PendingSource::Sequencer => "sequencer",
@@ -410,9 +395,9 @@ pub(crate) enum PendingSource {
 pub(crate) fn add_leaf_to_pool(
     state: &mut PoolState,
     cache: &impl CacheRead,
-    leaf: &LogEntry,
+    entry: &PendingLogEntry,
 ) -> AddLeafResult {
-    let hash = leaf.lookup_key();
+    let hash = entry.lookup_key();
 
     if let Some(result) = state.check(&hash) {
         // Entry is already pending or being sequenced.
@@ -422,7 +407,7 @@ pub(crate) fn add_leaf_to_pool(
         AddLeafResult::Cached(v)
     } else {
         // This is a new entry. Add it to the pool.
-        state.add(hash, leaf)
+        state.add(hash, entry)
     }
 }
 
@@ -477,7 +462,7 @@ pub(crate) async fn sequence(
     cache: &mut impl CacheWrite,
     metrics: &Metrics,
 ) -> Result<(), anyhow::Error> {
-    let (pending_entries, pending_done) = pool_state.take();
+    let (pending_entries, callbacks) = pool_state.take();
 
     metrics
         .seq_pool_size
@@ -490,7 +475,7 @@ pub(crate) async fn sequence(
         lock,
         cache,
         pending_entries,
-        pending_done,
+        callbacks,
         metrics,
     )
     .await
@@ -543,8 +528,8 @@ async fn sequence_entries(
     object: &impl ObjectBackend,
     lock: &impl LockBackend,
     cache: &mut impl CacheWrite,
-    mut pending_entries: Vec<LogEntry>,
-    done: Sender<SequenceMetadata>,
+    pending_entries: Vec<PendingLogEntry>,
+    callbacks: Vec<Sender<SequenceMetadata>>,
     metrics: &Metrics,
 ) -> Result<(), SequenceError> {
     let start = now_millis();
@@ -572,13 +557,16 @@ async fn sequence_entries(
     }
     let mut overlay = HashMap::new();
     let mut n = old_size;
-    let mut sequenced_leaves: Vec<LogEntry> = Vec::new();
+    let mut sequenced_entries: Vec<LogEntry> = Vec::new();
 
-    for leaf in &mut pending_entries {
-        leaf.leaf_index = n;
-        leaf.timestamp = timestamp;
-        sequenced_leaves.push(leaf.clone());
-        let tile_leaf = leaf.tile_leaf();
+    for entry in pending_entries {
+        let sequenced_entry = LogEntry {
+            inner: entry,
+            leaf_index: n,
+            timestamp,
+        };
+        let tile_leaf = sequenced_entry.tile_leaf();
+        let merkle_tree_leaf = sequenced_entry.merkle_tree_leaf();
         metrics.seq_leaf_size.observe(tile_leaf.len().as_f64());
         data_tile.extend(tile_leaf);
 
@@ -587,15 +575,15 @@ async fn sequence_entries(
         // the new tiles).
         let hashes = tlog_tiles::stored_hashes(
             n,
-            &leaf.merkle_tree_leaf(),
+            &merkle_tree_leaf,
             &HashReaderWithOverlay {
                 edge_tiles: &edge_tiles,
                 overlay: &overlay,
             },
         )
         .map_err(|e| {
             SequenceError::NonFatal(format!(
-                "couldn't compute new hashes for leaf {leaf:?}: {e}"
+                "couldn't compute new hashes for leaf {sequenced_entry:?}: {e}",
             ))
         })?;
         for (i, h) in hashes.iter().enumerate() {
@@ -611,6 +599,8 @@ async fn sequence_entries(
             metrics.seq_data_tile_size.observe(data_tile.len().as_f64());
             data_tile.clear();
         }
+
+        sequenced_entries.push(sequenced_entry);
     }
 
     // Stage leftover partial data tile, if any.
@@ -731,24 +721,31 @@ async fn sequence_entries(
 
     // Return SCTs to clients. Clients can recover the leaf index
     // from the old tree size and their index in the sequenced pool.
-    done.send_replace((old_size, timestamp));
+    for (pool_index, entry) in sequenced_entries.iter().enumerate() {
+        callbacks[pool_index].send_replace((entry.leaf_index, entry.timestamp));
+    }
 
     // At this point if the cache put fails, there's no reason to return errors to users. The
     // only consequence of cache false negatives are duplicated leaves anyway. In fact, an
     // error might cause the clients to resubmit, producing more cache false negatives and
     // duplicates.
     if let Err(e) = cache
         .put_entries(
-            &sequenced_leaves
+            &sequenced_entries
                 .iter()
-                .map(|entry| (entry.lookup_key(), (entry.leaf_index, entry.timestamp)))
+                .map(|entry| {
+                    (
+                        entry.inner.lookup_key(),
+                        (entry.leaf_index, entry.timestamp),
+                    )
+                })
                 .collect::<Vec<_>>(),
         )
         .await
     {
         warn!(
             "{name}: Cache put failed (entries={}): {e}",
-            sequenced_leaves.len()
+            sequenced_entries.len()
         );
     }
 
@@ -1130,7 +1127,7 @@ mod tests {
         for i in 0..500_u64 {
             for k in 0..3000_u64 {
                 let certificate = (i * 3000 + k).to_be_bytes().to_vec();
-                let leaf = LogEntry {
+                let leaf = PendingLogEntry {
                     certificate,
                     ..Default::default()
                 };
@@ -1305,9 +1302,9 @@ mod tests {
 
         // A pair of duplicates from the in_sequencing pool.
         let res21 = log.add_with_seed(is_precert, 2); // 7
-        let (pending_entries, pending_done) = log.pool_state.take();
+        let (pending_entries, callbacks) = log.pool_state.take();
         let res22 = log.add_with_seed(is_precert, 2);
-        log.sequence_finish(pending_entries, pending_done);
+        log.sequence_finish(pending_entries, callbacks);
         let entry21 = block_on(res21.resolve()).unwrap();
         let entry22 = block_on(res22.resolve()).unwrap();
         assert_eq!(entry21, entry22);
@@ -1880,8 +1877,8 @@ mod tests {
         }
         fn sequence_finish(
             &mut self,
-            pending_entries: Vec<LogEntry>,
-            pending_done: Sender<SequenceMetadata>,
+            pending_entries: Vec<PendingLogEntry>,
+            callbacks: Vec<Sender<SequenceMetadata>>,
         ) {
             block_on(sequence_entries(
                 &mut self.sequence_state,
@@ -1890,7 +1887,7 @@ mod tests {
                 &self.lock,
                 &mut self.cache,
                 pending_entries,
-                pending_done,
+                callbacks,
                 &self.metrics,
             ))
             .unwrap();
@@ -1919,14 +1916,12 @@ mod tests {
                 pre_certificate = Vec::new();
             }
             let issuers = CHAINS[rng.gen_range(0..CHAINS.len())];
-            let leaf = LogEntry {
+            let leaf = PendingLogEntry {
                 certificate,
                 pre_certificate,
                 is_precert,
                 issuer_key_hash,
                 chain_fingerprints: issuers.iter().map(|&x| Sha256::digest(x).into()).collect(),
-                leaf_index: 0,
-                timestamp: 0,
             };
 
             block_on(upload_issuers(&self.object, issuers, &self.config.name)).unwrap();
@@ -2020,16 +2015,16 @@ mod tests {
                         tlog_tiles::record_hash(&entry.merkle_tree_leaf())
                     );
 
-                    assert!(!entry.certificate.is_empty());
-                    if entry.is_precert {
-                        assert!(!entry.pre_certificate.is_empty());
-                        assert_ne!(entry.issuer_key_hash, [0; 32]);
+                    assert!(!entry.inner.certificate.is_empty());
+                    if entry.inner.is_precert {
+                        assert!(!entry.inner.pre_certificate.is_empty());
+                        assert_ne!(entry.inner.issuer_key_hash, [0; 32]);
                     } else {
-                        assert!(entry.pre_certificate.is_empty());
-                        assert_eq!(entry.issuer_key_hash, [0; 32]);
+                        assert!(entry.inner.pre_certificate.is_empty());
+                        assert_eq!(entry.inner.issuer_key_hash, [0; 32]);
                     }
 
-                    for fp in entry.chain_fingerprints {
+                    for fp in entry.inner.chain_fingerprints {
                         let b = block_on(self.object.fetch(&format!("issuer/{}", hex::encode(fp))))
                             .unwrap()
                             .unwrap();