feat: allow custom indexer

Author: lherman-cs

src/main.rs

@@ -3,8 +3,8 @@ use tokio_util::sync::CancellationToken;
 use tracing::{info, level_filters::LevelFilter};
 use tracing_subscriber::fmt::format::FmtSpan;
 
-use pulsebeam_server_foss::proto::signaling_server::SignalingServer;
 use pulsebeam_server_foss::server::Server;
+use pulsebeam_server_foss::{manager::IndexManager, proto::signaling_server::SignalingServer};
 use std::time::Duration;
 use tonic::service::LayerExt;
 use tower_http::cors::{AllowOrigin, CorsLayer};
@@ -29,15 +29,15 @@ async fn main() -> anyhow::Result<()> {
     let token = CancellationToken::new();
     let (mut health_reporter, health_service) = tonic_health::server::health_reporter();
     health_reporter
-        .set_serving::<SignalingServer<crate::Server>>()
+        .set_serving::<SignalingServer<crate::Server<IndexManager>>>()
         .await;
     // https://github.com/hyperium/tonic/discussions/1784
     // switch to v1 after tooling supports it
     let reflector = tonic_reflection::server::Builder::configure()
         .register_encoded_file_descriptor_set(pulsebeam_server_foss::proto::FILE_DESCRIPTOR_SET)
         .register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
         .build_v1()?;
-    let server = Server::spawn(token, CONNECTION_CAPACITY);
+    let server = Server::spawn_default(token, CONNECTION_CAPACITY);
     let grpc_server = tower::ServiceBuilder::new()
         .layer(cors)
         .layer(tonic_web::GrpcWebLayer::new())

src/manager.rs

@@ -93,6 +93,26 @@ pub enum ConnEvent {
     Removed(PeerInfo),
 }
 
+pub trait Indexer: Send + Sync + 'static {
+    fn select(&self, range: impl RangeBounds<PeerInfo>) -> Vec<(PeerInfo, PeerStats)>;
+    fn select_one(&self, range: impl RangeBounds<PeerInfo>) -> Option<PeerInfo>;
+
+    fn select_group(&self, group_id: GroupId) -> Vec<(PeerInfo, PeerStats)> {
+        let start = PeerInfo {
+            group_id: group_id.clone(),
+            peer_id: "".to_string(),
+            conn_id: u32::MIN,
+        };
+
+        let end = PeerInfo {
+            group_id,
+            peer_id: "~".to_string(),
+            conn_id: u32::MAX,
+        };
+        self.select(start..=end)
+    }
+}
+
 #[derive(Clone)]
 pub struct IndexManager {
     state: Arc<RwLock<IndexManagerState>>,
@@ -108,6 +128,27 @@ impl Default for IndexManager {
     }
 }
 
+impl Indexer for IndexManager {
+    fn select(&self, range: impl RangeBounds<PeerInfo>) -> Vec<(PeerInfo, PeerStats)> {
+        let state = self.state.read();
+        let mut result = Vec::new();
+        for (k, v) in state.index.range(range) {
+            result.push((k.clone(), v.clone()))
+        }
+        result
+    }
+
+    fn select_one(&self, range: impl RangeBounds<PeerInfo>) -> Option<PeerInfo> {
+        let state = self.state.read();
+        // pick the youngest connection
+        let found = state
+            .index
+            .range(range)
+            .max_by_key(|(_, p)| p.inserted_at)?;
+        Some(found.0.clone())
+    }
+}
+
 impl IndexManager {
     pub async fn run_until_cancelled_owned(self, mut event_ch: mpsc::UnboundedReceiver<ConnEvent>) {
         tracing::info!("spawned index worker");
@@ -128,40 +169,6 @@ impl IndexManager {
             }
         }
     }
-
-    pub fn select(&self, range: impl RangeBounds<PeerInfo>) -> Vec<(PeerInfo, PeerStats)> {
-        let state = self.state.read();
-        let mut result = Vec::new();
-        for (k, v) in state.index.range(range) {
-            result.push((k.clone(), v.clone()))
-        }
-        result
-    }
-
-    pub fn select_one(&self, range: impl RangeBounds<PeerInfo>) -> Option<PeerInfo> {
-        let state = self.state.read();
-        // pick the youngest connection
-        let found = state
-            .index
-            .range(range)
-            .max_by_key(|(_, p)| p.inserted_at)?;
-        Some(found.0.clone())
-    }
-
-    pub fn select_group(&self, group_id: GroupId) -> Vec<(PeerInfo, PeerStats)> {
-        let start = PeerInfo {
-            group_id: group_id.clone(),
-            peer_id: "".to_string(),
-            conn_id: u32::MIN,
-        };
-
-        let end = PeerInfo {
-            group_id,
-            peer_id: "~".to_string(),
-            conn_id: u32::MAX,
-        };
-        self.select(start..=end)
-    }
 }
 
 #[derive(Clone, Debug, Serialize)]

src/server.rs

@@ -9,21 +9,21 @@ use tokio_util::sync::CancellationToken;
 use tracing::field::valuable;
 use valuable::Enumerable;
 
-use crate::manager::{IndexManager, Manager};
+use crate::manager::{IndexManager, Indexer, Manager};
 const RESERVED_CONN_ID_DISCOVERY: u32 = 0;
 const RECV_STREAM_BUFFER: usize = 8;
 const KEEP_ALIVE_INTERVAL: Duration = Duration::from_secs(45);
 
 #[derive(Clone)]
-pub struct Server {
+pub struct Server<I> {
     pub manager: Manager,
-    pub index: IndexManager,
+    pub index: I,
 }
 
 pub type MessageStream = Pin<Box<dyn Stream<Item = proto::Message> + Send>>;
 
-impl Server {
-    pub fn spawn(token: CancellationToken, capacity: u64) -> Self {
+impl Server<IndexManager> {
+    pub fn spawn_default(token: CancellationToken, capacity: u64) -> Self {
         let event_ch = mpsc::unbounded_channel();
         let manager = Manager::new(capacity, event_ch.0);
         let index = IndexManager::default();
@@ -35,7 +35,9 @@ impl Server {
         }
         Self { manager, index }
     }
+}
 
+impl<I: Indexer> Server<I> {
     pub fn insert_recv_stream(&self, src: PeerInfo) -> MessageStream {
         let conn = self.manager.allocate(src);
         let payload_stream = ReceiverStream::new(conn);
@@ -55,7 +57,7 @@ impl Server {
 pub type RecvStream = Pin<Box<dyn Stream<Item = Result<proto::RecvResp, tonic::Status>> + Send>>;
 
 #[tonic::async_trait]
-impl Signaling for Server {
+impl<I: Indexer> Signaling for Server<I> {
     async fn prepare(
         &self,
         _req: tonic::Request<proto::PrepareReq>,
@@ -185,6 +187,8 @@ impl Signaling for Server {
 mod test {
     use std::iter::zip;
 
+    use crate::manager::IndexManager;
+
     use super::*;
     use proto::*;
 
@@ -220,8 +224,8 @@ mod test {
             .await
     }
 
-    fn setup() -> (Server, PeerInfo, PeerInfo) {
-        let s = Server::spawn(CancellationToken::new(), 65536);
+    fn setup() -> (Server<IndexManager>, PeerInfo, PeerInfo) {
+        let s = Server::spawn_default(CancellationToken::new(), 65536);
         let peer1 = PeerInfo {
             group_id: String::from("default"),
             peer_id: String::from("peer1"),