server.rs 27.4 KB
Newer Older
1
use crate::{
2
    config::{ConnectionMode, HistoryBackend, RouterConfig},
3
    core::{WorkerRegistry, WorkerType},
4
    data_connector::{MemoryResponseStorage, NoOpResponseStorage, SharedResponseStorage},
5
6
7
8
9
10
11
12
13
14
15
16
    logging::{self, LoggingConfig},
    metrics::{self, PrometheusConfig},
    middleware::{self, QueuedRequest, TokenBucket},
    policies::PolicyRegistry,
    protocols::{
        spec::{
            ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest,
            RerankRequest, ResponsesRequest, V1RerankReqInput,
        },
        worker_spec::{WorkerApiResponse, WorkerConfigRequest, WorkerErrorResponse},
    },
    reasoning_parser::ParserFactory,
17
    routers::WorkerInitializer,
18
19
20
21
22
23
24
    routers::{
        router_manager::{RouterId, RouterManager},
        RouterFactory, RouterTrait,
    },
    service_discovery::{start_service_discovery, ServiceDiscoveryConfig},
    tokenizer::{factory as tokenizer_factory, traits::Tokenizer},
    tool_parser::ParserRegistry,
25
};
26
use axum::{
27
    extract::{Path, Query, Request, State},
28
29
    http::StatusCode,
    response::{IntoResponse, Response},
30
    routing::{delete, get, post},
31
    serve, Json, Router,
32
};
33
use reqwest::Client;
34
35
36
37
38
39
40
41
use serde::Deserialize;
use serde_json::json;
use std::{
    sync::atomic::{AtomicBool, Ordering},
    sync::Arc,
    time::Duration,
};
use tokio::{net::TcpListener, signal, spawn};
42
use tracing::{error, info, warn, Level};
43

44
#[derive(Clone)]
45
pub struct AppContext {
46
    pub client: Client,
47
    pub router_config: RouterConfig,
48
    pub rate_limiter: Arc<TokenBucket>,
49
50
51
    pub tokenizer: Option<Arc<dyn Tokenizer>>,
    pub reasoning_parser_factory: Option<ParserFactory>,
    pub tool_parser_registry: Option<&'static ParserRegistry>,
52
53
54
    pub worker_registry: Arc<WorkerRegistry>,
    pub policy_registry: Arc<PolicyRegistry>,
    pub router_manager: Option<Arc<RouterManager>>,
55
    pub response_storage: SharedResponseStorage,
56
57
}

58
impl AppContext {
59
60
61
62
    pub fn new(
        router_config: RouterConfig,
        client: Client,
        max_concurrent_requests: usize,
63
        rate_limit_tokens_per_second: Option<usize>,
64
    ) -> Result<Self, String> {
65
66
        let rate_limit_tokens = rate_limit_tokens_per_second.unwrap_or(max_concurrent_requests);
        let rate_limiter = Arc::new(TokenBucket::new(max_concurrent_requests, rate_limit_tokens));
67
68
69

        // Initialize gRPC-specific components only when in gRPC mode
        let (tokenizer, reasoning_parser_factory, tool_parser_registry) =
70
            if router_config.connection_mode == ConnectionMode::Grpc {
71
72
73
74
75
76
77
78
79
80
81
82
83
                // Get tokenizer path (required for gRPC mode)
                let tokenizer_path = router_config
                    .tokenizer_path
                    .clone()
                    .or_else(|| router_config.model_path.clone())
                    .ok_or_else(|| {
                        "gRPC mode requires either --tokenizer-path or --model-path to be specified"
                            .to_string()
                    })?;

                // Initialize all gRPC components
                let tokenizer = Some(
                    tokenizer_factory::create_tokenizer(&tokenizer_path)
84
                        .map_err(|e| format!("Failed to create tokenizer: {e}"))?,
85
86
87
88
89
90
91
92
93
94
                );
                let reasoning_parser_factory = Some(ParserFactory::new());
                let tool_parser_registry = Some(ParserRegistry::new());

                (tokenizer, reasoning_parser_factory, tool_parser_registry)
            } else {
                // HTTP mode doesn't need these components
                (None, None, None)
            };

95
        let worker_registry = Arc::new(WorkerRegistry::new());
96
        let policy_registry = Arc::new(PolicyRegistry::new(router_config.policy.clone()));
97

98
        let router_manager = None;
99

100
101
102
103
104
105
        // Initialize response storage based on configuration
        let response_storage: SharedResponseStorage = match router_config.history_backend {
            HistoryBackend::Memory => Arc::new(MemoryResponseStorage::new()),
            HistoryBackend::None => Arc::new(NoOpResponseStorage::new()),
        };

106
        Ok(Self {
107
            client,
108
            router_config,
109
            rate_limiter,
110
111
112
            tokenizer,
            reasoning_parser_factory,
            tool_parser_registry,
113
114
115
            worker_registry,
            policy_registry,
            router_manager,
116
            response_storage,
117
        })
118
119
120
    }
}

