Repository: cloudflare/daphne
Pull request #330 (2023-05-25 17:49:02)

Improvements to HelperStateStore
This should take some pressure off of HelperStateStore. Note that deploying this will abort any aggregation jobs currently in flight.


Patch:
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index af96902..90042cf 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -764,13 +764,14 @@ pub trait DapHelper<'srv, 'req, S>: DapAggregator<'srv, 'req, S>
 where
     'srv: 'req,
 {
-    /// Store the Helper's aggregation-flow state.
-    async fn put_helper_state(
+    /// Store the Helper's aggregation-flow state unless it already exists. Returns a boolean
+    /// indicating if the operation succeeded.
+    async fn put_helper_state_if_not_exists(
         &self,
         task_id: &TaskId,
         agg_job_id: &MetaAggregationJobId,
         helper_state: &DapHelperState,
-    ) -> Result<(), DapError>;
+    ) -> Result<bool, DapError>;
 
     /// Fetch the Helper's aggregation-flow state. `None` is returned if the Helper has no state
     /// associated with the given task and aggregation job.
@@ -867,8 +868,6 @@ where
                     _ => unreachable!("unhandled resource {:?}", req.resource),
                 };
 
-                let helper_state = self.get_helper_state(task_id, &agg_job_id);
-
                 // Check whether the DAP version in the request matches the task config.
                 if task_config.version != req.version {
                     return Err(DapAbort::version_mismatch(req.version, task_config.version));
@@ -902,15 +901,6 @@ where
                     )
                     .await?;
 
-                // Check that helper state with the given task ID and aggregation job ID does not
-                // exist.
-                if helper_state.await?.is_some() {
-                    // TODO spec: Consider an explicit abort for this case.
-                    return Err(DapAbort::BadRequest(
-                        "unexpected message for aggregation job (already exists)".into(),
-                    ));
-                }
-
                 let agg_job_resp = match transition {
                     DapHelperTransition::Continue(mut state, mut agg_job_resp) => {
                         // Filter out early rejected reports.
@@ -951,7 +941,15 @@ where
                             }
                         }
 
