Repository: cloudflare/daphne
Pull request #61 (2022-07-25 16:56:20)

Implement LeaderState and Collect job related methods
Implements final missing pieces of #20, in particular methods used by the Leader to handle collect jobs.

This PR implements LeaderState, which stores collect jobs it received from the Collector.
Collect job related methods, `init_collect_job`, `poll_collect_job`, `get_pending_collect_jobs`, and `finish_collect_job` are also implemented.

The full e2e test is also implemented, starting from the Client submitting a report, the Leader and Helper processing it, and finally the Leader processing the collect job received from the Collector.


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index 4005dc6..853bd16 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -491,6 +491,7 @@ dependencies = [
  "getrandom",
  "hex",
  "hpke 0.8.0",
+ "matchit 0.6.0",
  "prio",
  "rand",
  "reqwest",
@@ -1428,6 +1429,12 @@ version = "0.4.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9376a4f0340565ad675d11fc1419227faf5f60cd7ac9cb2e7185a471f30af833"
 
+[[package]]
+name = "matchit"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3dfc802da7b1cf80aefffa0c7b2f77247c8b32206cc83c270b61264f5b360a80"
+
 [[package]]
 name = "md-5"
 version = "0.10.1"
@@ -3318,7 +3325,7 @@ dependencies = [
  "futures-util",
  "http",
  "js-sys",
- "matchit",
+ "matchit 0.4.6",
  "pin-project",
  "serde",
  "serde_json",
diff --git a/daphne/Cargo.toml b/daphne/Cargo.toml
index 11529ea..cf8961a 100644
--- a/daphne/Cargo.toml
+++ b/daphne/Cargo.toml
@@ -37,4 +37,5 @@ reqwest = { version = "0.11.11", features = ["blocking"] } # Required for daph C
 anyhow = "1.0.58" # Required for daph CLI
 
 [dev-dependencies]
+matchit = "0.6.0"
 tokio = { version = "^1.19", features = ["macros","rt"] }
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index 2762ac3..4200fb6 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -517,7 +517,7 @@ pub struct DapResponse {
 }
 
 /// Status of a collect job.
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Debug, Deserialize, PartialEq, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub enum DapCollectJob {
     Done(CollectResp),
diff --git a/daphne/src/messages.rs b/daphne/src/messages.rs
index b616303..e419a39 100644
--- a/daphne/src/messages.rs
+++ b/daphne/src/messages.rs
@@ -397,7 +397,7 @@ impl Decode for Interval {
 /// A collect request.
 //
 // TODO Add serialization tests.
-#[derive(Clone, Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
 pub struct CollectReq {
     pub task_id: Id,
     pub batch_interval: Interval,
@@ -425,7 +425,7 @@ impl Decode for CollectReq {
 /// A collect response.
 //
 // TODO Add serialization tests.
-#[derive(Clone, Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
 pub struct CollectResp {
     pub encrypted_agg_shares: Vec<HpkeCiphertext>,
 }
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index b65dc17..3da4fd4 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -170,7 +170,7 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
     async fn get_pending_collect_jobs(
         &self,
         task_id: &Id,
-    ) -> Result<Vec<(Id, CollectReq)>, DapError>;
+    ) -> Result<HashMap<Id, Vec<CollectReq>>, DapError>;
 
     /// Complete a collect job by assigning it the completed [`CollectResp`](crate::messages::CollectResp).
     async fn finish_collect_job(
@@ -403,10 +403,12 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
             // share computed during a collect job and any output shares computed during an aggregation
             // job.
             let pending = self.get_pending_collect_jobs(&task_id).await?;
-            for (collect_id, collect_req) in pending.into_iter() {
-                telem.reports_collected += self
-                    .run_collect_job(&collect_id, &task_id, task_config, collect_req)
-                    .await?;
+            for (collect_id, collect_reqs) in pending.into_iter() {
+                for collect_req in collect_reqs {
+                    telem.reports_collected += self
+                        .run_collect_job(&collect_id, &task_id, task_config, collect_req)
+                        .await?;
+                }
             }
         }
 
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index 418c568..c1df0d6 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -359,12 +359,12 @@ impl<D> DapLeader<BearerToken> for DaphneWorkerConfig<D> {
     async fn get_pending_collect_jobs(
         &self,
         task_id: &Id,
-    ) -> std::result::Result<Vec<(Id, CollectReq)>, DapError> {
+    ) -> std::result::Result<HashMap<Id, Vec<CollectReq>>, DapError> {
         let leader_state_namespace = self.durable_object("DAP_LEADER_STATE_STORE")?;
         let leader_state_stub = leader_state_namespace
             .id_from_name(&durable_leader_state_name(task_id))?
             .get_stub()?;
-        let res: Vec<(Id, CollectReq)> =
+        let res: HashMap<Id, Vec<CollectReq>> =
             durable_get!(leader_state_stub, DURABLE_LEADER_STATE_GET_COLLECT_REQS)
                 .await?
                 .json()
diff --git a/daphne_worker/src/durable/leader_state_store.rs b/daphne_worker/src/durable/leader_state_store.rs
index 1b3def7..ff1c6ed 100644
--- a/daphne_worker/src/durable/leader_state_store.rs
+++ b/daphne_worker/src/durable/leader_state_store.rs
@@ -14,7 +14,10 @@ use prio::{
     vdaf::prg::{Prg, PrgAes128, Seed, SeedStream},
 };
 use serde::{Deserialize, Serialize};
-use std::convert::TryInto;
+use std::{
+    collections::{HashMap, VecDeque},
+    convert::TryInto,
+};
 use worker::*;
 
 pub(crate) fn durable_leader_state_name(task_id: &Id) -> String {
@@ -130,7 +133,7 @@ impl DurableObject for LeaderStateStore {
                 let opt = ListOptions::new().prefix("pending/");
                 let iter = self.state.storage().list_with_options(opt).await?.entries();
                 let mut item = iter.next()?;
-                let mut res = Vec::new();
+                let mut list = Vec::new();
                 while !item.done() {
                     let (pending_key, ordered_collect_req): (String, OrderedCollectReq) =
                         item.value().into_serde()?;
@@ -140,15 +143,21 @@ impl DurableObject for LeaderStateStore {
                             .map_err(int_err)?
                             .try_into()
                             .map_err(|_| int_err("malformed key for pending CollectReq"))?);
-                    res.push((collect_id, ordered_collect_req));
+                    list.push((collect_id, ordered_collect_req));
                     item = iter.next()?;
                 }
 
-                res.sort_unstable_by_key(|(_id, prioritized)| prioritized.priority);
-                let res: Vec<(Id, CollectReq)> = res
+                list.sort_unstable_by_key(|(_id, prioritized)| prioritized.priority);
+                let mut list: VecDeque<(Id, CollectReq)> = list
                     .into_iter()
-                    .map(|(id, prioritized)| (id, prioritized.collect_req))
+                    .map(|(id, ordered_collect_req)| (id, ordered_collect_req.collect_req))
                     .collect();
+                let mut res: HashMap<Id, Vec<CollectReq>> = HashMap::default();
+                while !list.is_empty() {
+                    let (id, collect_req) = list.pop_front().unwrap();
+                    let collect_reqs = res.entry(id).or_insert_with(Vec::default);
+                    collect_reqs.push(collect_req);
+                }
                 Response::from_json(&res)
             }
 


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index 2bee4f3..d8781bf 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -10,17 +10,19 @@ use crate::{
     hpke::HpkeReceiverConfig,
     messages::{
         AggregateContinueReq, AggregateInitializeReq, AggregateResp, AggregateShareReq,
-        AggregateShareResp, CollectReq, HpkeCiphertext, Id, Interval, Nonce, Report, ReportShare,
-        Transition, TransitionFailure, TransitionVar,
+        AggregateShareResp, CollectReq, CollectResp, HpkeCiphertext, Id, Interval, Nonce, Report,
+        ReportShare, Transition, TransitionFailure, TransitionVar,
     },
     roles::{DapAggregator, DapHelper, DapLeader},
     testing::{
-        BucketInfo, MockAggregateInfo, MockAggregator, ReportStore, HPKE_RECEIVER_CONFIG_LIST,
-        LEADER_BEARER_TOKEN,
+        BucketInfo, MockAggregateInfo, MockAggregator, ReportStore, COLLECTOR_BEARER_TOKEN,
+        HPKE_RECEIVER_CONFIG_LIST, LEADER_BEARER_TOKEN,
     },
-    DapAbort, DapError, DapLeaderTransition, DapMeasurement, DapRequest, Prio3Config, VdafConfig,
+    DapAbort, DapCollectJob, DapError, DapLeaderTransition, DapMeasurement, DapRequest,
+    Prio3Config, VdafConfig,
 };
 use assert_matches::assert_matches;
+use matchit::Router;
 use prio::codec::{Decode, Encode};
 use rand::{thread_rng, Rng};
 use std::{ops::DerefMut, vec};
@@ -499,6 +501,112 @@ async fn get_reports_empty_response() {
     assert_eq!(reports.len(), 0);
 }
 
+#[tokio::test]
+async fn poll_collect_job_test_results() {
+    let leader = MockAggregator::new();
+    let task_id = leader.nominal_task_id();
+    let task_config = leader.get_task_config_for(task_id).unwrap();
+    let now = 1637361337;
+
+    // Collector: Create a CollectReq.
+    let collector_collect_req = CollectReq {
+        task_id: task_id.clone(),
+        batch_interval: task_config.current_batch_window(now),
+        agg_param: Vec::default(),
+    };
+    let req = DapRequest {
+        media_type: Some(MEDIA_TYPE_COLLECT_REQ),
+        payload: collector_collect_req.get_encoded(),
+        url: task_config.helper_url.join("/collect").unwrap(),
+        sender_auth: Some(BearerToken::from(COLLECTOR_BEARER_TOKEN.to_string())),
+    };
+
+    // Leader: Handle the CollectReq received from Collector.
+    leader.http_post_collect(&req).await.unwrap();
+
+    // Expect DapCollectJob::Unknown due to invalid collect ID.
+    assert_eq!(
+        leader
+            .poll_collect_job(task_id, &Id::default())
+            .await
+            .unwrap(),
+        DapCollectJob::Unknown
+    );
+
+    // Leader: Get pending collect job to obtain collect_id
+    let resp = leader.get_pending_collect_jobs(&task_id).await.unwrap();
+    let collect_id = resp.keys().collect::<Vec<&Id>>()[0];
+    let collect_resp = CollectResp {
+        encrypted_agg_shares: Vec::default(),
+    };
+
+    // Expect DapCollectJob::Pending due to pending collect job.
+    assert_eq!(
+        leader.poll_collect_job(task_id, &collect_id).await.unwrap(),
+        DapCollectJob::Pending
+    );
+
+    // Leader: Complete the collect job by storing CollectResp in LeaderStore.processed.
+    leader
+        .finish_collect_job(&task_id, &collect_id, &collect_resp)
+        .await
+        .unwrap();
+
+    // Expect DapCollectJob::Done due to processed collect job.
+    assert_matches!(
+        leader.poll_collect_job(task_id, &collect_id).await.unwrap(),
+        DapCollectJob::Done(..)
+    );
+}
+
+// Test a successful collect request submission.
+// This checks that the Leader reponds with the collect ID with the ID associated to the request.
+//
+// TODO(nakatsuka-y) Implement a test case when a Leader rejects a collect request.
+#[tokio::test]
+async fn http_post_collect_success() {
+    let leader = MockAggregator::new();
+    let task_id = leader.nominal_task_id();
+    let task_config = leader.get_task_config_for(task_id).unwrap();
+    let now = 1637361337;
+
+    // Collector: Create a CollectReq.
+    let collector_collect_req = CollectReq {
+        task_id: task_id.clone(),
+        batch_interval: task_config.current_batch_window(now),
+        agg_param: Vec::default(),
+    };
+    let req = DapRequest {
+        media_type: Some(MEDIA_TYPE_COLLECT_REQ),
+        payload: collector_collect_req.get_encoded(),
+        url: task_config.helper_url.join("/collect").unwrap(),
+        sender_auth: Some(BearerToken::from(COLLECTOR_BEARER_TOKEN.to_string())),
+    };
+
+    // Leader: Handle the CollectReq received from Collector.
+    let url = leader.http_post_collect(&req).await.unwrap();
+    let mut resp = leader.get_pending_collect_jobs(&task_id).await.unwrap();
+    let leader_collect_id = resp.keys().collect::<Vec<&Id>>()[0].clone();
+    let leader_collect_req = resp.remove(&leader_collect_id).unwrap();
+
+    // Check that the CollectReq sent by Collector is the same that is received by Leader.
+    assert_eq!(collector_collect_req, leader_collect_req[0]);
+
+    // Check that the collect_id included in the URI is the same with the one received
+    // by Leader.
+    let path = url.path().to_string();
+    let mut router = Router::new();
+    router
+        .insert("/collect/task/:task_id/req/:collect_id", true)
+        .unwrap();
+    let url_match = router.at(&path).unwrap();
+    let collector_collect_id = url_match.params.get("collect_id").unwrap();
+    assert_eq!(
+        collector_collect_id.to_string(),
+        leader_collect_id.to_base64url()
+    );
+}
+
 // Test the upload sub-protocol.
 #[tokio::test]
 async fn successful_http_post_upload() {
@@ -591,13 +699,24 @@ async fn e2e() {
         .unwrap();
     leader.put_out_shares(&task_id, out_shares).await.unwrap();
 
-    // TODO(nakatsuka-y) Call http_post_collect to post a collect job.
-    //                   Assume for now that the leader has already picked a collect job.
+    // Collector: Create a CollectReq.
     let collect_req = CollectReq {
         task_id: task_id.clone(),
         batch_interval: task_config.current_batch_window(now),
         agg_param: Vec::default(),
     };
+    let req = DapRequest {
+        media_type: Some(MEDIA_TYPE_COLLECT_REQ),
+        payload: collect_req.get_encoded(),
+        url: task_config.helper_url.join("/collect").unwrap(),
+        sender_auth: Some(BearerToken::from(COLLECTOR_BEARER_TOKEN.to_string())),
+    };
+
+    // Leader: Handle the CollectReq received from Collector.
+    leader.http_post_collect(&req).await.unwrap();
+    let mut resp = leader.get_pending_collect_jobs(&task_id).await.unwrap();
+    let collect_id = resp.keys().collect::<Vec<&Id>>()[0].clone();
+    let collect_req = resp.remove(&collect_id).unwrap()[0].clone();
 
     // Leader: Get Leader's encrypted aggregate share.
     let leader_agg_share = leader
@@ -605,7 +724,7 @@ async fn e2e() {
         .await
         .unwrap();
 
-    let _leader_enc_agg_share = task_config
+    let leader_enc_agg_share = task_config
         .vdaf
         .produce_leader_encrypted_agg_share(
             &task_config.collector_hpke_config,
@@ -617,9 +736,9 @@ async fn e2e() {
 
     // Leader: Prepare AggregateShareReq.
     let agg_share_req = AggregateShareReq {
-        task_id: collect_req.task_id.clone(),
-        batch_interval: collect_req.batch_interval.clone(),
-        agg_param: collect_req.agg_param.clone(),
+        task_id: collect_req.task_id,
+        batch_interval: collect_req.batch_interval,
+        agg_param: collect_req.agg_param,
         report_count: leader_agg_share.report_count,
         checksum: leader_agg_share.checksum,
     };
@@ -633,5 +752,28 @@ async fn e2e() {
     };
     let res = leader.http_post_aggregate_share(&req).await.unwrap();
     let agg_share_resp = AggregateShareResp::get_decoded(&res.payload).unwrap();
-    let _helper_enc_agg_share = agg_share_resp.encrypted_agg_share;
+    let helper_enc_agg_share = agg_share_resp.encrypted_agg_share;
+
+    // Leader: Complete the collect job by storing CollectResp in LeaderStore.processed.
+    let collect_resp = CollectResp {
+        encrypted_agg_shares: vec![leader_enc_agg_share, helper_enc_agg_share],
+    };
+
+    leader
+        .finish_collect_job(&task_id, &collect_id, &collect_resp)
+        .await
+        .unwrap();
+
+    // Leader: Mark the reports as collected.
+    leader
+        .mark_collected(&task_id, &agg_share_req.batch_interval)
+        .await
+        .unwrap();
+
+    // Leader: Respond to poll request from Collector.
+    let collect_job = leader
+        .poll_collect_job(&task_id, &collect_id)
+        .await
+        .unwrap();
+    assert_matches!(collect_job, DapCollectJob::Done(..))
 }
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index d43c389..d8053ba 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -15,9 +15,10 @@ use crate::{
     DapResponse, DapTaskConfig,
 };
 use async_trait::async_trait;
+use rand::{thread_rng, Rng};
 use serde::{Deserialize, Serialize};
 use std::{
-    collections::{HashMap, HashSet},
+    collections::{HashMap, HashSet, VecDeque},
     ops::DerefMut,
     sync::{Arc, Mutex},
     time::SystemTime,
@@ -74,6 +75,7 @@ pub(crate) struct MockAggregator {
     tasks: HashMap<Id, DapTaskConfig>,
     hpke_receiver_config_list: Vec<HpkeReceiverConfig>,
     pub(crate) report_store: Arc<Mutex<HashMap<BucketInfo, ReportStore>>>,
+    leader_state_store: Arc<Mutex<HashMap<Id, LeaderState>>>,
     helper_state_store: Arc<Mutex<HashMap<HelperStateInfo, DapHelperState>>>,
     agg_store: Arc<Mutex<HashMap<BucketInfo, DapAggregateShare>>>,
 }
@@ -89,6 +91,7 @@ impl MockAggregator {
                 .expect("failed to parse hpke_receiver_config_list");
 
         let report_store = Arc::new(Mutex::new(HashMap::new()));
+        let leader_state_store = Arc::new(Mutex::new(HashMap::new()));
         let helper_state_store = Arc::new(Mutex::new(HashMap::new()));
         let agg_store = Arc::new(Mutex::new(HashMap::new()));
 
@@ -96,6 +99,7 @@ impl MockAggregator {
             tasks,
             hpke_receiver_config_list,
             report_store,
+            leader_state_store,
             helper_state_store,
             agg_store,
         }
@@ -193,7 +197,7 @@ impl DapAggregator<BearerToken> for MockAggregator {
     ) -> Result<(), DapError> {
         let task_config = self
             .get_task_config_for(task_id)
-            .ok_or_else(|| DapError::fatal("no task found"))?;
+            .ok_or_else(|| DapError::fatal("task not found"))?;
 
         let agg_shares =
             DapAggregateShare::batches_from_out_shares(out_shares, task_config.min_batch_duration)?;
@@ -290,7 +294,7 @@ impl DapHelper<BearerToken> for MockAggregator {
             let bucket_info = BucketInfo::new(
                 self.tasks
                     .get(task_id)
-                    .ok_or_else(|| DapError::fatal("no task found"))?,
+                    .ok_or_else(|| DapError::fatal("task not found"))?,
                 task_id,
                 &report_share.nonce,
             );
@@ -383,7 +387,7 @@ impl DapLeader<BearerToken> for MockAggregator {
         let bucket_info = BucketInfo::new(
             self.tasks
                 .get(task_id)
-                .ok_or_else(|| DapError::fatal("no task found"))?,
+                .ok_or_else(|| DapError::fatal("task not found"))?,
             task_id,
             &report.nonce,
         );
@@ -414,7 +418,7 @@ impl DapLeader<BearerToken> for MockAggregator {
     ) -> Result<HashMap<Id, Vec<Report>>, DapError> {
         let task_config = self
             .get_task_config_for(&selector.task_id)
-            .ok_or_else(|| DapError::fatal("no task found"))?;
+            .ok_or_else(|| DapError::fatal("task not found"))?;
 
         let now = SystemTime::now()
             .duration_since(SystemTime::UNIX_EPOCH)
@@ -470,32 +474,137 @@ impl DapLeader<BearerToken> for MockAggregator {
         Ok(HashMap::from([(selector.task_id.clone(), reports)]))
     }
 
-    async fn init_collect_job(&self, _collect_req: &CollectReq) -> Result<Url, DapError> {
-        unreachable!("not implemented");
+    // Called after receiving a CollectReq from Collector.
+    async fn init_collect_job(&self, collect_req: &CollectReq) -> Result<Url, DapError> {
+        let mut rng = thread_rng();
+        let task_config = self
+            .get_task_config_for(&collect_req.task_id)
+            .ok_or_else(|| DapError::fatal("task not found"))?;
+
+        let mut leader_state_store_mutex_guard = self
+            .leader_state_store
+            .lock()
+            .map_err(|e| DapError::Fatal(e.to_string()))?;
+        let leader_state_store = leader_state_store_mutex_guard.deref_mut();
+
+        // Construct a new Collect URI for this CollectReq.
+        let collect_id = Id(rng.gen());
+        let collect_uri = task_config
+            .leader_url
+            .join(&format!(
+                "/collect/task/{}/req/{}",
+                collect_req.task_id.to_base64url(),
+                collect_id.to_base64url(),
+            ))
+            .map_err(|e| DapError::Fatal(e.to_string()))?;
+
+        // Store Collect ID and CollectReq into LeaderState.
+        let leader_state = leader_state_store
+            .entry(collect_req.task_id.clone())
+            .or_insert_with(LeaderState::new);
+        leader_state.collect_ids.push_back(collect_id.clone());
+        let collect_job_state = CollectJobState::Pending(collect_req.clone());
+        leader_state
+            .collect_jobs
+            .insert(collect_id, collect_job_state);
+
+        Ok(collect_uri)
     }
 
+    // Called to retrieve completed CollectResp at the request of Collector.
     async fn poll_collect_job(
         &self,
-        _task_id: &Id,
-        _collect_id: &Id,
+        task_id: &Id,
+        collect_id: &Id,
     ) -> Result<DapCollectJob, DapError> {
-        unreachable!("not implemented");
+        let mut leader_state_store_mutex_guard = self
+            .leader_state_store
+            .lock()
+            .map_err(|e| DapError::Fatal(e.to_string()))?;
+        let leader_state_store = leader_state_store_mutex_guard.deref_mut();
+
+        let leader_state = leader_state_store
+            .get(task_id)
+            .ok_or_else(|| DapError::fatal("collect job not found for task_id"))?;
+        if let Some(collect_job_state) = leader_state.collect_jobs.get(collect_id) {
+            match collect_job_state {
+                CollectJobState::Pending(_) => Ok(DapCollectJob::Pending),
+                CollectJobState::Processed(resp) => Ok(DapCollectJob::Done(resp.clone())),
+            }
+        } else {
+            Ok(DapCollectJob::Unknown)
+        }
     }
 
+    // Called to retrieve pending CollectReq.
     async fn get_pending_collect_jobs(
         &self,
-        _task_id: &Id,
-    ) -> Result<Vec<(Id, CollectReq)>, DapError> {
-        unreachable!("not implemented");
+        task_id: &Id,
+    ) -> Result<HashMap<Id, Vec<CollectReq>>, DapError> {
+        let mut leader_state_store_mutex_guard = self
+            .leader_state_store
+            .lock()
+            .map_err(|e| DapError::Fatal(e.to_string()))?;
+        let leader_state_store = leader_state_store_mutex_guard.deref_mut();
+
+        let leader_state = leader_state_store
+            .get_mut(task_id)
+            .ok_or_else(|| DapError::fatal("collect job not found for task_id"))?;
+        let mut res: HashMap<Id, Vec<CollectReq>> = HashMap::default();
+
+        // Iterate over collect IDs and copy them and their associated requests to the response.
+        for collect_id in leader_state.collect_ids.iter() {
+            if let CollectJobState::Pending(collect_req) =
+                leader_state.collect_jobs.get(collect_id).unwrap()
+            {
+                res.entry(collect_id.clone())
+                    .or_insert_with(Vec::default)
+                    .push(collect_req.clone());
+            }
+        }
+        Ok(res)
     }
 
+    // Called after finishing aggregation job to put resuts into LeaderState.
     async fn finish_collect_job(
         &self,
-        _task_id: &Id,
-        _collect_id: &Id,
-        _collect_resp: &CollectResp,
+        task_id: &Id,
+        collect_id: &Id,
+        collect_resp: &CollectResp,
     ) -> Result<(), DapError> {
-        unreachable!("not implemented");
+        let mut leader_state_store_mutex_guard = self
+            .leader_state_store
+            .lock()
+            .map_err(|e| DapError::Fatal(e.to_string()))?;
+        let leader_state_store = leader_state_store_mutex_guard.deref_mut();
+
+        let leader_state = leader_state_store
+            .get_mut(task_id)
+            .ok_or_else(|| DapError::fatal("collect job not found for task_id"))?;
+        let collect_job = leader_state
+            .collect_jobs
+            .get_mut(collect_id)
+            .ok_or_else(|| DapError::fatal("collect job not found for collect_id"))?;
+
+        match collect_job {
+            CollectJobState::Pending(_) => {
+                // Mark collect job as Processed.
+                *collect_job = CollectJobState::Processed(collect_resp.clone());
+
+                // Remove collect ID from queue.
+                let index = leader_state
+                    .collect_ids
+                    .iter()
+                    .position(|r| r == collect_id)
+                    .unwrap();
+                leader_state.collect_ids.remove(index);
+
+                Ok(())
+            }
+            CollectJobState::Processed(_) => {
+                Err(DapError::fatal("tried to overwrite collect response"))
+            }
+        }
     }
 
     async fn send_http_post(&self, _req: DapRequest<BearerToken>) -> Result<DapResponse, DapError> {
@@ -590,3 +699,25 @@ impl ReportStore {
         self.collected = true;
     }
 }
+
+/// Stores the state of the collect job.
+pub(crate) enum CollectJobState {
+    Pending(CollectReq),
+    Processed(CollectResp),
+}
+
+/// LeaderState keeps track of Collect IDs in their order of arrival.
+/// It also holds the state of the collect job associated to the Collect ID.
+pub(crate) struct LeaderState {
+    collect_ids: VecDeque<Id>,
+    collect_jobs: HashMap<Id, CollectJobState>,
+}
+
+impl LeaderState {
+    pub(crate) fn new() -> LeaderState {
+        Self {
+            collect_ids: VecDeque::default(),
+            collect_jobs: HashMap::default(),
+        }
+    }
+}