121
122
123
124
#[derive(Clone)]
pub struct AppState {
    pub router: Arc<dyn RouterTrait>,
    pub context: Arc<AppContext>,
125
    pub concurrency_queue_tx: Option<tokio::sync::mpsc::Sender<QueuedRequest>>,
126
    pub router_manager: Option<Arc<RouterManager>>,
127
128
}

129
130
131
// Fallback handler for unmatched routes
async fn sink_handler() -> Response {
    StatusCode::NOT_FOUND.into_response()
132
133
}

134
135
136
// Health check endpoints
async fn liveness(State(state): State<Arc<AppState>>) -> Response {
    state.router.liveness()
137
138
}

139
140
async fn readiness(State(state): State<Arc<AppState>>) -> Response {
    state.router.readiness()
141
142
}

143
async fn health(State(state): State<Arc<AppState>>, req: Request) -> Response {
144
    state.router.health(req).await
145
146
}

147
async fn health_generate(State(state): State<Arc<AppState>>, req: Request) -> Response {
148
    state.router.health_generate(req).await
149
150
}

151
async fn get_server_info(State(state): State<Arc<AppState>>, req: Request) -> Response {
152
    state.router.get_server_info(req).await
153
154
}

155
async fn v1_models(State(state): State<Arc<AppState>>, req: Request) -> Response {
156
    state.router.get_models(req).await
157
158
}

159
async fn get_model_info(State(state): State<Arc<AppState>>, req: Request) -> Response {
160
    state.router.get_model_info(req).await
161
}
162

163
164
// Generation endpoints
// The RouterTrait now accepts optional headers and typed body directly
165
async fn generate(
166
167
168
169
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<GenerateRequest>,
) -> Response {
170
171
172
173
    state
        .router
        .route_generate(Some(&headers), &body, None)
        .await
174
175
176
}

async fn v1_chat_completions(
177
178
179
180
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<ChatCompletionRequest>,
) -> Response {
181
    state.router.route_chat(Some(&headers), &body, None).await
182
183
184
}

async fn v1_completions(
185
186
187
188
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<CompletionRequest>,
) -> Response {
189
190
191
192
    state
        .router
        .route_completion(Some(&headers), &body, None)
        .await
193
194
}

195
196
197
198
199
async fn rerank(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<RerankRequest>,
) -> Response {
200
    state.router.route_rerank(Some(&headers), &body, None).await
201
202
203
204
205
206
207
208
209
}

async fn v1_rerank(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<V1RerankReqInput>,
) -> Response {
    state
        .router
210
        .route_rerank(Some(&headers), &body.into(), None)
211
212
213
        .await
}

214
215
216
217
218
async fn v1_responses(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<ResponsesRequest>,
) -> Response {
219
220
221
222
    state
        .router
        .route_responses(Some(&headers), &body, None)
        .await
223
224
}

225
226
227
228
229
230
231
232
233
234
235
async fn v1_embeddings(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<EmbeddingRequest>,
) -> Response {
    state
        .router
        .route_embeddings(Some(&headers), &body, None)
        .await
}

