Implement the async worker

Author: mendess

Cargo.lock

@@ -90,8 +90,9 @@ tracing = "0.1.40"
 tracing-core = "0.1.32"
 tracing-subscriber = "0.3.18"
 url = { version = "2.5.4", features = ["serde"] }
+wasm-bindgen = "0.2.99"
 webpki = "0.22.4"
-worker = { version = "0.5", features = ["http"] }
+worker = "0.5"
 x509-parser = "0.15.1"
 
 [workspace.dependencies.sentry]

Cargo.toml

@@ -21,6 +21,16 @@ helper:
 		-c ./crates/daphne-server/examples/configuration-helper.toml
 h: helper
 
+compute-offload:
+	RUST_LOG=hyper=off,debug cargo run \
+		--profile release-symbols \
+		--features test-utils \
+		--example service \
+		-- \
+		-c ./crates/daphne-server/examples/configuration-cpu-offload.toml
+co: compute-offload
+
+
 helper-worker:
 	cd ./crates/daphne-worker-test/ && \
 		wrangler dev -c wrangler.aggregator.toml --port 8788 -e helper

Makefile

@@ -27,6 +27,7 @@ use daphne::{
         BatchId, BatchSelector, PartialBatchSelector, TaskId,
     },
     metrics::DaphneMetrics,
+    protocol::ReadyAggregationJobResp,
     testing::report_generator::ReportGenerator,
     vdaf::VdafConfig,
     DapAggregateShare, DapAggregateSpan, DapAggregationParam, DapBatchMode, DapMeasurement,
