Use bitcode instead of json serde, and remove generics from batcher (#68)

- Use bitcode instead of json for more efficient serialization for
  communication across Workers and DOs. This is an intermediate step
  short of full RPC support (#3).
- Use bitcode instead of json for serializing the staging bundle.
- Introduce `PendingLogEntryBlob` struct that wraps the serialized bytes
  and lookup key of a `PendingLogEntry` so that the Batcher doesn't need
  to fully deserialize entries as all it needs is the lookup key. This
  allows us to remote the generic type from the Batcher altogether
  (fixes #67).

Other:
- Pass argument to durable_object macro to reduce generated code size
  (Batchers only need fetch methods, and Sequencers need fetch and
  alarm, but not websocket.)
- Clean up code in various places.

Author: lukevalenta

Cargo.lock

@@ -32,6 +32,7 @@ debug = 1
 [workspace.dependencies]
 anyhow = "1.0"
 base64 = "0.22"
+bitcode = { version = "0.6.6", features = ["serde"] }
 byteorder = "1.5"
 chrono = { version = "0.4", features = ["serde"] }
 console_error_panic_hook = "0.1.1"

Cargo.toml

@@ -5,7 +5,6 @@
 
 use chrono::Months;
 use config::AppConfig;
-use serde_json::from_str;
 use std::env;
 use std::fs;
 use url::Url;
@@ -19,10 +18,10 @@ fn main() {
     });
 
     // Validate the config json against the schema.
-    let json = from_str(config_contents).unwrap_or_else(|e| {
+    let json = serde_json::from_str(config_contents).unwrap_or_else(|e| {
         panic!("failed to deserialize JSON config '{config_file}': {e}");
     });
-    let schema = from_str(include_str!("config.schema.json")).unwrap_or_else(|e| {
+    let schema = serde_json::from_str(include_str!("config.schema.json")).unwrap_or_else(|e| {
         panic!("failed to deserialize JSON schema 'config.schema.json': {e}");
     });
     jsonschema::validate(&schema, &json).unwrap_or_else(|e| {

crates/ct_worker/build.rs

@@ -1,11 +1,10 @@
 use crate::CONFIG;
 use generic_log_worker::{get_durable_object_stub, load_cache_kv, BatcherConfig, GenericBatcher};
-use static_ct_api::StaticCTPendingLogEntry;
 #[allow(clippy::wildcard_imports)]
 use worker::*;
 
-#[durable_object]
-struct Batcher(GenericBatcher<StaticCTPendingLogEntry>);
+#[durable_object(fetch)]
+struct Batcher(GenericBatcher);
 
 impl DurableObject for Batcher {
     fn new(state: State, env: Env) -> Self {

crates/ct_worker/src/batcher_do.rs

@@ -5,12 +5,13 @@
 
 use std::sync::LazyLock;
 
-use crate::{load_signing_key, LookupKey, SequenceMetadata, CONFIG, ROOTS};
+use crate::{load_signing_key, SequenceMetadata, CONFIG, ROOTS};
 use config::TemporalInterval;
 use futures_util::future::try_join_all;
 use generic_log_worker::{
-    get_cached_metadata, get_durable_object_stub, init_logging, load_cache_kv, load_public_bucket,
-    log_ops::UploadOptions, put_cache_entry_metadata, ObjectBackend, ObjectBucket, ENTRY_ENDPOINT,
+    batcher_id_from_lookup_key, deserialize, get_cached_metadata, get_durable_object_stub,
+    init_logging, load_cache_kv, load_public_bucket, log_ops::UploadOptions,
+    put_cache_entry_metadata, serialize, ObjectBackend, ObjectBucket, ENTRY_ENDPOINT,
     METRICS_ENDPOINT,
 };
 use log::{debug, info, warn};
@@ -19,7 +20,7 @@ use serde::Serialize;
 use serde_with::{base64::Base64, serde_as};
 use sha2::{Digest, Sha256};
 use static_ct_api::{AddChainRequest, GetRootsResponse, StaticCTLogEntry};
-use tlog_tiles::{LogEntry, PendingLogEntry};
+use tlog_tiles::{LogEntry, PendingLogEntry, PendingLogEntryBlob};
 #[allow(clippy::wildcard_imports)]
 use worker::*;
 
@@ -176,7 +177,6 @@ async fn main(req: Request, env: Env, _ctx: Context) -> Result<Response> {
         })
 }
 
-#[allow(clippy::too_many_lines)]
 async fn add_chain_or_pre_chain(
     mut req: Request,
     env: &Env,
@@ -215,8 +215,7 @@ async fn add_chain_or_pre_chain(
 
     // Check if entry is cached and return right away if so.
     if params.enable_dedup {
-        if let Some(metadata) =
-            get_cached_metadata(&load_cache_kv(env, name)?, &pending_entry).await?
+        if let Some(metadata) = get_cached_metadata(&load_cache_kv(env, name)?, &lookup_key).await?
         {
             debug!("{name}: Entry is cached");
             let entry = StaticCTLogEntry::new(pending_entry, metadata);
@@ -242,30 +241,31 @@ async fn add_chain_or_pre_chain(
 
     // Submit entry to be sequenced, either via a batcher or directly to the
     // sequencer.
-    let stub = if params.num_batchers > 0 {
-        let batcher_id = batcher_id_from_lookup_key(&lookup_key, params.num_batchers);
+    let stub = {
+        let shard_id = batcher_id_from_lookup_key(&lookup_key, params.num_batchers);
         get_durable_object_stub(
             env,
             name,
-            Some(batcher_id),
-            "BATCHER",
-            params.location_hint.as_deref(),
-        )?
-    } else {
-        get_durable_object_stub(
-            env,
-            name,
-            None,
-            "SEQUENCER",
+            shard_id,
+            if shard_id.is_some() {
+                "BATCHER"
+            } else {
+                "SEQUENCER"
+            },
             params.location_hint.as_deref(),
         )?
     };
+
+    let serialized = serialize(&PendingLogEntryBlob {
+        lookup_key,
+        data: serialize(&pending_entry)?,
+    })?;
     let mut response = stub
         .fetch_with_request(Request::new_with_init(
             &format!("http://fake_url.com{ENTRY_ENDPOINT}"),
             &RequestInit {
                 method: Method::Post,
-                body: Some(serde_json::to_string(&pending_entry)?.into()),
+                body: Some(serialized.into()),
                 ..Default::default()
             },
         )?)
@@ -274,7 +274,7 @@ async fn add_chain_or_pre_chain(
         // Return the response from the sequencing directly to the client.
         return Ok(response);
     }
-    let metadata = response.json::<SequenceMetadata>().await?;
+    let metadata = deserialize::<SequenceMetadata>(&response.bytes().await?)?;
     if params.num_batchers == 0 && params.enable_dedup {
         // Write sequenced entry to the long-term deduplication cache in Workers
         // KV as there are no batchers configured to do it for us.
@@ -295,10 +295,6 @@ async fn add_chain_or_pre_chain(
     Response::from_json(&sct)
 }
 
-fn batcher_id_from_lookup_key(key: &LookupKey, num_batchers: u8) -> u8 {
-    key[0] % num_batchers
-}
-
 fn headers_from_http_metadata(meta: HttpMetadata) -> Headers {
     let h = Headers::new();
     if let Some(hdr) = meta.cache_control {

crates/ct_worker/src/frontend_worker.rs

@@ -10,7 +10,7 @@ use signed_note::KeyName;
 use static_ct_api::StaticCTCheckpointSigner;
 use std::collections::HashMap;
 use std::sync::{LazyLock, OnceLock};
-use tlog_tiles::{CheckpointSigner, Ed25519CheckpointSigner, LookupKey, SequenceMetadata};
+use tlog_tiles::{CheckpointSigner, Ed25519CheckpointSigner, SequenceMetadata};
 #[allow(clippy::wildcard_imports)]
 use worker::*;
 use x509_cert::Certificate;

crates/ct_worker/src/lib.rs

@@ -12,7 +12,7 @@ use static_ct_api::StaticCTLogEntry;
 #[allow(clippy::wildcard_imports)]
 use worker::*;
 
-#[durable_object]
+#[durable_object(alarm)]
 struct Sequencer(GenericSequencer<StaticCTLogEntry>);
 
 impl DurableObject for Sequencer {

crates/ct_worker/src/sequencer_do.rs

@@ -22,6 +22,7 @@ static_ct_api.workspace = true
 [dependencies]
 anyhow.workspace = true
 base64.workspace = true
+bitcode.workspace = true
 byteorder.workspace = true
 console_error_panic_hook.workspace = true
 console_log.workspace = true
@@ -35,7 +36,6 @@ rand.workspace = true
 serde-wasm-bindgen.workspace = true
 serde.workspace = true
 serde_bytes.workspace = true
-serde_json.workspace = true
 sha2.workspace = true
 signed_note.workspace = true
 thiserror.workspace = true

crates/generic_log_worker/Cargo.toml

@@ -6,25 +6,25 @@
 //!
 //! Entries are assigned to Batcher shards with consistent hashing on the cache key.
 
-use crate::{LookupKey, SequenceMetadata, BATCH_ENDPOINT, ENTRY_ENDPOINT};
+use crate::{deserialize, serialize, LookupKey, SequenceMetadata, BATCH_ENDPOINT, ENTRY_ENDPOINT};
 use base64::prelude::*;
 use futures_util::future::{join_all, select, Either};
 use std::{
     cell::RefCell,
     collections::{HashMap, HashSet},
     time::Duration,
 };
-use tlog_tiles::PendingLogEntry;
+use tlog_tiles::PendingLogEntryBlob;
 use tokio::sync::watch::{self, Sender};
 use worker::kv::KvStore;
 #[allow(clippy::wildcard_imports)]
 use worker::*;
 
-pub struct GenericBatcher<E: PendingLogEntry> {
+pub struct GenericBatcher {
     config: BatcherConfig,
     kv: Option<KvStore>,
     sequencer: Stub,
-    batch: RefCell<Batch<E>>,
+    batch: RefCell<Batch>,
     in_flight: RefCell<usize>,
     processed: RefCell<usize>,
 }
@@ -36,13 +36,13 @@ pub struct BatcherConfig {
 }
 
 // A batch of entries to be submitted to the Sequencer together.
-struct Batch<E: PendingLogEntry> {
-    entries: Vec<E>,
+struct Batch {
+    entries: Vec<PendingLogEntryBlob>,
     by_hash: HashSet<LookupKey>,
     done: Sender<HashMap<LookupKey, SequenceMetadata>>,
 }
 
-impl<E: PendingLogEntry> Default for Batch<E> {
+impl Default for Batch {
     /// Returns a batch initialized with a watch channel.
     fn default() -> Self {
         Self {
@@ -53,7 +53,7 @@ impl<E: PendingLogEntry> Default for Batch<E> {
     }
 }
 
-impl<E: PendingLogEntry> GenericBatcher<E> {
+impl GenericBatcher {
     /// Returns a new batcher with the given config.
     pub fn new(config: BatcherConfig, kv: Option<KvStore>, sequencer: Stub) -> Self {
         Self {
@@ -77,8 +77,8 @@ impl<E: PendingLogEntry> GenericBatcher<E> {
     pub async fn fetch(&self, mut req: Request) -> Result<Response> {
         match req.path().as_str() {
             ENTRY_ENDPOINT => {
-                let entry: E = req.json().await?;
-                let key = entry.lookup_key();
+                let entry: PendingLogEntryBlob = deserialize(&req.bytes().await?)?;
+                let key = entry.lookup_key;
 
                 *self.in_flight.borrow_mut() += 1;
                 *self.processed.borrow_mut() += 1;
@@ -135,7 +135,7 @@ impl<E: PendingLogEntry> GenericBatcher<E> {
 
                 let resp = if let Some(value) = recv.borrow().get(&key) {
                     // The entry has been sequenced!
-                    Response::from_json(&value)
+                    Response::from_bytes(serialize(&value)?)
                 } else {
                     // Failed to sequence this entry, either due to an error
                     // submitting the batch or rate limiting at the Sequencer.
@@ -150,7 +150,7 @@ impl<E: PendingLogEntry> GenericBatcher<E> {
     }
 }
 
-impl<E: PendingLogEntry> GenericBatcher<E> {
+impl GenericBatcher {
     /// Submit the current pending batch to be sequenced.
     ///
     /// # Errors
@@ -174,16 +174,19 @@ impl<E: PendingLogEntry> GenericBatcher<E> {
             &format!("http://fake_url.com{BATCH_ENDPOINT}"),
             &RequestInit {
                 method: Method::Post,
-                body: Some(serde_json::to_string(&batch.entries)?.into()),
+                body: Some(serialize(&batch.entries)?.into()),
                 ..Default::default()
             },
         )?;
-        let sequenced_entries: HashMap<LookupKey, SequenceMetadata> = self
-            .sequencer
-            .fetch_with_request(req)
-            .await?
-            .json::<Vec<(LookupKey, SequenceMetadata)>>()
-            .await?
+        let sequenced_entries: HashMap<LookupKey, SequenceMetadata> =
+            deserialize::<Vec<(LookupKey, SequenceMetadata)>>(
+                &self
+                    .sequencer
+                    .fetch_with_request(req)
+                    .await?
+                    .bytes()
+                    .await?,
+            )?
             .into_iter()
             .collect();
 
@@ -206,3 +209,12 @@ impl<E: PendingLogEntry> GenericBatcher<E> {
         Ok(())
     }
 }
+
+/// Return a batcher ID to which the provided entry should be assigned.
+pub fn batcher_id_from_lookup_key(key: &LookupKey, num_batchers: u8) -> Option<u8> {
+    if num_batchers > 0 {
+        Some(key[0] % num_batchers)
+    } else {
+        None
+    }
+}