236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
async fn v1_responses_get(
    State(state): State<Arc<AppState>>,
    Path(response_id): Path<String>,
    headers: http::HeaderMap,
) -> Response {
    state
        .router
        .get_response(Some(&headers), &response_id)
        .await
}

async fn v1_responses_cancel(
    State(state): State<Arc<AppState>>,
    Path(response_id): Path<String>,
    headers: http::HeaderMap,
) -> Response {
    state
        .router
        .cancel_response(Some(&headers), &response_id)
        .await
}

async fn v1_responses_delete(
    State(state): State<Arc<AppState>>,
    Path(response_id): Path<String>,
    headers: http::HeaderMap,
) -> Response {
    // Python server does not support this yet
    state
        .router
        .delete_response(Some(&headers), &response_id)
        .await
}

async fn v1_responses_list_input_items(
    State(state): State<Arc<AppState>>,
    Path(response_id): Path<String>,
    headers: http::HeaderMap,
) -> Response {
    // Python server does not support this yet
    state
        .router
        .list_response_input_items(Some(&headers), &response_id)
        .await
}

282
283
284
285
286
287
288
// ---------- Worker management endpoints (Legacy) ----------

#[derive(Deserialize)]
struct UrlQuery {
    url: String,
}

289
async fn add_worker(
290
    State(state): State<Arc<AppState>>,
291
    Query(UrlQuery { url }): Query<UrlQuery>,
292
) -> Response {
293
    match state.router.add_worker(&url).await {
294
295
        Ok(message) => (StatusCode::OK, message).into_response(),
        Err(error) => (StatusCode::BAD_REQUEST, error).into_response(),
296
    }
297
298
}

299
300
301
async fn list_workers(State(state): State<Arc<AppState>>) -> Response {
    let worker_list = state.router.get_worker_urls();
    Json(serde_json::json!({ "urls": worker_list })).into_response()
302
303
}

304
async fn remove_worker(
305
    State(state): State<Arc<AppState>>,
306
    Query(UrlQuery { url }): Query<UrlQuery>,
307
) -> Response {
308
    state.router.remove_worker(&url);
309
310
    (
        StatusCode::OK,
311
        format!("Successfully removed worker: {url}"),
312
313
    )
        .into_response()
314
315
}

316
async fn flush_cache(State(state): State<Arc<AppState>>, _req: Request) -> Response {
317
    state.router.flush_cache().await
318
319
}

320
async fn get_loads(State(state): State<Arc<AppState>>, _req: Request) -> Response {
321
    state.router.get_worker_loads().await
322
323
}

324
// ---------- Worker management endpoints (RESTful) ----------
325
326
327
328
329
330

/// POST /workers - Add a new worker with full configuration
async fn create_worker(
    State(state): State<Arc<AppState>>,
    Json(config): Json<WorkerConfigRequest>,
) -> Response {
331
332
    // Check if we have a RouterManager (enable_igw=true)
    if let Some(router_manager) = &state.router_manager {
333
        // Call RouterManager's add_worker method directly with the full config
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
        match router_manager.add_worker(config).await {
            Ok(response) => (StatusCode::OK, Json(response)).into_response(),
            Err(error) => (StatusCode::BAD_REQUEST, Json(error)).into_response(),
        }
    } else {
        // In single router mode, use the router's add_worker with basic config
        match state.router.add_worker(&config.url).await {
            Ok(message) => {
                let response = WorkerApiResponse {
                    success: true,
                    message,
                    worker: None,
                };
                (StatusCode::OK, Json(response)).into_response()
            }
            Err(error) => {
                let error_response = WorkerErrorResponse {
                    error,
                    code: "ADD_WORKER_FAILED".to_string(),
                };
                (StatusCode::BAD_REQUEST, Json(error_response)).into_response()
            }
        }
    }
}

