fix(xds): add mutiple async connections to the server

hanshal101 · hanshal101 · commit 56e00120f5ab · 2025-08-12T16:21:29.000+05:30
Signed-off-by: hanshal101 &lt;hanshalmehta10@gmail.com&gt;
diff --git a/orion-xds/examples/server.rs b/orion-xds/examples/server.rs
@@ -11,13 +11,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .with(tracing_subscriber::fmt::layer())
         .init();
 
-    let (delta_resource_tx, delta_resources_rx) = tokio::sync::mpsc::channel(100);
-    let (_stream_resource_tx, stream_resources_rx) = tokio::sync::mpsc::channel(100);
+    let (delta_resource_tx, _delta_resources_rx) = tokio::sync::broadcast::channel::<ServerAction>(100);
+    let (stream_resource_tx, _stream_resources_rx) = tokio::sync::broadcast::channel::<ServerAction>(100);
     let addr = "127.0.0.1:50051".parse()?;
+    let delta_tx_clone = delta_resource_tx.clone();
+    let stream_tx_clone = stream_resource_tx.clone();
 
     let grpc_server = tokio::spawn(async move {
         info!("Server started");
-        let res = start_aggregate_server(addr, delta_resources_rx, stream_resources_rx).await;
+        let res = start_aggregate_server(addr, delta_tx_clone, stream_tx_clone).await;
         info!("Server stopped {res:?}");
     });
     tokio::time::sleep(std::time::Duration::from_secs(10)).await;
@@ -37,7 +39,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             info!("Adding cluster {cluster_id}");
             let cluster_resource = resources::create_cluster_resource(&cluster);
 
-            if delta_resource_tx.send(ServerAction::Add(cluster_resource.clone())).await.is_err() {
+            if delta_resource_tx.send(ServerAction::Add(cluster_resource.clone())).is_err() {
                 break;
             };
             tokio::time::sleep(Duration::from_secs(5)).await;
@@ -50,13 +52,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             );
             let listener_resource = resources::create_listener_resource(&listener);
             info!("Adding listener {listener_resource:?}");
