Add Phase 1 frxd implementation with conformance test suite

src/relay.rs

@@ -0,0 +1,191 @@
+use std::collections::{HashMap, VecDeque};
+use std::net::SocketAddr;
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
+
+use anyhow::Result;
+use axum::extract::{Query, State};
+use axum::http::{StatusCode, header};
+use axum::response::{IntoResponse, Response};
+use axum::routing::{get, post};
+use axum::{Json, Router};
+use serde::Deserialize;
+use serde_json::json;
+use tokio::net::TcpListener;
+
+use crate::message::{Envelope, TYPE_AGGREGATE, TYPE_QUERY, TYPE_RESPONSE};
+
+pub const DEFAULT_CAPACITY: usize = 256;
+
+#[derive(Default)]
+struct MemberQueue {
+    items: VecDeque<(u64, Envelope)>,
+}
+
+struct Inner {
+    members: HashMap<String, MemberQueue>,
+    seq: u64,
+}
+
+pub struct Relay {
+    inner: Mutex<Inner>,
+    capacity: usize,
+}
+
+impl Relay {
+    pub fn new(capacity: usize) -> Arc<Self> {
+        Arc::new(Self {
+            inner: Mutex::new(Inner {
+                members: HashMap::new(),
+                seq: 0,
+            }),
+            capacity,
+        })
+    }
+
+    fn push(&self, targets: Option<&str>, envelope: Envelope) -> Result<usize, StatusCode> {
+        let mut inner = self.inner.lock().expect("relay lock");
+        if inner
+            .members
+            .values()
+            .any(|q| q.items.len() >= self.capacity)
+        {
+            return Err(StatusCode::TOO_MANY_REQUESTS);
+        }
+        inner.seq += 1;
+        let seq = inner.seq;
+        match targets {
+            Some(member) => {
+                let Some(queue) = inner.members.get_mut(member) else {
+                    return Err(StatusCode::NOT_FOUND);
+                };
+                queue.items.push_back((seq, envelope));
+                Ok(1)
+            }
+            None => {
+                let from = envelope.from.clone();
+                inner.members.entry(from).or_default();
+                let mut delivered = 0;
+                for queue in inner.members.values_mut() {
+                    queue.items.push_back((seq, envelope.clone()));
+                    delivered += 1;
+                }
+                Ok(delivered)
+            }
+        }
+    }
+}
+
+pub fn router(relay: Arc<Relay>) -> Router {
+    Router::new()
+        .route("/health", get(health))
+        .route("/v1/publish", post(publish))
+        .route("/v1/unicast", post(unicast))
+        .route("/v1/poll", get(poll))
+        .with_state(relay)
+}
+
+pub async fn run(listener: TcpListener, capacity: usize) -> Result<()> {
+    let relay = Relay::new(capacity);
+    axum::serve(listener, router(relay)).await?;
+    Ok(())
+}
+
+pub async fn bind(addr: &str) -> Result<(TcpListener, SocketAddr)> {
+    let listener = TcpListener::bind(addr).await?;
+    let local = listener.local_addr()?;
+    Ok((listener, local))
+}
+
+async fn health() -> &'static str {
+    "ok"
+}
+
+async fn publish(State(relay): State<Arc<Relay>>, Json(envelope): Json<Envelope>) -> Response {
+    if envelope.verify().is_err() {
+        return bad_request("invalid signature");
+    }
+    if envelope.msg_type != TYPE_QUERY {
+        return bad_request("relay carries broadcast queries only");
+    }
+    match relay.push(None, envelope) {
+        Ok(delivered) => (
+            StatusCode::ACCEPTED,
+            Json(json!({ "delivered": delivered })),
+        )
+            .into_response(),
+        Err(status) => backpressure(status),
+    }
+}
+
+#[derive(Deserialize)]
+struct UnicastParams {
+    to: String,
+}
+
+async fn unicast(
+    State(relay): State<Arc<Relay>>,
+    Query(params): Query<UnicastParams>,
+    Json(envelope): Json<Envelope>,
+) -> Response {
+    if envelope.verify().is_err() {
+        return bad_request("invalid signature");
+    }
+    if envelope.msg_type != TYPE_RESPONSE && envelope.msg_type != TYPE_AGGREGATE {
+        return bad_request("unicast carries responses and aggregates only");
+    }
+    match relay.push(Some(&params.to), envelope) {
+        Ok(_) => (StatusCode::OK, Json(json!({ "delivered": true }))).into_response(),
+        Err(StatusCode::NOT_FOUND) => (
+            StatusCode::NOT_FOUND,
+            Json(json!({ "error": "member not connected" })),
+        )
+            .into_response(),
+        Err(status) => backpressure(status),
+    }
+}
+
+#[derive(Deserialize)]
+struct PollParams {
+    member: String,
+    timeout_ms: Option<u64>,
+}
+
+async fn poll(State(relay): State<Arc<Relay>>, Query(params): Query<PollParams>) -> Response {
+    if params.member.is_empty() || hex::decode(&params.member).is_err() {
+        return bad_request("valid member key required");
+    }
+    let timeout = Duration::from_millis(params.timeout_ms.unwrap_or(25_000).min(60_000));
+    let deadline = Instant::now() + timeout;
+    loop {
+        let batch: Vec<Envelope> = {
+            let mut inner = relay.inner.lock().expect("relay lock");
+            let queue = inner.members.entry(params.member.clone()).or_default();
+            queue
+                .items
+                .drain(..)
+                .map(|(_, envelope)| envelope)
+                .collect()
+        };
+        if !batch.is_empty() {
+            return (StatusCode::OK, Json(json!({ "messages": batch }))).into_response();
+        }
+        if Instant::now() >= deadline {
+            return StatusCode::NO_CONTENT.into_response();
+        }
+        tokio::time::sleep(Duration::from_millis(25)).await;
+    }
+}
+
+fn bad_request(message: &str) -> Response {
+    (StatusCode::BAD_REQUEST, Json(json!({ "error": message }))).into_response()
+}
+
+fn backpressure(status: StatusCode) -> Response {
+    (
+        status,
+        [(header::RETRY_AFTER, "1")],
+        Json(json!({ "error": "transport backpressure" })),
+    )
+        .into_response()
+}