/// GET /workers - List all workers with details
async fn list_workers_rest(State(state): State<Arc<AppState>>) -> Response {
362
    if let Some(router_manager) = &state.router_manager {
363
364
365
366
367
368
369
370
371
372
        let response = router_manager.list_workers();
        Json(response).into_response()
    } else {
        // In single router mode, get detailed worker info from registry
        let workers = state.context.worker_registry.get_all();
        let response = serde_json::json!({
            "workers": workers.iter().map(|worker| {
                let mut worker_info = serde_json::json!({
                    "url": worker.url(),
                    "model_id": worker.model_id(),
373
374
375
376
377
                    "worker_type": match worker.worker_type() {
                        WorkerType::Regular => "regular",
                        WorkerType::Prefill { .. } => "prefill",
                        WorkerType::Decode => "decode",
                    },
378
379
380
381
382
383
384
385
                    "is_healthy": worker.is_healthy(),
                    "load": worker.load(),
                    "connection_mode": format!("{:?}", worker.connection_mode()),
                    "priority": worker.priority(),
                    "cost": worker.cost(),
                });

                // Add bootstrap_port for Prefill workers
386
                if let WorkerType::Prefill { bootstrap_port } = worker.worker_type() {
387
388
389
390
391
392
393
394
395
                    worker_info["bootstrap_port"] = serde_json::json!(bootstrap_port);
                }

                worker_info
            }).collect::<Vec<_>>(),
            "total": workers.len(),
            "stats": {
                "prefill_count": state.context.worker_registry.get_prefill_workers().len(),
                "decode_count": state.context.worker_registry.get_decode_workers().len(),
396
                "regular_count": state.context.worker_registry.get_by_type(&WorkerType::Regular).len(),
397
398
399
400
401
402
403
            }
        });
        Json(response).into_response()
    }
}

/// GET /workers/{url} - Get specific worker info
404
async fn get_worker(State(state): State<Arc<AppState>>, Path(url): Path<String>) -> Response {
405
    if let Some(router_manager) = &state.router_manager {
406
407
408
409
        if let Some(worker) = router_manager.get_worker(&url) {
            Json(worker).into_response()
        } else {
            let error = WorkerErrorResponse {
410
                error: format!("Worker {url} not found"),
411
412
413
414
415
416
417
                code: "WORKER_NOT_FOUND".to_string(),
            };
            (StatusCode::NOT_FOUND, Json(error)).into_response()
        }
    } else {
        let workers = state.router.get_worker_urls();
        if workers.contains(&url) {
418
            Json(json!({
419
420
421
                "url": url,
                "model_id": "unknown",
                "is_healthy": true
422
423
            }))
            .into_response()
424
425
        } else {
            let error = WorkerErrorResponse {
426
                error: format!("Worker {url} not found"),
427
428
429
430
431
432
433
434
                code: "WORKER_NOT_FOUND".to_string(),
            };
            (StatusCode::NOT_FOUND, Json(error)).into_response()
        }
    }
}

/// DELETE /workers/{url} - Remove a worker
435
async fn delete_worker(State(state): State<Arc<AppState>>, Path(url): Path<String>) -> Response {
436
    if let Some(router_manager) = &state.router_manager {
437
438
439
440
441
442
443
444
445
        match router_manager.remove_worker_from_registry(&url) {
            Ok(response) => (StatusCode::OK, Json(response)).into_response(),
            Err(error) => (StatusCode::BAD_REQUEST, Json(error)).into_response(),
        }
    } else {
        // In single router mode, use router's remove_worker
        state.router.remove_worker(&url);
        let response = WorkerApiResponse {
            success: true,
446
            message: format!("Worker {url} removed successfully"),
447
448
449
450
451
452
            worker: None,
        };
        (StatusCode::OK, Json(response)).into_response()
    }
}