-                        self.put_helper_state(task_id, &agg_job_id, &state).await?;
+                        if !self
+                            .put_helper_state_if_not_exists(task_id, &agg_job_id, &state)
+                            .await?
+                        {
+                            // TODO spec: Consider an explicit abort for this case.
+                            return Err(DapAbort::BadRequest(
+                                "unexpected message for aggregation job (already exists)".into(),
+                            ));
+                        }
                         agg_job_resp
                     }
                     DapHelperTransition::Finish(..) => {
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index 4b058d8..e95b5ea 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -20,7 +20,8 @@ use crate::{
         },
         durable_name_agg_store, durable_name_queue, durable_name_task,
         helper_state_store::{
-            durable_helper_state_name, DURABLE_HELPER_STATE_GET, DURABLE_HELPER_STATE_PUT,
+            durable_helper_state_name, DURABLE_HELPER_STATE_GET,
+            DURABLE_HELPER_STATE_PUT_IF_NOT_EXISTS,
         },
         leader_agg_job_queue::DURABLE_LEADER_AGG_JOB_QUEUE_GET,
         leader_batch_queue::{
@@ -991,25 +992,25 @@ impl<'srv, 'req> DapHelper<'srv, 'req, DaphneWorkerAuth> for DaphneWorker<'srv>
 where
     'srv: 'req,
 {
-    async fn put_helper_state(
+    async fn put_helper_state_if_not_exists(
         &self,
         task_id: &TaskId,
         agg_job_id: &MetaAggregationJobId,
         helper_state: &DapHelperState,
-    ) -> std::result::Result<(), DapError> {
+    ) -> std::result::Result<bool, DapError> {
         let task_config = self.try_get_task_config(task_id).await?;
         let helper_state_hex = hex::encode(helper_state.get_encoded(&task_config.as_ref().vdaf)?);
-        self.durable()
+        Ok(self
+            .durable()
             .with_retry()
             .post(
                 BINDING_DAP_HELPER_STATE_STORE,
-                DURABLE_HELPER_STATE_PUT,
+                DURABLE_HELPER_STATE_PUT_IF_NOT_EXISTS,
                 durable_helper_state_name(&task_config.as_ref().version, task_id, agg_job_id),
                 helper_state_hex,
             )
             .await
-            .map_err(dap_err)?;
-        Ok(())
+            .map_err(dap_err)?)
     }
 
     async fn get_helper_state(
@@ -1023,11 +1024,10 @@ where
         let res: Option<String> = self
             .durable()
             .with_retry()
-            .post(
+            .get(
                 BINDING_DAP_HELPER_STATE_STORE,
                 DURABLE_HELPER_STATE_GET,
                 durable_helper_state_name(&task_config.as_ref().version, task_id, agg_job_id),
-                (),
             )
             .await
             .map_err(dap_err)?;
diff --git a/daphne_worker/src/durable/helper_state_store.rs b/daphne_worker/src/durable/helper_state_store.rs
index b86b455..52c2cbf 100644
--- a/daphne_worker/src/durable/helper_state_store.rs
+++ b/daphne_worker/src/durable/helper_state_store.rs
@@ -1,7 +1,11 @@
 // Copyright (c) 2022 Cloudflare, Inc. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
-use crate::{config::DaphneWorkerConfig, durable::state_get, initialize_tracing, int_err};
+use crate::{
+    config::DaphneWorkerConfig,
+    durable::{state_get, state_set_if_not_exists},
+    initialize_tracing, int_err,
+};
 use daphne::{messages::TaskId, DapVersion, MetaAggregationJobId};
 use tracing::{trace, warn};
 use worker::*;
@@ -19,14 +23,16 @@ pub(crate) fn durable_helper_state_name(
     )
 }
 
-pub(crate) const DURABLE_HELPER_STATE_PUT: &str = "/internal/do/helper_state/put";
+pub(crate) const DURABLE_HELPER_STATE_PUT_IF_NOT_EXISTS: &str =
+    "/internal/do/helper_state/put_if_not_exists";
 pub(crate) const DURABLE_HELPER_STATE_GET: &str = "/internal/do/helper_state/get";
 
 /// Durable Object (DO) for storing the Helper's state for a given aggregation job.
 ///
 /// This object implements the following API endpoints:
 ///
-/// - `DURABLE_HELPER_STATE_PUT`: Stores Helper's hex-encoded state.
+/// - `DURABLE_HELPER_STATE_PUT_IF_NOT_EXISTS`: Stores Helper's hex-encoded state unless the state
+///    already exists. Returns a boolean indicating whether the operation succeeded.
 /// - `DURABLE_HELPER_STATE_GET`: Drains the Helper's hex-encoded state.
 ///
 /// The state blob is stored in `helper_state`.
@@ -62,35 +68,24 @@ impl DurableObject for HelperStateStore {
         match (req.path().as_ref(), req.method()) {
             // Store the Helper's state.
             //
-            // Idempotent
+            // Non-idempotent
             // Input: `helper_state_hex: String` (hex-encoded state)
-            // Output: `()`
-            (DURABLE_HELPER_STATE_PUT, Method::Post) => {
-                // The state is handled as an opaque hex string.
-                let mut helper_state_hex: Option<String> =
-                    state_get(&self.state, "helper_state").await?;
-                if helper_state_hex.is_some() {
-                    // TODO spec: Handle this as an abort rather than an internal error.
-                    return Err(int_err("tried to overwrite helper state"));
-                }
-
-                helper_state_hex = req.json().await?;
-                self.state
-                    .storage()
-                    .put("helper_state", helper_state_hex)
-                    .await?;
-                Response::from_json(&())
+            // Output: `bool`
+            (DURABLE_HELPER_STATE_PUT_IF_NOT_EXISTS, Method::Post) => {
+                let helper_state_hex: String = req.json().await?;
+                let success =
+                    state_set_if_not_exists(&self.state, "helper_state", &helper_state_hex)
+                        .await?
+                        .is_none();
+                Response::from_json(&success)
             }
 
-            // Drain the Helper's state.
+            // Get the Helper's state.
             //
-            // Non-idempotent (do not retry)
+            // Idempotent
             // Output: `String` (hex-encoded state)
-            (DURABLE_HELPER_STATE_GET, Method::Post) => {
+            (DURABLE_HELPER_STATE_GET, Method::Get) => {
                 let helper_state: Option<String> = state_get(&self.state, "helper_state").await?;
-                if helper_state.is_some() {
-                    self.state.storage().delete("helper_state").await?;
-                }
                 Response::from_json(&helper_state)
             }
 
diff --git a/daphne_worker/src/durable/reports_processed.rs b/daphne_worker/src/durable/reports_processed.rs
index 5917268..459d440 100644
--- a/daphne_worker/src/durable/reports_processed.rs
+++ b/daphne_worker/src/durable/reports_processed.rs
@@ -80,7 +80,7 @@ impl DurableObject for ReportsProcessed {
             // Mark a set of reports as aggregated. Return the set of report IDs that already
             // exist.
             //
-            // Non-idempotent (do not retry)
+            // Non-idempotent
             // Input: `report_id_hex_set: Vec<String>` (hex-encoded report IDs)
             // Output: `Vec<String>` (subset of the inputs that already exist).
             (DURABLE_REPORTS_PROCESSED_MARK_AGGREGATED, Method::Post) => {


Test patch:
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index 56b9157..545861d 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -554,12 +554,12 @@ impl<'srv, 'req> DapHelper<'srv, 'req, BearerToken> for MockAggregator
 where
     'srv: 'req,
 {
-    async fn put_helper_state(
+    async fn put_helper_state_if_not_exists(
         &self,
         task_id: &TaskId,
         agg_job_id: &MetaAggregationJobId,
         helper_state: &DapHelperState,
-    ) -> Result<(), DapError> {
+    ) -> Result<bool, DapError> {
         let helper_state_info = HelperStateInfo {
             task_id: task_id.clone(),
             agg_job_id_owned: agg_job_id.into(),
@@ -573,16 +573,14 @@ where
         let helper_state_store = helper_state_store_mutex_guard.deref_mut();
 
         if helper_state_store.contains_key(&helper_state_info) {
-            return Err(DapError::Fatal(
-                "overwriting existing helper state".to_string(),
-            ));
+            return Ok(false);
         }
 
         // NOTE: This code is only correct for VDAFs with exactly one round of preparation.
         // For VDAFs with more rounds, the helper state blob will need to be updated here.
         helper_state_store.insert(helper_state_info, helper_state.clone());
 
-        Ok(())
+        Ok(true)
     }
 
     async fn get_helper_state(