Implement the async worker

Cargo.toml

@@ -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]

Makefile

@@ -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

crates/dapf/src/acceptance/mod.rs

@@ -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/daphne-server/docker/example-service.Dockerfile

@@ -1,7 +1,7 @@
 # Copyright (c) 2024 Cloudflare, Inc. All rights reserved.
 # SPDX-License-Identifier: BSD-3-Clause
 
-FROM rust:1.80-bookworm AS builder
+FROM rust:1.84.1-bookworm AS builder
 
 RUN apt update && \
     apt install -y \

crates/daphne-server/src/roles/aggregator.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/helper.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/leader.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/router/extractor.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-service-utils/Cargo.toml

@@ -21,7 +21,7 @@ serde_json = { workspace = true, optional = true }
 url = { workspace = true, optional = true }
 
 [dev-dependencies]
-daphne = { path = "../daphne", default-features = false, features = ["prometheus"] }
+daphne = { path = "../daphne", default-features = false, features = ["prometheus", "test-utils"] }
 rand.workspace = true
 
 [build-dependencies]

crates/daphne-service-utils/build.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/src/capnproto/base.capnp

@@ -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;

crates/daphne-service-utils/src/capnproto/mod.rs

@@ -4,6 +4,7 @@
 use crate::base_capnp::{self, partial_batch_selector, u8_l16, u8_l32};
 use capnp::struct_list;
 use capnp::traits::{FromPointerBuilder, FromPointerReader};
+use daphne::messages;
 use daphne::{
     messages::{AggregationJobId, BatchId, PartialBatchSelector, ReportId, TaskId},
     DapVersion,
@@ -204,6 +205,44 @@ impl CapnprotoPayloadDecode for PartialBatchSelector {
     }
 }
 
+impl From<messages::ReportError> for base_capnp::ReportError {
+    fn from(failure: messages::ReportError) -> Self {
+        match failure {
+            messages::ReportError::Reserved => Self::Reserved,
+            messages::ReportError::BatchCollected => Self::BatchCollected,
+            messages::ReportError::ReportReplayed => Self::ReportReplayed,
+            messages::ReportError::ReportDropped => Self::ReportDropped,
+            messages::ReportError::HpkeUnknownConfigId => Self::HpkeUnknownConfigId,
+            messages::ReportError::HpkeDecryptError => Self::HpkeDecryptError,
+            messages::ReportError::VdafPrepError => Self::VdafPrepError,
+            messages::ReportError::BatchSaturated => Self::BatchSaturated,
+            messages::ReportError::TaskExpired => Self::TaskExpired,
+            messages::ReportError::InvalidMessage => Self::InvalidMessage,
+            messages::ReportError::ReportTooEarly => Self::ReportTooEarly,
+            messages::ReportError::TaskNotStarted => Self::TaskNotStarted,
+        }
+    }
+}
+
+impl From<base_capnp::ReportError> for messages::ReportError {
+    fn from(val: base_capnp::ReportError) -> Self {
+        match val {
+            base_capnp::ReportError::Reserved => Self::Reserved,
+            base_capnp::ReportError::BatchCollected => Self::BatchCollected,
+            base_capnp::ReportError::ReportReplayed => Self::ReportReplayed,
+            base_capnp::ReportError::ReportDropped => Self::ReportDropped,
+            base_capnp::ReportError::HpkeUnknownConfigId => Self::HpkeUnknownConfigId,
+            base_capnp::ReportError::HpkeDecryptError => Self::HpkeDecryptError,
+            base_capnp::ReportError::VdafPrepError => Self::VdafPrepError,
+            base_capnp::ReportError::BatchSaturated => Self::BatchSaturated,
+            base_capnp::ReportError::TaskExpired => Self::TaskExpired,
+            base_capnp::ReportError::InvalidMessage => Self::InvalidMessage,
+            base_capnp::ReportError::ReportTooEarly => Self::ReportTooEarly,
+            base_capnp::ReportError::TaskNotStarted => Self::TaskNotStarted,
+        }
+    }
+}
+
 pub fn encode_list<I, O>(list: I, mut builder: struct_list::Builder<'_, O>)
 where
     I: IntoIterator<Item: CapnprotoPayloadEncode>,
@@ -229,3 +268,80 @@ pub fn usize_to_capnp_len(u: usize) -> u32 {
     u.try_into()
         .expect("capnp can't encode more that u32::MAX of something")
 }
+
+#[cfg(test)]
+pub fn roundtrip_test<T>(before: T)
+where
+    T: CapnprotoPayloadDecode + CapnprotoPayloadEncode + PartialEq + std::fmt::Debug,
+{
+    assert_eq!(
+        before,
+        T::decode_from_bytes(&before.encode_to_bytes()).unwrap()
+    );
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rand::Rng;
+
+    #[test]
+    fn test_u8_array_serialize_deserialize() {
+        roundtrip_test(rand::thread_rng().gen::<[u8; 32]>());
+        roundtrip_test(rand::thread_rng().gen::<[u8; 16]>());
+    }
+
+    #[test]
+    fn test_partial_batch_selector_serialize_deserialize() {
+        roundtrip_test(PartialBatchSelector::TimeInterval);
+
+        roundtrip_test(PartialBatchSelector::LeaderSelectedByBatchId {
+            batch_id: BatchId(rand::thread_rng().gen()),
+        });
+    }
+
+    #[test]
+    fn test_report_error_conversion() {
+        // cause a compilation error if the variants change
+        const _: () = {
+            #[allow(clippy::match_same_arms)]
+            match messages::ReportError::Reserved {
+                messages::ReportError::Reserved => (),
+                messages::ReportError::BatchCollected => (),
+                messages::ReportError::ReportReplayed => (),
+                messages::ReportError::ReportDropped => (),
+                messages::ReportError::HpkeUnknownConfigId => (),
+                messages::ReportError::HpkeDecryptError => (),
+                messages::ReportError::VdafPrepError => (),
+                messages::ReportError::BatchSaturated => (),
+                messages::ReportError::TaskExpired => (),
+                messages::ReportError::InvalidMessage => (),
+                messages::ReportError::ReportTooEarly => (),
+                messages::ReportError::TaskNotStarted => (),
+            }
+        };
+        let all_errors = vec![
+            messages::ReportError::Reserved,
+            messages::ReportError::BatchCollected,
+            messages::ReportError::ReportReplayed,
+            messages::ReportError::ReportDropped,
+            messages::ReportError::HpkeUnknownConfigId,
+            messages::ReportError::HpkeDecryptError,
+            messages::ReportError::VdafPrepError,
+            messages::ReportError::BatchSaturated,
+            messages::ReportError::TaskExpired,
+            messages::ReportError::InvalidMessage,
+            messages::ReportError::ReportTooEarly,
+            messages::ReportError::TaskNotStarted,
+        ];
+
+        for error in all_errors {
+            let converted: base_capnp::ReportError = error.into();
+            let back_converted: messages::ReportError = converted.into();
+            assert_eq!(
+                error, back_converted,
+                "Conversion symmetry failed for {error:?}",
+            );
+        }
+    }
+}