453
454
455
pub struct ServerConfig {
    pub host: String,
    pub port: u16,
456
    pub router_config: RouterConfig,
457
    pub max_payload_size: usize,
458
    pub log_dir: Option<String>,
459
    pub log_level: Option<String>,
460
    pub service_discovery_config: Option<ServiceDiscoveryConfig>,
461
    pub prometheus_config: Option<PrometheusConfig>,
462
    pub request_timeout_secs: u64,
463
    pub request_id_headers: Option<Vec<String>>,
464
465
}

466
467
468
469
470
471
472
473
474
475
476
/// Build the Axum application with all routes and middleware
pub fn build_app(
    app_state: Arc<AppState>,
    max_payload_size: usize,
    request_id_headers: Vec<String>,
    cors_allowed_origins: Vec<String>,
) -> Router {
    // Create routes
    let protected_routes = Router::new()
        .route("/generate", post(generate))
        .route("/v1/chat/completions", post(v1_chat_completions))
477
        .route("/v1/completions", post(v1_completions))
478
479
        .route("/rerank", post(rerank))
        .route("/v1/rerank", post(v1_rerank))
480
        .route("/v1/responses", post(v1_responses))
481
        .route("/v1/embeddings", post(v1_embeddings))
482
483
484
485
486
487
488
489
490
491
        .route("/v1/responses/{response_id}", get(v1_responses_get))
        .route(
            "/v1/responses/{response_id}/cancel",
            post(v1_responses_cancel),
        )
        .route("/v1/responses/{response_id}", delete(v1_responses_delete))
        .route(
            "/v1/responses/{response_id}/input",
            get(v1_responses_list_input_items),
        )
492
493
        .route_layer(axum::middleware::from_fn_with_state(
            app_state.clone(),
494
            middleware::concurrency_limit_middleware,
495
        ));
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512

    let public_routes = Router::new()
        .route("/liveness", get(liveness))
        .route("/readiness", get(readiness))
        .route("/health", get(health))
        .route("/health_generate", get(health_generate))
        .route("/v1/models", get(v1_models))
        .route("/get_model_info", get(get_model_info))
        .route("/get_server_info", get(get_server_info));

    let admin_routes = Router::new()
        .route("/add_worker", post(add_worker))
        .route("/remove_worker", post(remove_worker))
        .route("/list_workers", get(list_workers))
        .route("/flush_cache", post(flush_cache))
        .route("/get_loads", get(get_loads));

513
514
515
516
517
    // Worker management routes
    let worker_routes = Router::new()
        .route("/workers", post(create_worker))
        .route("/workers", get(list_workers_rest))
        .route("/workers/{url}", get(get_worker))
518
        .route("/workers/{url}", delete(delete_worker));
519

520
521
522
523
524
    // Build app with all routes and middleware
    Router::new()
        .merge(protected_routes)
        .merge(public_routes)
        .merge(admin_routes)
525
        .merge(worker_routes)
526
527
528
529
        // Request body size limiting
        .layer(tower_http::limit::RequestBodyLimitLayer::new(
            max_payload_size,
        ))
530
531
        .layer(middleware::create_logging_layer())
        .layer(middleware::RequestIdLayer::new(request_id_headers))
532
533
534
535
536
537
        .layer(create_cors_layer(cors_allowed_origins))
        .fallback(sink_handler)
        .with_state(app_state)
}

pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Error>> {
538
539
540
541
542
    // Only initialize logging if not already done (for Python bindings support)
    static LOGGING_INITIALIZED: AtomicBool = AtomicBool::new(false);

    let _log_guard = if !LOGGING_INITIALIZED.swap(true, Ordering::SeqCst) {
        Some(logging::init_logging(LoggingConfig {
543
544
545
546
547
548
            level: config
                .log_level
                .as_deref()
                .and_then(|s| match s.to_uppercase().parse::<Level>() {
                    Ok(l) => Some(l),
                    Err(_) => {
549
                        warn!("Invalid log level string: '{s}'. Defaulting to INFO.");
550
551
552
553
                        None
                    }
                })
                .unwrap_or(Level::INFO),
554
555
556
557
558
559
560
561
562
            json_format: false,
            log_dir: config.log_dir.clone(),
            colorize: true,
            log_file_name: "sgl-router".to_string(),
            log_targets: None,
        }))
    } else {
        None
    };
563

564
565
    // Initialize prometheus metrics exporter
    if let Some(prometheus_config) = config.prometheus_config {
566
        metrics::start_prometheus(prometheus_config);
567
568
    }

569
    info!(
570
571
572
573
574
        "Starting router on {}:{} | mode: {:?} | policy: {:?} | max_payload: {}MB",
        config.host,
        config.port,
        config.router_config.mode,
        config.router_config.policy,
575
576
577
        config.max_payload_size / (1024 * 1024)
    );

578
    let client = Client::builder()
579
        .pool_idle_timeout(Some(Duration::from_secs(50)))
580
        .pool_max_idle_per_host(500)
581
        .timeout(Duration::from_secs(config.request_timeout_secs))
582
        .connect_timeout(Duration::from_secs(10))
583
        .tcp_nodelay(true)
584
        .tcp_keepalive(Some(Duration::from_secs(30)))
585
586
587
        .build()
        .expect("Failed to create HTTP client");

588
    // Create the application context with all dependencies
589
    let app_context = AppContext::new(
590
591
592
        config.router_config.clone(),
        client.clone(),
        config.router_config.max_concurrent_requests,
593
        config.router_config.rate_limit_tokens_per_second,
594
595
596
597
    )?;

    let app_context = Arc::new(app_context);

598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
    // Initialize workers before creating routers
    // This separates worker lifecycle from router lifecycle
    info!(
        "Initializing workers for routing mode: {:?}",
        config.router_config.mode
    );
    WorkerInitializer::initialize_workers(&config.router_config, &app_context.worker_registry)
        .await
        .map_err(|e| format!("Failed to initialize workers: {}", e))?;

    let worker_stats = app_context.worker_registry.stats();
    info!(
        "Workers initialized: {} total, {} healthy",
        worker_stats.total_workers, worker_stats.healthy_workers
    );

614
    // Create the appropriate router based on enable_igw flag
615
616
617
618
619
620
621
622
623
624
625
626
627
    let (router, router_manager): (Arc<dyn RouterTrait>, Option<Arc<RouterManager>>) =
        if config.router_config.enable_igw {
            info!("Multi-router mode enabled (enable_igw=true)");

            // Create RouterManager with shared registries from AppContext
            let router_manager = Arc::new(RouterManager::new(
                config.router_config.clone(),
                client.clone(),
                app_context.worker_registry.clone(),
                app_context.policy_registry.clone(),
            ));

            // 1. HTTP Regular Router
628
            match RouterFactory::create_regular_router(&app_context).await {
629
630
631
632
633
634
635
636
637
638
                Ok(http_regular) => {
                    info!("Created HTTP Regular router");
                    router_manager.register_router(
                        RouterId::new("http-regular".to_string()),
                        Arc::from(http_regular),
                    );
                }
                Err(e) => {
                    warn!("Failed to create HTTP Regular router: {e}");
                }
639
640
            }

641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
            // 2. HTTP PD Router
            match RouterFactory::create_pd_router(
                None,
                None,
                &config.router_config.policy,
                &app_context,
            )
            .await
            {
                Ok(http_pd) => {
                    info!("Created HTTP PD router");
                    router_manager
                        .register_router(RouterId::new("http-pd".to_string()), Arc::from(http_pd));
                }
                Err(e) => {
                    warn!("Failed to create HTTP PD router: {e}");
                }
658
            }
659

660
            // TODO: Add gRPC routers once we have dynamic tokenizer loading
661

662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
            info!(
                "RouterManager initialized with {} routers",
                router_manager.router_count()
            );
            (
                router_manager.clone() as Arc<dyn RouterTrait>,
                Some(router_manager),
            )
        } else {
            info!("Single router mode (enable_igw=false)");
            // Create single router with the context
            (
                Arc::from(RouterFactory::create_router(&app_context).await?),
                None,
            )
        };
678
679
680
681
682
683
684
685
686

    // Start health checker for all workers in the registry
    let _health_checker = app_context
        .worker_registry
        .start_health_checker(config.router_config.health_check.check_interval_secs);
    info!(
        "Started health checker for workers with {}s interval",
        config.router_config.health_check.check_interval_secs
    );
687

688
    // Set up concurrency limiter with queue if configured
689
    let (limiter, processor) = middleware::ConcurrencyLimiter::new(
690
691
692
693
694
695
696
        app_context.rate_limiter.clone(),
        config.router_config.queue_size,
        Duration::from_secs(config.router_config.queue_timeout_secs),
    );

    // Start queue processor if enabled
    if let Some(processor) = processor {
697
        spawn(processor.run());
698
699
700
701
702
703
        info!(
            "Started request queue with size: {}, timeout: {}s",
            config.router_config.queue_size, config.router_config.queue_timeout_secs
        );
    }

704
705
    // Create app state with router and context
    let app_state = Arc::new(AppState {
706
        router,
707
        context: app_context.clone(),
708
        concurrency_queue_tx: limiter.queue_tx.clone(),
709
        router_manager,
710
    });
711
    let router_arc = Arc::clone(&app_state.router);
712

713
714
715
    // Start the service discovery if enabled
    if let Some(service_discovery_config) = config.service_discovery_config {
        if service_discovery_config.enabled {
716
            match start_service_discovery(service_discovery_config, router_arc).await {
717
                Ok(handle) => {
718
                    info!("Service discovery started");
719
720
721
722
723
724
725
726
                    // 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) => {
727
                    error!("Failed to start service discovery: {e}");
728
729
730
731
732
733
                    warn!("Continuing without service discovery");
                }
            }
        }
    }

734
    info!(
735
        "Router ready | workers: {:?}",
736
        app_state.router.get_worker_urls()
737
    );
738

739
740
741
742
743
744
745
746
747
    let request_id_headers = config.request_id_headers.clone().unwrap_or_else(|| {
        vec![
            "x-request-id".to_string(),
            "x-correlation-id".to_string(),
            "x-trace-id".to_string(),
            "request-id".to_string(),
        ]
    });

748
749
750
751
752
753
754
    // Build the application
    let app = build_app(
        app_state,
        config.max_payload_size,
        request_id_headers,
        config.router_config.cors_allowed_origins.clone(),
    );
755

756
757
758
    let addr = format!("{}:{}", config.host, config.port);
    let listener = TcpListener::bind(&addr).await?;
    info!("Starting server on {}", addr);
759
    serve(listener, app)
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
        .with_graceful_shutdown(shutdown_signal())
        .await
        .map_err(|e| Box::new(e) as Box<dyn std::error::Error>)?;

    Ok(())
}

// Graceful shutdown handler
async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {
            info!("Received Ctrl+C, starting graceful shutdown");
        },
        _ = terminate => {
            info!("Received terminate signal, starting graceful shutdown");
        },
    }
}

// CORS Layer Creation
fn create_cors_layer(allowed_origins: Vec<String>) -> tower_http::cors::CorsLayer {
    use tower_http::cors::Any;

    let cors = if allowed_origins.is_empty() {
        // Allow all origins if none specified
        tower_http::cors::CorsLayer::new()
            .allow_origin(Any)
            .allow_methods(Any)
            .allow_headers(Any)
            .expose_headers(Any)
    } else {
        // Restrict to specific origins
        let origins: Vec<http::HeaderValue> = allowed_origins
            .into_iter()
            .filter_map(|origin| origin.parse().ok())
            .collect();

        tower_http::cors::CorsLayer::new()
            .allow_origin(origins)
            .allow_methods([http::Method::GET, http::Method::POST, http::Method::OPTIONS])
            .allow_headers([http::header::CONTENT_TYPE, http::header::AUTHORIZATION])
            .expose_headers([http::header::HeaderName::from_static("x-request-id")])
    };

    cors.max_age(Duration::from_secs(3600))
822
}