Add wshim support (cloudflare#167)

Author: mendess

Cargo.lock

@@ -25,7 +25,7 @@ impl DurableObject for Batcher {
             enable_dedup: params.enable_dedup,
             location_hint: params.location_hint.clone(),
         };
-        Batcher(GenericBatcher::new(env, config))
+        Batcher(GenericBatcher::new(state, env, config))
     }
 
     async fn fetch(&self, req: Request) -> Result<Response> {

crates/ct_worker/src/batcher_do.rs

@@ -35,7 +35,7 @@ impl DurableObject for Cleaner {
             clean_interval: Duration::from_secs(params.clean_interval_secs),
         };
 
-        Cleaner(GenericCleaner::new(&state, &env, config))
+        Cleaner(GenericCleaner::new(state, &env, config))
     }
 
     async fn fetch(&self, req: Request) -> Result<Response> {

crates/ct_worker/src/cleaner_do.rs

@@ -7,8 +7,8 @@ use crate::{load_roots, load_signing_key, SequenceMetadata, CONFIG};
 use config::TemporalInterval;
 use generic_log_worker::{
     batcher_id_from_lookup_key, deserialize, get_cached_metadata, get_durable_object_stub,
-    init_logging, load_cache_kv, load_public_bucket, put_cache_entry_metadata, serialize,
-    ObjectBucket, ENTRY_ENDPOINT, METRICS_ENDPOINT,
+    init_logging, load_cache_kv, load_public_bucket, obs::Wshim, put_cache_entry_metadata,
+    serialize, ObjectBucket, ENTRY_ENDPOINT,
 };
 use p256::pkcs8::EncodePublicKey;
 use serde::Serialize;
@@ -62,9 +62,10 @@ fn start() {
 ///
 /// Panics if there are issues parsing route parameters, which should never happen.
 #[event(fetch, respond_with_errors)]
-async fn main(req: Request, env: Env, _ctx: Context) -> Result<Response> {
+async fn main(req: Request, env: Env, ctx: Context) -> Result<Response> {
+    let wshim = Wshim::from_env(&env);
     // Use an outer router as middleware to check that the log name is valid.
-    Router::new()
+    let response = Router::new()
         .or_else_any_method_async("/logs/:log/*route", |req, ctx| async move {
             let name = if let Some(name) = ctx.param("log") {
                 if CONFIG.logs.contains_key(name) {
@@ -117,18 +118,6 @@ async fn main(req: Request, env: Env, _ctx: Context) -> Result<Response> {
                         temporal_interval: &params.temporal_interval,
                     })
                 })
-                .get_async("/logs/:log/metrics", |_req, ctx| async move {
-                    let name = ctx.data;
-                    let stub = get_durable_object_stub(
-                        &ctx.env,
-                        name,
-                        None,
-                        "SEQUENCER",
-                        CONFIG.logs[name].location_hint.as_deref(),
-                    )?;
-                    stub.fetch_with_str(&format!("http://fake_url.com{METRICS_ENDPOINT}"))
-                        .await
-                })
                 .get("/logs/:log/sequencer_id", |_req, ctx| {
                     // Print out the Durable Object ID of the sequencer to allow
                     // looking it up in internal Cloudflare dashboards. This
@@ -181,7 +170,11 @@ async fn main(req: Request, env: Env, _ctx: Context) -> Result<Response> {
                 log::warn!("Internal error: {e}");
                 Response::error("Internal error", 500)
             }
-        })
+        });
+    if let Ok(wshim) = wshim {
+        ctx.wait_until(async move { wshim.flush(&generic_log_worker::obs::logs::LOGGER).await });
+    }
+    response
 }
 
 #[allow(clippy::too_many_lines)]

crates/ct_worker/src/frontend_worker.rs

@@ -17,7 +17,7 @@
             "build": {
                 // Change '--release' to '--dev' to compile with debug symbols.
                 // DEPLOY_ENV is used in build.rs to select per-environment config and roots.
-                "command": "cargo install -q worker-build && DEPLOY_ENV=dev worker-build --release"
+                "command": "cargo install -q worker-build@0.1.14 && DEPLOY_ENV=dev worker-build --release"
             },
             "route": {
                 "pattern": "static-ct-dev.cloudflareresearch.com",
@@ -117,7 +117,7 @@
             "build": {
                 // Change '--release' to '--dev' to compile with debug symbols.
                 // DEPLOY_ENV is used in build.rs to select per-environment config and roots.
-                "command": "cargo install -q worker-build && DEPLOY_ENV=cftest worker-build --release"
+                "command": "cargo install -q worker-build@0.1.14 && DEPLOY_ENV=cftest worker-build --release"
             },
             "route": {
                 "pattern": "static-ct.cloudflareresearch.com",
@@ -185,7 +185,7 @@
             "build": {
                 // Change '--release' to '--dev' to compile with debug symbols.
                 // DEPLOY_ENV is used in build.rs to select per-environment config and roots.
-                "command": "cargo install -q worker-build && DEPLOY_ENV=raio worker-build --release"
+                "command": "cargo install -q worker-build@0.1.14 && DEPLOY_ENV=raio worker-build --release"
             },
             "routes": [
                 "https://ct.cloudflare.com/logs/raio2025h2b/*",

crates/ct_worker/wrangler.jsonc

@@ -29,13 +29,14 @@ console_log.workspace = true
 ed25519-dalek.workspace = true
 futures-util.workspace = true
 hex.workspace = true
-log.workspace = true
+log = { workspace = true, features = ["kv"] }
 p256.workspace = true
 prometheus.workspace = true
 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

@@ -7,8 +7,8 @@
 //! Entries are assigned to Batcher shards with consistent hashing on the cache key.
 
 use crate::{
-    deserialize, get_durable_object_stub, load_cache_kv, serialize, LookupKey, SequenceMetadata,
-    BATCH_ENDPOINT, ENTRY_ENDPOINT, SEQUENCER_BINDING,
+    deserialize, get_durable_object_stub, load_cache_kv, obs, serialize, LookupKey,
+    SequenceMetadata, BATCH_ENDPOINT, ENTRY_ENDPOINT, SEQUENCER_BINDING,
 };
 use base64::prelude::*;
 use futures_util::future::{join_all, select, Either};
@@ -26,10 +26,12 @@ use worker::*;
 pub struct GenericBatcher {
     env: Env,
     config: BatcherConfig,
+    state: State,
     kv: Option<KvStore>,
     batch: RefCell<Batch>,
     in_flight: RefCell<usize>,
     processed: RefCell<usize>,
+    wshim: Option<obs::Wshim>,
 }
 
 pub struct BatcherConfig {
@@ -64,20 +66,24 @@ impl GenericBatcher {
     /// # Panics
     ///
     /// Panics if we can't get a handle for the sequencer or KV store.
-    pub fn new(env: Env, config: BatcherConfig) -> Self {
+    pub fn new(state: State, env: Env, config: BatcherConfig) -> Self {
         let kv = if config.enable_dedup {
             Some(load_cache_kv(&env, &config.name).unwrap())
         } else {
             None
         };
 
+        let wshim = obs::Wshim::from_env(&env).ok();
+
         Self {
             env,
             config,
+            state,
             kv,
             batch: RefCell::new(Batch::default()),
             in_flight: RefCell::new(0),
             processed: RefCell::new(0),
+            wshim,
         }
     }
 
@@ -89,7 +95,11 @@ impl GenericBatcher {
     ///
     /// Returns an error if the request cannot be parsed or response cannot be
     /// constructed.
-    pub async fn fetch(&self, mut req: Request) -> Result<Response> {
+    pub async fn fetch(&self, req: Request) -> Result<Response> {
+        self.with_obs(async || self.fetch_impl(req).await).await
+    }
+
+    async fn fetch_impl(&self, mut req: Request) -> Result<Response> {
         match req.path().as_str() {
             ENTRY_ENDPOINT => {
                 let entry: PendingLogEntryBlob = deserialize(&req.bytes().await?)?;
@@ -163,6 +173,19 @@ impl GenericBatcher {
             _ => Response::error("not found", 404),
         }
     }
+
+    async fn with_obs<F, R>(&self, f: F) -> R
+    where
+        F: AsyncFnOnce() -> R,
+    {
+        let r = f().await;
+        if let Some(wshim) = self.wshim.clone() {
+            self.state.wait_until(async move {
+                wshim.flush(&obs::logs::LOGGER).await;
+            });
+        }
+        r
+    }
 }
 
 impl GenericBatcher {

crates/generic_log_worker/src/batcher_do.rs

@@ -9,7 +9,7 @@ use std::{cell::RefCell, mem, time::Duration};
 use tlog_tiles::{PathElem, TlogTile};
 use worker::{Bucket, Env, Error as WorkerError, Object, Request, Response, State, Storage};
 
-use crate::{load_public_bucket, log_ops::CHECKPOINT_KEY, util::now_millis};
+use crate::{load_public_bucket, log_ops::CHECKPOINT_KEY, obs, util::now_millis};
 
 // Workers are limited to 1000 subrequests per invocation (including R2 operations).
 // For each log, we'll need to perform the following subrequests:
@@ -36,13 +36,15 @@ pub struct CleanerConfig {
 }
 
 pub struct GenericCleaner {
+    state: State,
     config: CleanerConfig,
     storage: Storage,
     bucket: Bucket,
     cleaned_size: RefCell<u64>,
     current_size: RefCell<u64>,
     subrequests: RefCell<usize>,
     initialized: RefCell<bool>,
+    wshim: Option<obs::Wshim>,
 }
 
 impl GenericCleaner {
@@ -51,16 +53,19 @@ impl GenericCleaner {
     /// # Panics
     ///
     /// Panics if we can't get a handle for the public bucket.
-    pub fn new(state: &State, env: &Env, config: CleanerConfig) -> Self {
+    pub fn new(state: State, env: &Env, config: CleanerConfig) -> Self {
         let bucket = load_public_bucket(env, &config.name).unwrap();
+        let wshim = obs::Wshim::from_env(env).ok();
         Self {
             storage: state.storage(),
+            state,
             config,
             bucket,
             cleaned_size: RefCell::new(0),
             current_size: RefCell::new(0),
             subrequests: RefCell::new(0),
             initialized: RefCell::new(false),
+            wshim,
         }
     }
 
@@ -96,10 +101,13 @@ impl GenericCleaner {
     /// # Errors
     /// Will return an error if initialization fails.
     pub async fn fetch(&self, _req: Request) -> Result<Response, WorkerError> {
-        if !*self.initialized.borrow() {
-            self.initialize().await?;
-        }
-        Response::ok("Started cleaner")
+        self.with_obs(async || {
+            if !*self.initialized.borrow() {
+                self.initialize().await?;
+            }
+            Response::ok("Started cleaner")
+        })
+        .await
     }
 
     /// Alarm handler for the partial tile cleaner. This runs in a loop
@@ -109,6 +117,10 @@ impl GenericCleaner {
     /// # Errors
     /// Will return an error if initialization or cleaning fails.
     pub async fn alarm(&self) -> Result<Response, WorkerError> {
+        self.with_obs(async || self.alarm_impl().await).await
+    }
+
+    async fn alarm_impl(&self) -> Result<Response, WorkerError> {
         // Reset the subrequest count.
         *self.subrequests.borrow_mut() = 0;
 
@@ -293,4 +305,17 @@ impl GenericCleaner {
         *self.subrequests.borrow_mut() += new;
         Ok(())
     }
+
+    async fn with_obs<F, R>(&self, f: F) -> R
+    where
+        F: AsyncFnOnce() -> R,
+    {
+        let r = f().await;
+        if let Some(wshim) = self.wshim.clone() {
+            self.state.wait_until(async move {
+                wshim.flush(&obs::logs::LOGGER).await;
+            });
+        }
+        r
+    }
 }