@@ -511,7 +512,7 @@ impl Test {
         let _guard = load_control.wait().await;
         info!("Starting AggregationJobInitReq");
         let start = Instant::now();
-        let agg_job_resp = self
+        let mut agg_job_resp = self
             .http_client
             .submit_aggregation_job_init_req(
                 self.helper_url.join(&format!(
@@ -528,14 +529,42 @@ impl Test {
             )
             .await?;
         let duration = start.elapsed();
-        info!("Finished AggregationJobInitReq in {duration:#?}");
+        info!("Finished submitting AggregationJobInitReq in {duration:#?}");
+        let mut poll_count = 1;
+        let ready = loop {
+            agg_job_resp = match agg_job_resp {
+                messages::AggregationJobResp::Ready { prep_resps } => {
+                    if poll_count != 1 {
+                        info!(
+                            "Finished polling for AggregationJobResp after {:#?}",
+                            start.elapsed()
+                        );
+                    }
+                    break ReadyAggregationJobResp { prep_resps };
+                }
+                messages::AggregationJobResp::Processing => {
+                    if poll_count == 1 {
+                        info!("Polling for AggregationJobResp");
+                    }
+                    tokio::time::sleep(Duration::from_millis(poll_count * 200)).await;
+                    poll_count += 1;
+                    self.http_client
+                        .poll_aggregation_job_init(
+                            self.helper_url
+                                .join(&format!("tasks/{task_id}/aggregation_jobs/{agg_job_id}"))?,
+                            task_config.version,
+                            functions::helper::Options {
+                                taskprov_advertisement: taskprov_advertisement.as_ref(),
+                                bearer_token: self.bearer_token.as_ref(),
+                            },
+                        )
+                        .await?
+                }
+            };
+        };
 
-        let agg_share_span = task_config.consume_agg_job_resp(
-            task_id,
-            agg_job_state,
-            agg_job_resp.unwrap_ready(), // TODO: implement polling
-            self.metrics(),
-        )?;
+        let agg_share_span =
+            task_config.consume_agg_job_resp(task_id, agg_job_state, ready, self.metrics())?;
 
         let aggregated_report_count = agg_share_span
             .iter()

crates/dapf/src/acceptance/mod.rs

@@ -1,4 +1,4 @@
-// Copyright (c) 2024 Cloudflare, Inc. All rights reserved.
+// Copyright (c) 2025 Cloudflare, Inc. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
 use std::{future::ready, num::NonZeroUsize, ops::Range, time::SystemTime};
@@ -79,6 +79,7 @@ impl DapAggregator for crate::App {
     #[tracing::instrument(skip(self))]
     async fn get_agg_share(
         &self,
+        _version: DapVersion,
         task_id: &TaskId,
         batch_sel: &BatchSelector,
     ) -> Result<DapAggregateShare, DapError> {
@@ -115,6 +116,7 @@ impl DapAggregator for crate::App {
     #[tracing::instrument(skip(self))]
     async fn mark_collected(
         &self,
+        _version: DapVersion,
         task_id: &TaskId,
         batch_sel: &BatchSelector,
     ) -> Result<(), DapError> {
@@ -255,6 +257,7 @@ impl DapAggregator for crate::App {
 
     async fn is_batch_overlapping(
         &self,
+        _version: DapVersion,
         task_id: &TaskId,
         batch_sel: &BatchSelector,
     ) -> Result<bool, DapError> {
@@ -288,7 +291,12 @@ impl DapAggregator for crate::App {
         )
     }
 
-    async fn batch_exists(&self, task_id: &TaskId, batch_id: &BatchId) -> Result<bool, DapError> {
+    async fn batch_exists(
+        &self,
+        _version: DapVersion,
+        task_id: &TaskId,
+        batch_id: &BatchId,
+    ) -> Result<bool, DapError> {
         let task_config = self
             .get_task_config_for(task_id)
             .await?

crates/daphne-server/src/roles/aggregator.rs

@@ -3,7 +3,8 @@
 
 use axum::async_trait;
 use daphne::{
-    messages::{AggregationJobId, TaskId},
+    fatal_error,
+    messages::{AggregationJobId, AggregationJobResp, TaskId},
     roles::{helper::AggregationJobRequestHash, DapHelper},
     DapError, DapVersion,
 };
@@ -20,4 +21,13 @@ impl DapHelper for crate::App {
         // the server implementation can't check for this
         Ok(())
     }
+
+    async fn poll_aggregated(
+        &self,
+        _version: DapVersion,
+        _task_id: &TaskId,
+        _agg_job_id: &AggregationJobId,
+    ) -> Result<AggregationJobResp, DapError> {
+        Err(fatal_error!(err = "polling not implemented"))
+    }
 }

crates/daphne-server/src/roles/helper.rs

@@ -121,6 +121,10 @@ impl DapLeader for crate::App {
     {
         self.send_http(meta, Method::PUT, url, payload).await
     }
+
+    async fn send_http_get(&self, meta: DapRequestMeta, url: Url) -> Result<DapResponse, DapError> {
+        self.send_http(meta, Method::PUT, url, ()).await
+    }
 }
 
 impl crate::App {

crates/daphne-server/src/roles/leader.rs

@@ -72,12 +72,12 @@ impl DecodeFromDapHttpBody for HashedAggregationJobReq {
     }
 }
 
-/// Using `()` ignores the body of a request.
 impl DecodeFromDapHttpBody for CollectionPollReq {
     fn decode_from_http_body(_bytes: Bytes, _meta: &DapRequestMeta) -> Result<Self, DapAbort> {
         Ok(Self)
     }
 }
+
 /// Using `()` ignores the body of a request.
 impl DecodeFromDapHttpBody for () {
     fn decode_from_http_body(_bytes: Bytes, _meta: &DapRequestMeta) -> Result<Self, DapAbort> {

crates/daphne-server/src/router/extractor.rs

@@ -11,7 +11,10 @@ fn main() {
         #[cfg(feature = "durable_requests")]
         compiler
             .file("./src/durable_requests/durable_request.capnp")
-            .file("./src/durable_requests/bindings/aggregation_job_store.capnp");
+            .file("./src/durable_requests/bindings/aggregation_job_store.capnp")
+            .file("./src/durable_requests/bindings/aggregate_store_v2.capnp")
+            .file("./src/durable_requests/bindings/agg_job_response_store.capnp")
+            .file("./src/durable_requests/bindings/replay_checker.capnp");
 
         #[cfg(feature = "compute-offload")]
         compiler.file("./src/compute_offload/compute_offload.capnp");

crates/daphne-service-utils/build.rs

@@ -22,13 +22,29 @@ struct U8L16 @0x9e3f65b13f71cfcb {
     snd @1 :UInt64;
 }
 
-struct PartialBatchSelector {
+struct PartialBatchSelector @0xae86084e56c22fc0 {
     union {
         timeInterval @0 :Void;
         leaderSelectedByBatchId @1 :BatchId;
     }
 }
 
+enum ReportError @0xa76428617779e659 {
+    reserved @0;
+    batchCollected @1;
+    reportReplayed @2;
+    reportDropped @3;
+    hpkeUnknownConfigId @4;
+    hpkeDecryptError @5;
+    vdafPrepError @6;
+    batchSaturated @7;
+    taskExpired @8;
+    invalidMessage @9;
+    reportTooEarly @10;
+    taskNotStarted @11;
+}
+
+
 using ReportId = U8L16;
 using BatchId = U8L32;
 using TaskId = U8L32;