daphne_worker 0.2.0

// Copyright (c) 2022 Cloudflare, Inc. All rights reserved.
// SPDX-License-Identifier: BSD-3-Clause

use crate::{
    durable::{DurableOrdered, BINDING_DAP_LEADER_AGG_JOB_QUEUE},
    int_err,
};
use worker::*;

pub(crate) const DURABLE_LEADER_AGG_JOB_QUEUE_PUT: &str = "/internal/do/agg_job_queue/put";
pub(crate) const DURABLE_LEADER_AGG_JOB_QUEUE_GET: &str = "/internal/do/agg_job_queue/get";
pub(crate) const DURABLE_LEADER_AGG_JOB_QUEUE_FINISH: &str = "/internal/do/agg_job_queue/finish";

/// Durable Object (DO) representing an aggregation job queue.
///
/// This object defines the following API endpoints:
///
/// - `DURABLE_LEADER_AGG_JOB_QUEUE_PUT`: Adds a job to the queue. This is called by an instance of
///   `ReportsPending`.
/// - `DURABLE_LEADER_AGG_JOB_QUEUE_PUT`: Fetches the desired number of jobs from the front of the
///    queue.
/// - `DURABLE_LEADER_AGG_JOB_QUEUE_FINISH`: Removes the indicated job from the queue.
///
/// The schemea for data stored in instances of this DO is as follows:
///
/// ```text
///     agg_job/item/time/<time>/nonce/<nonce> -> String
/// ```
///
/// where `<time>` and `<nonce>` were generated by the `ReportsPending` instance at creation time.
/// The value stored is the unique name of the `ReportsPending` instance. Note that this schema
/// matches the ordinal generated by [`DurableOrdered::new_roughly_ordered`].
#[durable_object]
pub struct LeaderAggregationJobQueue {
    #[allow(dead_code)]
    state: State,
    env: Env,
    touched: bool,
}

#[durable_object]
impl DurableObject for LeaderAggregationJobQueue {
    fn new(state: State, env: Env) -> Self {
        Self {
            state,
            env,
            touched: false,
        }
    }

    async fn fetch(&mut self, mut req: Request) -> Result<Response> {
        let id_hex = self.state.id().to_string();
        ensure_garbage_collected!(req, self, id_hex, BINDING_DAP_LEADER_AGG_JOB_QUEUE);

        match (req.path().as_ref(), req.method()) {
            // Put a job (near) the back of the queue.
            //
            // Input: `agg_job: DurableOrdered<String>` (the `String` is the name of the
            // `ReportsPending` instance)
            (DURABLE_LEADER_AGG_JOB_QUEUE_PUT, Method::Post) => {
                let agg_job: DurableOrdered<String> = req.json().await?;
                agg_job.put(&self.state).await?;
                console_debug!(
                    "LeaderAggregationJobQueue: {} has been scheduled",
                    agg_job.as_ref(),
                );
                Response::from_json(&())
            }

            // Fetch the aggregation jobs at the fron tf the queue.
            //
            // Input: `max_agg_jobs: usize`,
            // Output: `Vec<String>` (the names of the `ReportsPending` instances from which to
            // drain reports)
            (DURABLE_LEADER_AGG_JOB_QUEUE_GET, Method::Post) => {
                let max_agg_jobs: usize = req.json().await?;
                let res: Vec<String> =
                    DurableOrdered::get_front(&self.state, "agg_job", max_agg_jobs)
                        .await?
                        .into_iter()
                        .map(|agg_job| agg_job.into_item())
                        .collect();

                console_debug!("agg job queue: {:?}", res);
                Response::from_json(&res)
            }

            // Remove a job from the queue.
            //
            // Input: `agg_job: DurableOrdered<String>` (the `String` is the name of the
            // `ReportsPending` instance that has become empty)
            (DURABLE_LEADER_AGG_JOB_QUEUE_FINISH, Method::Post) => {
                let agg_job: DurableOrdered<String> = req.json().await?;
                agg_job.delete(&self.state).await?;
                Response::from_json(&())
            }

            _ => Err(int_err(format!(
                "LeaderAggregationJobQueue: unexpected request: method={:?}; path={:?}",
                req.method(),
                req.path()
            ))),
        }
    }
}