-            if delta_resource_tx.send(ServerAction::Add(listener_resource)).await.is_err() {
+            if delta_resource_tx.send(ServerAction::Add(listener_resource)).is_err() {
                 break;
             };
             tokio::time::sleep(Duration::from_secs(15)).await;
 
             info!("Removing cluster {cluster_id}");
-            if delta_resource_tx.send(ServerAction::Remove(cluster_resource)).await.is_err() {
+            if delta_resource_tx.send(ServerAction::Remove(cluster_resource)).is_err() {
                 break;
             };
             tokio::time::sleep(Duration::from_secs(5)).await;
@@ -69,7 +71,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             );
             let listener_resource = resources::create_listener_resource(&listener);
             info!("Removing listener {listener_resource:?}");
-            if delta_resource_tx.send(ServerAction::Remove(listener_resource)).await.is_err() {
+            if delta_resource_tx.send(ServerAction::Remove(listener_resource)).is_err() {
                 break;
             };
         }
diff --git a/orion-xds/examples/server_routes_and_loads.rs b/orion-xds/examples/server_routes_and_loads.rs
@@ -14,13 +14,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .with(tracing_subscriber::fmt::layer())
         .init();
 
-    let (delta_resource_tx, delta_resources_rx) = tokio::sync::mpsc::channel(100);
-    let (_stream_resource_tx, stream_resources_rx) = tokio::sync::mpsc::channel(100);
+    let (delta_resource_tx, _delta_resources_rx) = tokio::sync::broadcast::channel::<ServerAction>(100);
+    let (stream_resource_tx, _stream_resources_rx) = tokio::sync::broadcast::channel::<ServerAction>(100);
     let addr = "127.0.0.1:50051".parse()?;
+    let delta_tx_clone = delta_resource_tx.clone();
+    let stream_tx_clone = stream_resource_tx.clone();
 
     let grpc_server = tokio::spawn(async move {
         info!("Server started");
-        let res = start_aggregate_server(addr, delta_resources_rx, stream_resources_rx).await;
+        let res = start_aggregate_server(addr, delta_tx_clone, stream_tx_clone).await;
         info!("Server stopped {res:?}");
     });
     tokio::time::sleep(std::time::Duration::from_secs(10)).await;
@@ -38,7 +40,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         info!("Adding Cluster Load Assignment for cluster {cluster_id}");
         let load_assigment_resource = resources::create_load_assignment_resource(&cluster_id, &cla);
 
-        if delta_resource_tx.send(ServerAction::Add(load_assigment_resource.clone())).await.is_err() {
+        if delta_resource_tx.send(ServerAction::Add(load_assigment_resource.clone())).is_err() {
             return;
         };
         tokio::time::sleep(Duration::from_secs(5)).await;
@@ -49,20 +51,20 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         let route_configuration_resource =
             resources::create_route_configuration_resource(&route_id, &route_configuration);
 
-        if delta_resource_tx.send(ServerAction::Add(route_configuration_resource.clone())).await.is_err() {
+        if delta_resource_tx.send(ServerAction::Add(route_configuration_resource.clone())).is_err() {
             return;
         };
 
         tokio::time::sleep(Duration::from_secs(15)).await;
 
         info!("Removing cluster load assignment {cluster_id}");
-        if delta_resource_tx.send(ServerAction::Remove(load_assigment_resource)).await.is_err() {
+        if delta_resource_tx.send(ServerAction::Remove(load_assigment_resource)).is_err() {
             return;
         };
         tokio::time::sleep(Duration::from_secs(5)).await;
 
         info!("Removing route configuration {route_id}");
-        if delta_resource_tx.send(ServerAction::Remove(route_configuration_resource)).await.is_err() {
+        if delta_resource_tx.send(ServerAction::Remove(route_configuration_resource)).is_err() {
             return;
         };
         tokio::time::sleep(Duration::from_secs(5)).await;
diff --git a/orion-xds/examples/server_secret_rotation.rs b/orion-xds/examples/server_secret_rotation.rs
@@ -18,13 +18,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .with(tracing_subscriber::fmt::layer())
         .init();
 
-    let (delta_resource_tx, delta_resources_rx) = tokio::sync::mpsc::channel(100);
-    let (_stream_resource_tx, stream_resources_rx) = tokio::sync::mpsc::channel(100);
+    let (delta_resource_tx, _delta_resources_rx) = tokio::sync::broadcast::channel::<ServerAction>(100);
+    let (stream_resource_tx, _stream_resources_rx) = tokio::sync::broadcast::channel::<ServerAction>(100);
     let addr = "127.0.0.1:50051".parse()?;
+    let delta_tx_clone = delta_resource_tx.clone();
+    let stream_tx_clone = stream_resource_tx.clone();
 
     let grpc_server = tokio::spawn(async move {
         info!("Server started");
-        let res = start_aggregate_server(addr, delta_resources_rx, stream_resources_rx).await;
+        let res = start_aggregate_server(addr, delta_tx_clone, stream_tx_clone).await;
         info!("Server stopped {res:?}");
     });
     tokio::time::sleep(std::time::Duration::from_secs(10)).await;
@@ -55,7 +57,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         info!("Adding downstream secret {secret_id}");
         let secret_resource = resources::create_secret_resource(secret_id, &secret);
 
-        if delta_resource_tx.send(ServerAction::Add(secret_resource.clone())).await.is_err() {
+        if delta_resource_tx.send(ServerAction::Add(secret_resource.clone())).is_err() {
             return;
         };
 
@@ -77,7 +79,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         info!("Adding upstream secret {secret_id}");
         let secret_resource = resources::create_secret_resource(secret_id, &secret);
 
-        if delta_resource_tx.send(ServerAction::Add(secret_resource.clone())).await.is_err() {
+        if delta_resource_tx.send(ServerAction::Add(secret_resource.clone())).is_err() {
             return;
         };
 
diff --git a/orion-xds/examples/server_secret_rotation_simple.rs b/orion-xds/examples/server_secret_rotation_simple.rs
@@ -4,7 +4,10 @@ use orion_data_plane_api::envoy_data_plane_api::envoy::{
     config::core::v3::{data_source::Specifier, DataSource},
     extensions::transport_sockets::tls::v3::{secret, CertificateValidationContext},
 };
-use orion_xds::xds::{resources, server::start_aggregate_server};
+use orion_xds::xds::{
+    resources,
+    server::{start_aggregate_server, ServerAction},
+};
 use tracing::info;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 
@@ -15,13 +18,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .with(tracing_subscriber::fmt::layer())
         .init();
 
-    let (_, delta_resources_rx) = tokio::sync::mpsc::channel(100);
-    let (_stream_resource_tx, stream_resources_rx) = tokio::sync::mpsc::channel(100);
+    let (delta_resource_tx, _delta_resources_rx) = tokio::sync::broadcast::channel::<ServerAction>(100);
+    let (stream_resource_tx, _stream_resources_rx) = tokio::sync::broadcast::channel::<ServerAction>(100);
     let addr = "127.0.0.1:50051".parse()?;
+    let delta_tx_clone = delta_resource_tx.clone();
+    let stream_tx_clone = stream_resource_tx.clone();
 
     let grpc_server = tokio::spawn(async move {
         info!("Server started");
-        let res = start_aggregate_server(addr, delta_resources_rx, stream_resources_rx).await;
+        let res = start_aggregate_server(addr, delta_tx_clone, stream_tx_clone).await;
         info!("Server stopped {res:?}");
     });
     tokio::time::sleep(std::time::Duration::from_secs(10)).await;
diff --git a/orion-xds/src/xds/server.rs b/orion-xds/src/xds/server.rs
@@ -20,20 +20,23 @@
 
 use std::{net::SocketAddr, pin::Pin};
 
-use atomic_take::AtomicTake;
 use orion_data_plane_api::envoy_data_plane_api::envoy::service::discovery::v3::{
     aggregated_discovery_service_server::{AggregatedDiscoveryService, AggregatedDiscoveryServiceServer},
     DeltaDiscoveryRequest, DeltaDiscoveryResponse, DiscoveryRequest, DiscoveryResponse, Resource, ResourceName,
 };
 use orion_data_plane_api::envoy_data_plane_api::tonic::{
     self, transport::Server, IntoStreamingRequest, Response, Status,
 };
-use tokio::sync::mpsc::{self, Receiver};
+use tokio::sync::{
+    broadcast::Sender,
+    mpsc::{self},
+};
 use tokio_stream::{wrappers::ReceiverStream, Stream, StreamExt};
 use tracing::info;
 
 use crate::xds::{self, model::XdsError};
 
+#[derive(Debug, Clone)]
 pub enum ServerAction {
     Add(Resource),
     Remove(Resource),
@@ -42,16 +45,13 @@ pub enum ServerAction {
 pub type ResourceAction = ServerAction;
 #[derive(Debug)]
 pub struct AggregateServer {
-    delta_resources_rx: AtomicTake<Receiver<ServerAction>>,
-    stream_resources_rx: AtomicTake<Receiver<ServerAction>>,
+    delta_resources_rx: Sender<ServerAction>,
+    stream_resources_rx: Sender<ServerAction>,
 }
 
 impl AggregateServer {
-    pub fn new(delta_resources_rx: Receiver<ServerAction>, stream_resources_rx: Receiver<ServerAction>) -> Self {
-        Self {
-            delta_resources_rx: AtomicTake::new(delta_resources_rx),
-            stream_resources_rx: AtomicTake::new(stream_resources_rx),
-        }
+    pub fn new(delta_resources_rx: Sender<ServerAction>, stream_resources_rx: Sender<ServerAction>) -> Self {
+        Self { delta_resources_rx, stream_resources_rx }
     }
 }
 
@@ -70,10 +70,9 @@ impl AggregatedDiscoveryService for AggregateServer {
         info!("\tclient connected from: {:?}", req.remote_addr());
 
         let (tx, rx) = mpsc::channel(128);
-        let mut resources_rx =
-            self.stream_resources_rx.take().ok_or(Status::internal("Resource stream is unavailable"))?;
+        let mut resources_rx = self.stream_resources_rx.subscribe();
         tokio::spawn(async move {
-            while let Some(action) = resources_rx.recv().await {
+            while let Ok(action) = resources_rx.recv().await {
                 let item = match action {
                     xds::server::ServerAction::Add(resource) => {
                         let Some(resource) = resource.resource else {
@@ -136,9 +135,9 @@ impl AggregatedDiscoveryService for AggregateServer {
         // spawn and channel are required if you want handle "disconnect" functionality
         // the `out_stream` will not be polled after client disconnect
         let (tx, rx) = mpsc::channel(128);
-        let mut resources_rx = self.delta_resources_rx.take().ok_or(Status::internal("Delta stream is unavailable"))?;
+        let mut resources_rx = self.delta_resources_rx.subscribe();
         tokio::spawn(async move {
-            while let Some(action) = resources_rx.recv().await {
+            while let Ok(action) = resources_rx.recv().await {
                 let item = match action {
                     xds::server::ServerAction::Add(r) => {
                         let Some(ref resource) = r.resource else {
@@ -198,8 +197,8 @@ impl AggregatedDiscoveryService for AggregateServer {
 
 pub async fn start_aggregate_server(
     addr: SocketAddr,
-    delta_resources_rx: Receiver<ResourceAction>,
-    stream_resources_rx: Receiver<ResourceAction>,
+    delta_resources_rx: Sender<ResourceAction>,
+    stream_resources_rx: Sender<ResourceAction>,
 ) -> Result<(), XdsError> {
     info!("Server started {addr:?}");
     let server = AggregateServer::new(delta_resources_rx, stream_resources_rx);
