[Misc] add service discovery for sgl router

sgl-router/src/server.rs

@@ -1,26 +1,30 @@
 use crate::logging::{self, LoggingConfig};
 use crate::router::PolicyConfig;
 use crate::router::Router;
+use crate::service_discovery::{start_service_discovery, ServiceDiscoveryConfig};
 use actix_web::{
     error, get, post, web, App, Error, HttpRequest, HttpResponse, HttpServer, Responder,
 };
 use bytes::Bytes;
 use futures_util::StreamExt;
+use reqwest::Client;
 use std::collections::HashMap;
 use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
 use std::time::Duration;
-use tracing::{info, Level};
+use tokio::spawn;
+use tracing::{error, info, warn, Level};
 
 #[derive(Debug)]
 pub struct AppState {
     router: Router,
-    client: reqwest::Client,
+    client: Client,
 }
 
 impl AppState {
     pub fn new(
         worker_urls: Vec<String>,
-        client: reqwest::Client,
+        client: Client,
         policy_config: PolicyConfig,
     ) -> Result<Self, String> {
         // Create router based on policy
@@ -149,6 +153,7 @@ pub struct ServerConfig {
     pub verbose: bool,
     pub max_payload_size: usize,
     pub log_dir: Option<String>,
+    pub service_discovery_config: Option<ServiceDiscoveryConfig>,
 }
 
 pub async fn startup(config: ServerConfig) -> std::io::Result<()> {
@@ -180,7 +185,15 @@ pub async fn startup(config: ServerConfig) -> std::io::Result<()> {
         config.max_payload_size / (1024 * 1024)
     );
 
-    let client = reqwest::Client::builder()
+    // Log service discovery status
+    if let Some(service_discovery_config) = &config.service_discovery_config {
+        info!("🚧 Service discovery enabled");
+        info!("🚧 Selector: {:?}", service_discovery_config.selector);
+    } else {
+        info!("🚧 Service discovery disabled");
+    }
+
+    let client = Client::builder()
         .pool_idle_timeout(Some(Duration::from_secs(50)))
         .build()
         .expect("Failed to create HTTP client");
@@ -194,6 +207,30 @@ pub async fn startup(config: ServerConfig) -> std::io::Result<()> {
         .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?,
     );
 
+    // Start the service discovery if enabled
+    if let Some(service_discovery_config) = config.service_discovery_config {
+        if service_discovery_config.enabled {
+            let worker_urls = Arc::clone(&app_state.router.get_worker_urls());
+
+            match start_service_discovery(service_discovery_config, worker_urls).await {
+                Ok(handle) => {
+                    info!("✅ Service discovery started successfully");
+
+                    // Spawn a task to handle the service discovery thread
+                    spawn(async move {
+                        if let Err(e) = handle.await {
+                            error!("Service discovery task failed: {:?}", e);
+                        }
+                    });
+                }
+                Err(e) => {
+                    error!("Failed to start service discovery: {}", e);
+                    warn!("Continuing without service discovery");
+                }
+            }
+        }
+    }
+
     info!("✅ Serving router on {}:{}", config.host, config.port);
     info!("✅ Serving workers on {:?}", config.worker_urls);