Add support for multiple OPRF instances in a single server, use duration strings instead of seconds for epoch durations

Author: DJAndries

Cargo.lock

@@ -10,6 +10,7 @@ edition = "2021"
 axum = "0.6.20"
 axum-prometheus = "0.4.0"
 base64 = "0.21.2"
+calendar-duration = { git = "https://github.com/brave-experiments/calendar-duration", branch = "init" }
 clap = { version = "4.3.24", features = ["derive"] }
 metrics-exporter-prometheus = "0.12.1"
 ppoprf = "0.3.1"

Cargo.toml

@@ -1,12 +1,14 @@
 //! STAR Randomness web service route implementation
 
-use axum::extract::{Json, State};
+use std::sync::RwLockReadGuard;
+
+use axum::extract::{Json, Path, State};
 use axum::http::StatusCode;
 use base64::prelude::{Engine as _, BASE64_STANDARD as BASE64};
 use serde::{Deserialize, Serialize};
-use tracing::debug;
+use tracing::{debug, instrument};
 
-use crate::OPRFState;
+use crate::state::{OPRFInstance, OPRFState};
 use ppoprf::ppoprf;
 
 /// Request format for the randomness endpoint
@@ -63,6 +65,8 @@ struct ErrorResponse {
 /// handling requests.
 #[derive(thiserror::Error, Debug)]
 pub enum Error {
+    #[error("instance '{0}' not found")]
+    InstanceNotFound(String),
     #[error("Couldn't lock state: RwLock poisoned")]
     LockFailure,
     #[error("Invalid point")]
@@ -93,6 +97,7 @@ impl axum::response::IntoResponse for Error {
     /// Construct an http response from our error type
     fn into_response(self) -> axum::response::Response {
         let code = match self {
+            Error::InstanceNotFound(_) => StatusCode::NOT_FOUND,
             // This indicates internal failure.
             Error::LockFailure => StatusCode::INTERNAL_SERVER_ERROR,
             // Other cases are the client's fault.
@@ -105,13 +110,28 @@ impl axum::response::IntoResponse for Error {
     }
 }
 
+type Result<T> = std::result::Result<T, Error>;
+
+fn get_server_from_state(
+    state: &OPRFState,
+    instance_name: String,
+) -> Result<RwLockReadGuard<'_, OPRFInstance>> {
+    Ok(state
+        .instances
+        .get(&instance_name)
+        .ok_or(Error::InstanceNotFound(instance_name))?
+        .read()?)
+}
+
 /// Process PPOPRF evaluation requests
-pub async fn randomness(
-    State(state): State<OPRFState>,
-    Json(request): Json<RandomnessRequest>,
-) -> Result<Json<RandomnessResponse>, Error> {
+#[instrument(skip(state, request))]
+async fn randomness(
+    state: OPRFState,
+    instance_name: String,
+    request: RandomnessRequest,
+) -> Result<Json<RandomnessResponse>> {
     debug!("recv: {request:?}");
-    let state = state.read()?;
+    let state = get_server_from_state(&state, instance_name)?;
     let epoch = request.epoch.unwrap_or(state.epoch);
     if epoch != state.epoch {
         return Err(Error::BadEpoch(epoch));
@@ -138,12 +158,29 @@ pub async fn randomness(
     Ok(Json(response))
 }
 
-/// Process PPOPRF epoch and key requests
-pub async fn info(
+/// Process PPOPRF evaluation requests using default instance
+pub async fn default_instance_randomness(
+    State(state): State<OPRFState>,
+    Json(request): Json<RandomnessRequest>,
+) -> Result<Json<RandomnessResponse>> {
+    let instance_name = state.default_instance.clone();
+    randomness(state, instance_name, request).await
+}
+
+/// Process PPOPRF evaluation requests using specific instance
+pub async fn specific_instance_randomness(
     State(state): State<OPRFState>,
-) -> Result<Json<InfoResponse>, Error> {
+    Path(instance_name): Path<String>,
+    Json(request): Json<RandomnessRequest>,
+) -> Result<Json<RandomnessResponse>> {
+    randomness(state, instance_name, request).await
+}
+
+/// Provide PPOPRF epoch and key metadata
+#[instrument(skip(state))]
+async fn info(state: OPRFState, instance_name: String) -> Result<Json<InfoResponse>> {
     debug!("recv: info request");
-    let state = state.read()?;
+    let state = get_server_from_state(&state, instance_name)?;
     let public_key = state.server.get_public_key().serialize_to_bincode()?;
     let public_key = BASE64.encode(public_key);
     let response = InfoResponse {
@@ -155,3 +192,17 @@ pub async fn info(
     debug!("send: {response:?}");
     Ok(Json(response))
 }
+
+/// Provide PPOPRF epoch and key metadata using default instance
+pub async fn default_instance_info(State(state): State<OPRFState>) -> Result<Json<InfoResponse>> {
+    let instance_name = state.default_instance.clone();
+    info(state, instance_name).await
+}
+
+/// Provide PPOPRF epoch and key metadata using specific instance
+pub async fn specific_instance_info(
+    State(state): State<OPRFState>,
+    Path(instance_name): Path<String>,
+) -> Result<Json<InfoResponse>> {
+    info(state, instance_name).await
+}

src/handler.rs

@@ -2,23 +2,23 @@
 
 use axum::{routing::get, routing::post, Router};
 use axum_prometheus::PrometheusMetricLayer;
+use calendar_duration::CalendarDuration;
 use clap::Parser;
 use metrics_exporter_prometheus::PrometheusHandle;
 use rlimit::Resource;
-use std::sync::{Arc, RwLock};
+use state::{OPRFServer, OPRFState};
 use tikv_jemallocator::Jemalloc;
-use time::format_description::well_known::Rfc3339;
 use time::OffsetDateTime;
 use tracing::{debug, info, metadata::LevelFilter};
 use tracing_subscriber::EnvFilter;
+use util::{assert_unique_names, parse_timestamp};
 
 #[global_allocator]
 static GLOBAL: Jemalloc = Jemalloc;
 
 mod handler;
 mod state;
-
-pub use state::OPRFState;
+mod util;
 
 #[cfg(test)]
 mod tests;
@@ -27,15 +27,22 @@ mod tests;
 const MAX_POINTS: usize = 1024;
 
 /// Command line switches
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
 #[command(author, version, about, long_about = None)]
 pub struct Config {
     /// Host and port to listen for http connections
     #[arg(long, default_value = "127.0.0.1:8080")]
     listen: String,
-    /// Duration of each randomness epoch
-    #[arg(long, default_value_t = 5)]
-    epoch_seconds: u32,
+    /// Name of OPRF instance contained in server. Multiple instances may be defined
+    /// by defining this switch multiple times. The first defined instance will
+    /// become the default instance.
+    #[arg(long = "instance-name", default_value = "main")]
+    instance_names: Vec<String>,
+    /// Duration of each randomness epoch. This switch may be defined multiple times
+    /// to set the epoch length for each respective instance, if multiple instances
+    /// are defined.
+    #[arg(long = "epoch-duration", value_name = "Duration string i.e. 1mon5h2s", default_values = ["5s"])]
+    epoch_durations: Vec<CalendarDuration>,
     /// First epoch tag to make available
     #[arg(long, default_value_t = 0)]
     first_epoch: u8,
@@ -56,20 +63,24 @@ pub struct Config {
     prometheus_listen: Option<String>,
 }
 
-/// Parse a timestamp given as a config option
-fn parse_timestamp(stamp: &str) -> Result<OffsetDateTime, &'static str> {
-    OffsetDateTime::parse(stamp, &Rfc3339).map_err(|_| "Try something like '2023-05-15T04:30:00Z'.")
-}
-
 /// Initialize an axum::Router for our web service
 /// Having this as a separate function makes testing easier.
 fn app(oprf_state: OPRFState) -> Router {
     Router::new()
         // Friendly default route to identify the site
         .route("/", get(|| async { "STAR randomness server\n" }))
-        // Main endpoints
-        .route("/randomness", post(handler::randomness))
-        .route("/info", get(handler::info))
+        // Endpoints for all instances
+        .route(
+            "/instances/:instance/randomness",
+            post(handler::specific_instance_randomness),
+        )
+        .route(
+            "/instances/:instance/info",
+            get(handler::specific_instance_info),
+        )
+        // Endpoints for default instance
+        .route("/randomness", post(handler::default_instance_randomness))
+        .route("/info", get(handler::default_instance_info))
         // Attach shared state
         .with_state(oprf_state)
         // Logging must come after active routes
@@ -126,22 +137,28 @@ async fn main() {
         increase_nofile_limit();
     }
 
-    // Oblivious function state
-    info!("initializing OPRF state...");
-    let server = state::OPRFServer::new(&config).expect("Could not initialize PPOPRF state");
-    info!("epoch now {}", server.epoch);
-    let oprf_state = Arc::new(RwLock::new(server));
+    assert_unique_names(&config.instance_names);
+    assert!(
+        !config.epoch_durations.iter().any(|d| d.is_zero()),
+        "all epoch lengths must be non-zero"
+    );
+    assert!(
+        !config.instance_names.is_empty(),
+        "at least one instance name must be defined"
+    );
+    assert!(
+        config.instance_names.len() == config.epoch_durations.len(),
+        "instance-name switch count must match epoch-seconds switch count"
+    );
 
     let metric_layer = config.prometheus_listen.as_ref().map(|listen| {
         let (layer, handle) = PrometheusMetricLayer::pair();
         start_prometheus_server(handle, listen.clone());
         layer
     });
 
-    // Spawn a background process to advance the epoch
-    info!("Spawning background epoch rotation task...");
-    let background_state = oprf_state.clone();
-    tokio::spawn(async move { state::epoch_loop(background_state, &config).await });
+    let oprf_state = OPRFServer::new(&config);
+    oprf_state.start_background_tasks(&config);
 
     // Set up routes and middleware
     info!("initializing routes...");