WIP

Author: mendess

Cargo.lock

@@ -89,8 +89,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

@@ -24,7 +24,7 @@ use daphne::{
     hpke::{HpkeConfig, HpkeKemId, HpkeReceiverConfig},
     messages::{
         self, taskprov::TaskprovAdvertisement, AggregateShareReq, AggregationJobId, Base64Encode,
-        BatchId, BatchSelector, PartialBatchSelector, TaskId,
+        BatchId, BatchSelector, PartialBatchSelector, ReadyAggregationJobResp, TaskId,
     },
     metrics::DaphneMetrics,
     testing::report_generator::ReportGenerator,
@@ -508,7 +508,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!(
@@ -525,14 +525,30 @@ impl Test {
             )
             .await?;
         let duration = start.elapsed();
-        info!("Finished AggregationJobInitReq in {duration:#?}");
+        info!("Finished submitting AggregationJobInitReq in {duration:#?}");
+        let ready = loop {
+            agg_job_resp = match agg_job_resp {
+                messages::AggregationJobResp::Ready { transitions } => {
+                    break ReadyAggregationJobResp { transitions }
+                }
+                messages::AggregationJobResp::Processing => {
+                    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,
-            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,7 +1,7 @@
 // Copyright (c) 2024 Cloudflare, Inc. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
-use anyhow::{anyhow, Context as _};
+use anyhow::{anyhow, Context};
 use daphne::{
     constants::DapMediaType,
     error::aborts::ProblemDetails,
@@ -12,7 +12,7 @@ use daphne::{
     DapVersion,
 };
 use daphne_service_utils::{bearer_token::BearerToken, http_headers};
-use prio::codec::{ParameterizedDecode as _, ParameterizedEncode as _};
+use prio::codec::{ParameterizedDecode, ParameterizedEncode as _};
 use reqwest::header;
 use url::Url;
 
@@ -43,30 +43,30 @@ impl HttpClient {
             .send()
             .await
             .context("sending AggregationJobInitReq")?;
-        if resp.status() == 400 {
-            let text = resp.text().await?;
-            let problem_details: ProblemDetails =
-                serde_json::from_str(&text).with_context(|| {
-                    format!("400 Bad Request: failed to parse problem details document: {text:?}")
-                })?;
-            Err(anyhow!("400 Bad Request: {problem_details:?}"))
-        } else if resp.status() == 500 {
-            Err(anyhow::anyhow!(
-                "500 Internal Server Error: {}",
-                resp.text().await?
-            ))
-        } else if !resp.status().is_success() {
-            Err(response_to_anyhow(resp).await).context("while running an AggregationJobInitReq")
-        } else {
-            AggregationJobResp::get_decoded_with_param(
-                &version,
-                &resp
-                    .bytes()
-                    .await
-                    .context("transfering bytes from the AggregateInitReq")?,
-            )
-            .with_context(|| "failed to parse response to AggregateInitReq from Helper")
-        }
+        handle_response(resp, &version).await
+    }
+
+    pub async fn poll_aggregation_job_init(
+        &self,
+        url: Url,
+        version: DapVersion,
+        opts: Options<'_>,
+    ) -> anyhow::Result<AggregationJobResp> {
+        let resp = self
+            .get(url)
+            .headers(construct_request_headers(
+                DapMediaType::AggregationJobInitReq
+                    .as_str_for_version(version)
+                    .with_context(|| {
+                        format!("AggregationJobInitReq media type is not defined for {version}")
+                    })?,
+                version,
+                opts,
+            )?)
+            .send()
+            .await
+            .context("polling aggregation job init req")?;
+        handle_response(resp, &version).await
     }
 
     pub async fn get_aggregate_share(
@@ -144,3 +144,32 @@ fn construct_request_headers(
     }
     Ok(headers)
 }
+
+async fn handle_response<R, P>(resp: reqwest::Response, params: &P) -> anyhow::Result<R>
+where
+    R: ParameterizedDecode<P>,
+{
+    if resp.status() == 400 {
+        let text = resp.text().await?;
+        let problem_details: ProblemDetails = serde_json::from_str(&text).with_context(|| {
+            format!("400 Bad Request: failed to parse problem details document: {text:?}")
+        })?;
+        Err(anyhow!("400 Bad Request: {problem_details:?}"))
+    } else if resp.status() == 500 {
+        Err(anyhow::anyhow!(
+            "500 Internal Server Error: {}",
+            resp.text().await?
+        ))
+    } else if !resp.status().is_success() {
+        Err(response_to_anyhow(resp).await).context("while running an AggregationJobInitReq")
+    } else {
+        R::get_decoded_with_param(
+            params,
+            &resp
+                .bytes()
+                .await
+                .context("transfering bytes from the AggregateInitReq")?,
+        )
+        .with_context(|| "failed to parse response to AggregateInitReq from Helper")
+    }
+}

crates/dapf/src/functions/helper.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> {

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");