Support multiple quorums on a single LighthouseServer using gRPC metadata-based room assignment

Cargo.toml

@@ -9,8 +9,12 @@ askama = "0.12.1"
 atty = "0.2.14"
 axum = "0.7.7"
 chrono = "0.4.40"
+dashmap = "6.1"
 fern = {version = "0.7.1", features = ["colored"]}
+futures = "0.3"
 gethostname = "0.5.0"
+hyper = "0.14"
+http = "0.2"
 log = "0.4.22"
 prost = "0.13.3"
 prost-types = "0.13.3"
@@ -21,7 +25,9 @@ slog-stdlog = "4.1.1"
 stderrlog = "0.6.0"
 structopt = "0.3.26"
 tokio = {version = "1.40.0", features = ["full", "test-util", "tracing", "macros", "rt-multi-thread"] }
+tokio-stream = "0.1"
 tonic = "0.12.2"
+tower = "0.4"
 
 [build-dependencies]
 tonic-build = "0.12.2"

src/bin/lighthouse.rs

@@ -4,8 +4,11 @@
 // This source code is licensed under the BSD-style license found in the
 // LICENSE file in the root directory of this source tree.
 
+use std::net::SocketAddr;
 use structopt::StructOpt;
-use torchft::lighthouse::{Lighthouse, LighthouseOpt};
+use tonic::transport::Server;
+use torchft::lighthouse::LighthouseOpt;
+use torchft::router::Router;
 
 #[tokio::main(flavor = "multi_thread", worker_threads = 4)]
 async fn main() {
@@ -17,7 +20,12 @@ async fn main() {
         .unwrap();
 
     let opt = LighthouseOpt::from_args();
-    let lighthouse = Lighthouse::new(opt).await.unwrap();
+    let router = Router::new(opt.clone());
+    let addr: SocketAddr = opt.bind.parse().expect("invalid --bind address");
 
-    lighthouse.run().await.unwrap();
+    Server::builder()
+        .add_service(router)
+        .serve(addr)
+        .await
+        .unwrap();
 }

src/interceptor.rs

@@ -0,0 +1,23 @@
+use tonic::{service::Interceptor, metadata::MetadataValue, Request, Status};
+
+/// Attaches user-assigned room-id header to every outbound RPC
+#[derive(Clone)]
+pub struct RoomIdInterceptor {
+    room: String,
+}
+
+impl RoomIdInterceptor {
+    pub fn new(room: String) -> Self {
+        Self { room }
+    }
+}
+
+impl Interceptor for RoomIdInterceptor {
+    fn call(&mut self, mut req: Request<()>) -> Result<Request<()>, Status> {
+        req.metadata_mut().insert(
+            crate::router::ROOM_ID_HEADER,
+            MetadataValue::try_from(self.room.as_str()).expect("ascii header"),
+        );
+        Ok(req)
+    }
+}

src/interceptor.rs~

@@ -0,0 +1,12 @@
+use tonic::{Request, Status, service::Interceptor};
+use tonic::metadata::MetadataValue;
+
+pub fn room_id_interceptor(room: String) -> impl Interceptor {
+    move |mut req: Request<()>| {
+        req.metadata_mut().insert(
+            crate::router::ROOM_ID_HEADER,
+            MetadataValue::try_from(room.as_str()).expect("ascii header"),
+        );
+        Ok(req) // returning Err(Status) would cancel the call :contentReference[oaicite:0]{index=0}
+    }
+}

src/lib.rs

@@ -4,24 +4,31 @@
 // This source code is licensed under the BSD-style license found in the
 // LICENSE file in the root directory of this source tree.
 
+mod interceptor;
 pub mod lighthouse;
 pub mod manager;
 mod net;
 mod retry;
+pub mod router;
 mod timeout;
 
+pub use crate::router::Router;
+
 use anyhow::Result;
 use atty::Stream;
 use core::time::Duration;
 use pyo3::exceptions::{PyRuntimeError, PyTimeoutError};
 use std::cmp;
 use std::env;
+use std::net::SocketAddr;
 use std::sync::Arc;
 use std::thread::available_parallelism;
 use structopt::StructOpt;
 use tokio::runtime::Runtime;
 use tokio::task::JoinHandle;
-use tonic::transport::Channel;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::service::interceptor::InterceptedService;
+use tonic::transport::{Channel, Endpoint};
 use tonic::Status;
 
 use chrono::Local;
@@ -32,6 +39,7 @@ pub mod torchftpb {
     tonic::include_proto!("torchft");
 }
 
+use crate::interceptor::RoomIdInterceptor;
 use crate::torchftpb::lighthouse_service_client::LighthouseServiceClient;
 use crate::torchftpb::manager_service_client::ManagerServiceClient;
 use crate::torchftpb::{
@@ -339,9 +347,13 @@ fn lighthouse_main(py: Python<'_>) -> PyResult<()> {
 }
 
 async fn lighthouse_main_async(opt: lighthouse::LighthouseOpt) -> Result<()> {
-    let lighthouse = lighthouse::Lighthouse::new(opt).await?;
+    let router = Router::new(opt.clone());
+    let addr: SocketAddr = opt.bind.parse()?;
 
-    lighthouse.run().await?;
+    tonic::transport::Server::builder()
+        .add_service(router)
+        .serve(addr)
+        .await?;
 
     Ok(())
 }
@@ -477,28 +489,39 @@ fn convert_quorum(py: Python, q: &torchftpb::Quorum) -> PyResult<Quorum> {
 ///     connect_timeout (timedelta): The timeout for connecting to the lighthouse server.
 #[pyclass]
 struct LighthouseClient {
-    client: LighthouseServiceClient<Channel>,
+    client: LighthouseServiceClient<InterceptedService<Channel, RoomIdInterceptor>>,
     runtime: Runtime,
 }
 
 #[pymethods]
 impl LighthouseClient {
-    #[pyo3(signature = (addr, connect_timeout))]
+    #[pyo3(signature = (addr, connect_timeout, room_id = None))]
     #[new]
-    fn new(py: Python<'_>, addr: String, connect_timeout: Duration) -> PyResult<Self> {
+    fn new(
+        py: Python<'_>,
+        addr: String,
+        connect_timeout: Duration,
+        room_id: Option<String>,
+    ) -> PyResult<Self> {
         py.allow_threads(move || {
             let runtime = tokio::runtime::Builder::new_multi_thread()
                 .worker_threads(num_threads())
                 .thread_name("torchft-lhclnt")
                 .enable_all()
                 .build()?;
-            let client = runtime
-                .block_on(manager::lighthouse_client_new(addr, connect_timeout))
+
+            let endpoint = Endpoint::from_shared(addr.clone())
                 .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
-            Ok(Self {
-                client: client,
-                runtime: runtime,
-            })
+            let channel = runtime
+                .block_on(endpoint.connect_timeout(connect_timeout).connect())
+                .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+
+            let interceptor =
+                RoomIdInterceptor::new(room_id.unwrap_or_else(|| "default".to_owned()));
+
+            let client = LighthouseServiceClient::with_interceptor(channel, interceptor);
+
+            Ok(Self { client, runtime })
         })
     }
 
@@ -603,7 +626,7 @@ impl LighthouseClient {
 ///     heartbeat_timeout_ms (int): The timeout for heartbeats.
 #[pyclass]
 struct LighthouseServer {
-    lighthouse: Arc<lighthouse::Lighthouse>,
+    bind: String,
     handle: JoinHandle<Result<()>>,
     _runtime: Runtime,
 }
@@ -631,19 +654,31 @@ impl LighthouseServer {
                 .enable_all()
                 .build()?;
 
-            let lighthouse = rt
-                .block_on(lighthouse::Lighthouse::new(lighthouse::LighthouseOpt {
-                    bind: bind,
-                    min_replicas: min_replicas,
-                    join_timeout_ms: join_timeout_ms,
-                    quorum_tick_ms: quorum_tick_ms,
-                    heartbeat_timeout_ms: heartbeat_timeout_ms,
-                }))
-                .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+            let opt = lighthouse::LighthouseOpt {
+                bind: bind.clone(),
+                min_replicas,
+                join_timeout_ms,
+                quorum_tick_ms,
+                heartbeat_timeout_ms,
+            };
+
+            let listener = rt.block_on(tokio::net::TcpListener::bind(&bind))?;
+            let bound_sock = listener.local_addr()?;
+            let bound = format!("http://{}", bound_sock);
+            let incoming = TcpListenerStream::new(listener);
+            let router = Router::new(opt.clone());
+
+            let handle = rt.spawn(async move {
+                tonic::transport::Server::builder()
+                    .add_service(router)
+                    .serve_with_incoming(incoming)
+                    .await
+                    .map_err(|e: tonic::transport::Error| anyhow::anyhow!(e))
+            });
 
             Ok(Self {
-                handle: rt.spawn(lighthouse.clone().run()),
-                lighthouse: lighthouse,
+                bind: bound,
+                handle,
                 _runtime: rt,
             })
         })
@@ -654,7 +689,7 @@ impl LighthouseServer {
     /// Returns:
     ///    str: The address of the lighthouse server.
     fn address(&self) -> PyResult<String> {
-        Ok(self.lighthouse.address().to_string())
+        Ok(self.bind.clone())
     }
 
     /// shutdown shuts down the lighthouse server.

src/lighthouse.rs

@@ -83,7 +83,7 @@ impl ChangeLogger {
     }
 }
 
-#[derive(StructOpt, Debug)]
+#[derive(StructOpt, Debug, Clone)]
 #[structopt()]
 pub struct LighthouseOpt {
     // bind is the address to bind the server to.

src/router.rs

@@ -0,0 +1,110 @@
+use std::{
+    convert::Infallible,
+    future::Future,
+    pin::Pin,
+    sync::Arc,
+    task::{Context, Poll},
+};
+
+use dashmap::{mapref::entry::Entry, DashMap};
+use futures::FutureExt;
+use tonic::{
+    body::BoxBody,
+    codegen::http::{HeaderMap, Request, Response}, // http-0.2 types
+    server::NamedService,
+};
+use tower::Service;
+
+use crate::{
+    lighthouse::{Lighthouse, LighthouseOpt},
+    torchftpb::lighthouse_service_server::LighthouseServiceServer,
+};
+
+/// Metadata header recognised by both client interceptor and this router.
+pub const ROOM_ID_HEADER: &str = "room-id";
+
+/// gRPC server for a single room (inner state = `Arc<Lighthouse>`).
+type GrpcSvc = LighthouseServiceServer<Arc<Lighthouse>>;
+
+#[derive(Clone)]
+pub struct Router {
+    rooms: Arc<DashMap<String, Arc<GrpcSvc>>>,
+    tmpl_opt: LighthouseOpt,
+}
+
+impl Router {
+    pub fn new(tmpl_opt: LighthouseOpt) -> Self {
+        Self {
+            rooms: Arc::new(DashMap::new()),
+            tmpl_opt,
+        }
+    }
+
+    fn room_id(hdrs: &HeaderMap) -> &str {
+        hdrs.get(ROOM_ID_HEADER)
+            .and_then(|v| v.to_str().ok())
+            .unwrap_or("default")
+    }
+
+    async fn room_service(
+        rooms: Arc<DashMap<String, Arc<GrpcSvc>>>,
+        tmpl: LighthouseOpt,
+        id: &str,
+    ) -> Arc<GrpcSvc> {
+        if let Some(svc) = rooms.get(id) {
+            return svc.clone();
+        }
+
+        // Build room state once.
+        let lh = Lighthouse::new(tmpl.clone())
+            .await
+            .expect("failed to create Lighthouse");
+
+        let svc_new = Arc::new(LighthouseServiceServer::new(lh));
+
+        match rooms.entry(id.to_owned()) {
+            Entry::Occupied(e) => e.get().clone(),
+            Entry::Vacant(v) => {
+                v.insert(svc_new.clone());
+                svc_new
+            }
+        }
+    }
+}
+
+// Tower::Service implementation
+impl Service<Request<BoxBody>> for Router {
+    type Response = Response<BoxBody>;
+    type Error = Infallible;
+    type Future =
+        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
+
+    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn call(&mut self, req: Request<BoxBody>) -> Self::Future {
+        let rooms = self.rooms.clone();
+        let tmpl = self.tmpl_opt.clone();
+        let room = Self::room_id(req.headers()).to_owned();
+
+        async move {
+            let svc_arc = Self::room_service(rooms, tmpl, &room).await;
+
+            // `Arc<GrpcSvc>` itself isn’t a Service; clone the inner value.
+            let mut svc = (*svc_arc).clone();
+            let resp = svc
+                .call(req)
+                .await
+                .map_err(|_e| -> Infallible { unreachable!() })?;
+
+            Ok(resp)
+        }
+        .boxed()
+    }
+}
+
+// Forward tonic’s NamedService marker
+impl NamedService for Router {
+    const NAME: &'static str = <GrpcSvc as NamedService>::NAME;
+}