"examples/production_monitoring/Otel.md" did not exist on "d77d1901ad0c33f0dbfc74e118442a3cade0cbc2"
server.rs 27.8 KB
Newer Older
1
use crate::config::RouterConfig;
2
use crate::core::WorkerRegistry;
3
use crate::logging::{self, LoggingConfig};
4
use crate::metrics::{self, PrometheusConfig};
5
use crate::middleware::TokenBucket;
6
use crate::policies::PolicyRegistry;
7
use crate::protocols::spec::{
8
9
    ChatCompletionRequest, CompletionRequest, GenerateRequest, RerankRequest, ResponsesRequest,
    V1RerankReqInput,
10
};
11
use crate::protocols::worker_spec::{WorkerApiResponse, WorkerConfigRequest, WorkerErrorResponse};
12
use crate::reasoning_parser::ParserFactory;
13
use crate::routers::router_manager::{RouterId, RouterManager};
14
use crate::routers::{RouterFactory, RouterTrait};
15
use crate::service_discovery::{start_service_discovery, ServiceDiscoveryConfig};
16
17
use crate::tokenizer::{factory as tokenizer_factory, traits::Tokenizer};
use crate::tool_parser::ParserRegistry;
18
use axum::{
19
    extract::{Path, Query, Request, State},
20
21
    http::StatusCode,
    response::{IntoResponse, Response},
22
    routing::{delete, get, post},
23
    Json, Router,
24
};
25
use reqwest::Client;
26
use std::collections::HashMap;
27
use std::sync::atomic::{AtomicBool, Ordering};
28
use std::sync::Arc;
29
use std::time::Duration;
30
31
use tokio::net::TcpListener;
use tokio::signal;
32
33
use tokio::spawn;
use tracing::{error, info, warn, Level};
34

35
#[derive(Clone)]
36
pub struct AppContext {
37
    pub client: Client,
38
    pub router_config: RouterConfig,
39
    pub rate_limiter: Arc<TokenBucket>,
40
41
42
    pub tokenizer: Option<Arc<dyn Tokenizer>>,
    pub reasoning_parser_factory: Option<ParserFactory>,
    pub tool_parser_registry: Option<&'static ParserRegistry>,
43
44
45
    pub worker_registry: Arc<WorkerRegistry>, // Shared worker registry
    pub policy_registry: Arc<PolicyRegistry>, // Shared policy registry
    pub router_manager: Option<Arc<RouterManager>>, // Only present when enable_igw=true
46
47
}

48
impl AppContext {
49
50
51
52
    pub fn new(
        router_config: RouterConfig,
        client: Client,
        max_concurrent_requests: usize,
53
        rate_limit_tokens_per_second: Option<usize>,
54
    ) -> Result<Self, String> {
55
56
        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));
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

        // Initialize gRPC-specific components only when in gRPC mode
        let (tokenizer, reasoning_parser_factory, tool_parser_registry) =
            if router_config.connection_mode == crate::config::ConnectionMode::Grpc {
                // 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)
                        .map_err(|e| format!("Failed to create tokenizer: {}", e))?,
                );
                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)
            };

85
86
87
88
89
90
91
92
93
        // Initialize shared registries
        let worker_registry = Arc::new(WorkerRegistry::new());
        let policy_registry = Arc::new(PolicyRegistry::new(
            router_config.policy.clone(), // Use default policy from config
        ));

        // Initialize RouterManager only when enable_igw is true
        let router_manager = None; // Will be initialized in startup() based on config

94
        Ok(Self {
95
            client,
96
            router_config,
97
            rate_limiter,
98
99
100
            tokenizer,
            reasoning_parser_factory,
            tool_parser_registry,
101
102
103
            worker_registry,
            policy_registry,
            router_manager,
104
        })
105
106
107
    }
}

108
109
110
111
#[derive(Clone)]
pub struct AppState {
    pub router: Arc<dyn RouterTrait>,
    pub context: Arc<AppContext>,
112
    pub concurrency_queue_tx: Option<tokio::sync::mpsc::Sender<crate::middleware::QueuedRequest>>,
113
114
}

115
116
117
// Fallback handler for unmatched routes
async fn sink_handler() -> Response {
    StatusCode::NOT_FOUND.into_response()
118
119
}

120
121
122
// Health check endpoints
async fn liveness(State(state): State<Arc<AppState>>) -> Response {
    state.router.liveness()
123
124
}

125
126
async fn readiness(State(state): State<Arc<AppState>>) -> Response {
    state.router.readiness()
127
128
}

129
async fn health(State(state): State<Arc<AppState>>, req: Request) -> Response {
130
    state.router.health(req).await
131
132
}

133
async fn health_generate(State(state): State<Arc<AppState>>, req: Request) -> Response {
134
    state.router.health_generate(req).await
135
136
}

137
async fn get_server_info(State(state): State<Arc<AppState>>, req: Request) -> Response {
138
    state.router.get_server_info(req).await
139
140
}

141
async fn v1_models(State(state): State<Arc<AppState>>, req: Request) -> Response {
142
    state.router.get_models(req).await
143
144
}

145
async fn get_model_info(State(state): State<Arc<AppState>>, req: Request) -> Response {
146
    state.router.get_model_info(req).await
147
}
148

149
150
// Generation endpoints
// The RouterTrait now accepts optional headers and typed body directly
151
async fn generate(
152
153
154
155
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<GenerateRequest>,
) -> Response {
156
157
158
159
    state
        .router
        .route_generate(Some(&headers), &body, None)
        .await
160
161
162
}

async fn v1_chat_completions(
163
164
165
166
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<ChatCompletionRequest>,
) -> Response {
167
    state.router.route_chat(Some(&headers), &body, None).await
168
169
170
}

async fn v1_completions(
171
172
173
174
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<CompletionRequest>,
) -> Response {
175
176
177
178
    state
        .router
        .route_completion(Some(&headers), &body, None)
        .await
179
180
}

181
182
183
184
185
async fn rerank(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<RerankRequest>,
) -> Response {
186
    state.router.route_rerank(Some(&headers), &body, None).await
187
188
189
190
191
192
193
194
195
}

async fn v1_rerank(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<V1RerankReqInput>,
) -> Response {
    state
        .router
196
        .route_rerank(Some(&headers), &body.into(), None)
197
198
199
        .await
}

200
201
202
203
204
async fn v1_responses(
    State(state): State<Arc<AppState>>,
    headers: http::HeaderMap,
    Json(body): Json<ResponsesRequest>,
) -> Response {
205
206
207
208
    state
        .router
        .route_responses(Some(&headers), &body, None)
        .await
209
210
}

211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
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
}

257
// Worker management endpoints
258
async fn add_worker(
259
260
261
262
    State(state): State<Arc<AppState>>,
    Query(params): Query<HashMap<String, String>>,
) -> Response {
    let worker_url = match params.get("url") {
263
264
        Some(url) => url.to_string(),
        None => {
265
266
267
268
269
            return (
                StatusCode::BAD_REQUEST,
                "Worker URL required. Provide 'url' query parameter",
            )
                .into_response();
270
271
        }
    };
272

273
274
275
    match state.router.add_worker(&worker_url).await {
        Ok(message) => (StatusCode::OK, message).into_response(),
        Err(error) => (StatusCode::BAD_REQUEST, error).into_response(),
276
    }
277
278
}

279
280
281
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()
282
283
}

284
async fn remove_worker(
285
286
287
288
    State(state): State<Arc<AppState>>,
    Query(params): Query<HashMap<String, String>>,
) -> Response {
    let worker_url = match params.get("url") {
289
        Some(url) => url.to_string(),
290
        None => return StatusCode::BAD_REQUEST.into_response(),
291
    };
292

293
294
295
296
297
298
    state.router.remove_worker(&worker_url);
    (
        StatusCode::OK,
        format!("Successfully removed worker: {}", worker_url),
    )
        .into_response()
299
300
}

301
async fn flush_cache(State(state): State<Arc<AppState>>, _req: Request) -> Response {
302
    state.router.flush_cache().await
303
304
}

305
async fn get_loads(State(state): State<Arc<AppState>>, _req: Request) -> Response {
306
    state.router.get_worker_loads().await
307
308
}

309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
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
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
// New RESTful worker management endpoints (when enable_igw=true)

/// POST /workers - Add a new worker with full configuration
async fn create_worker(
    State(state): State<Arc<AppState>>,
    Json(config): Json<WorkerConfigRequest>,
) -> Response {
    // Check if RouterManager is available (enable_igw=true)
    if let Some(router_manager) = &state.context.router_manager {
        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 {
    if let Some(router_manager) = &state.context.router_manager {
        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(),
                    "worker_type": format!("{:?}", worker.worker_type()),
                    "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
                if let crate::core::WorkerType::Prefill { bootstrap_port } = worker.worker_type() {
                    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(),
                "regular_count": state.context.worker_registry.get_by_type(&crate::core::WorkerType::Regular).len(),
            }
        });
        Json(response).into_response()
    }
}

/// GET /workers/{url} - Get specific worker info
async fn get_worker(
    State(state): State<Arc<AppState>>,
    axum::extract::Path(url): axum::extract::Path<String>,
) -> Response {
    if let Some(router_manager) = &state.context.router_manager {
        if let Some(worker) = router_manager.get_worker(&url) {
            Json(worker).into_response()
        } else {
            let error = WorkerErrorResponse {
                error: format!("Worker {} not found", url),
                code: "WORKER_NOT_FOUND".to_string(),
            };
            (StatusCode::NOT_FOUND, Json(error)).into_response()
        }
    } else {
        // In single router mode, check if worker exists
        let workers = state.router.get_worker_urls();
        if workers.contains(&url) {
            let worker_info = serde_json::json!({
                "url": url,
                "model_id": "unknown",
                "is_healthy": true
            });
            Json(worker_info).into_response()
        } else {
            let error = WorkerErrorResponse {
                error: format!("Worker {} not found", url),
                code: "WORKER_NOT_FOUND".to_string(),
            };
            (StatusCode::NOT_FOUND, Json(error)).into_response()
        }
    }
}

/// DELETE /workers/{url} - Remove a worker
async fn delete_worker(
    State(state): State<Arc<AppState>>,
    axum::extract::Path(url): axum::extract::Path<String>,
) -> Response {
    if let Some(router_manager) = &state.context.router_manager {
        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,
            message: format!("Worker {} removed successfully", url),
            worker: None,
        };
        (StatusCode::OK, Json(response)).into_response()
    }
}

440
441
442
pub struct ServerConfig {
    pub host: String,
    pub port: u16,
443
    pub router_config: RouterConfig,
444
    pub max_payload_size: usize,
445
    pub log_dir: Option<String>,
446
    pub log_level: Option<String>,
447
    pub service_discovery_config: Option<ServiceDiscoveryConfig>,
448
    pub prometheus_config: Option<PrometheusConfig>,
449
    pub request_timeout_secs: u64,
450
    pub request_id_headers: Option<Vec<String>>,
451
452
}

453
454
455
456
457
458
459
460
461
462
463
/// 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))
464
        .route("/v1/completions", post(v1_completions))
465
466
        .route("/rerank", post(rerank))
        .route("/v1/rerank", post(v1_rerank))
467
        .route("/v1/responses", post(v1_responses))
468
469
470
471
472
473
474
475
476
477
        .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),
        )
478
479
480
481
        .route_layer(axum::middleware::from_fn_with_state(
            app_state.clone(),
            crate::middleware::concurrency_limit_middleware,
        ));
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498

    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));

499
500
501
502
503
504
505
    // 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))
        .route("/workers/{url}", axum::routing::delete(delete_worker));

506
507
508
509
510
    // Build app with all routes and middleware
    Router::new()
        .merge(protected_routes)
        .merge(public_routes)
        .merge(admin_routes)
511
        .merge(worker_routes)
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
        // Request body size limiting
        .layer(tower_http::limit::RequestBodyLimitLayer::new(
            max_payload_size,
        ))
        // Request ID layer - must be added AFTER logging layer in the code
        // so it executes BEFORE logging layer at runtime (layers execute bottom-up)
        .layer(crate::middleware::RequestIdLayer::new(request_id_headers))
        // Custom logging layer that can now see request IDs from extensions
        .layer(crate::middleware::create_logging_layer())
        // CORS (should be outermost)
        .layer(create_cors_layer(cors_allowed_origins))
        // Fallback
        .fallback(sink_handler)
        // State - apply last to get Router<Arc<AppState>>
        .with_state(app_state)
}

pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Error>> {
530
531
532
533
534
    // 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 {
535
536
537
538
539
540
541
542
543
544
545
            level: config
                .log_level
                .as_deref()
                .and_then(|s| match s.to_uppercase().parse::<Level>() {
                    Ok(l) => Some(l),
                    Err(_) => {
                        warn!("Invalid log level string: '{}'. Defaulting to INFO.", s);
                        None
                    }
                })
                .unwrap_or(Level::INFO),
546
547
548
549
550
551
552
553
554
            json_format: false,
            log_dir: config.log_dir.clone(),
            colorize: true,
            log_file_name: "sgl-router".to_string(),
            log_targets: None,
        }))
    } else {
        None
    };
555

556
557
    // Initialize prometheus metrics exporter
    if let Some(prometheus_config) = config.prometheus_config {
558
        metrics::start_prometheus(prometheus_config);
559
560
    }

561
    info!(
562
563
564
565
566
        "Starting router on {}:{} | mode: {:?} | policy: {:?} | max_payload: {}MB",
        config.host,
        config.port,
        config.router_config.mode,
        config.router_config.policy,
567
568
569
        config.max_payload_size / (1024 * 1024)
    );

570
    let client = Client::builder()
571
        .pool_idle_timeout(Some(Duration::from_secs(50)))
572
        .pool_max_idle_per_host(500) // Increase to 500 connections per host
573
574
575
576
        .timeout(Duration::from_secs(config.request_timeout_secs))
        .connect_timeout(Duration::from_secs(10)) // Separate connection timeout
        .tcp_nodelay(true)
        .tcp_keepalive(Some(Duration::from_secs(30))) // Keep connections alive
577
578
579
        .build()
        .expect("Failed to create HTTP client");

580
    // Create the application context with all dependencies
581
    let app_context = AppContext::new(
582
583
584
        config.router_config.clone(),
        client.clone(),
        config.router_config.max_concurrent_requests,
585
        config.router_config.rate_limit_tokens_per_second,
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
    )?;

    let app_context = Arc::new(app_context);

    // Create the appropriate router based on enable_igw flag
    let router: Box<dyn RouterTrait> = if config.router_config.enable_igw {
        info!("Multi-router mode enabled (enable_igw=true)");

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

        // Create HTTP routers at startup (with empty worker lists)
        // Workers will be added to these routers dynamically via RouterManager's worker registry

        // 1. HTTP Regular Router
        match RouterFactory::create_regular_router(
            &[], // Empty worker list - workers added later
            &app_context,
        )
        .await
        {
            Ok(http_regular) => {
                info!("Created HTTP Regular router");
                router_manager.register_router(
                    RouterId::new("http-regular".to_string()),
                    Arc::from(http_regular),
                    vec![], // Models will be determined by workers
                );
            }
            Err(e) => {
                warn!("Failed to create HTTP Regular router: {}", e);
            }
        }

        // 2. HTTP PD Router
        match RouterFactory::create_pd_router(
            &[],  // Empty prefill URLs
            &[],  // Empty decode URLs
            None, // Use default prefill policy
            None, // Use default decode policy
            &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),
                    vec![],
                );
            }
            Err(e) => {
                warn!("Failed to create HTTP PD router: {}", e);
            }
        }
648

649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
        // TODO: Add gRPC routers once we have dynamic tokenizer loading
        // Currently gRPC routers require tokenizer to be initialized first,
        // but each model needs its own tokenizer. Once we implement dynamic
        // tokenizer loading per model, we can enable gRPC routers here:
        // - RouterType::GrpcRegular (RouterId: "grpc-regular")
        // - RouterType::GrpcPd (RouterId: "grpc-pd")

        info!(
            "RouterManager initialized with {} routers",
            router_manager.router_count()
        );
        Box::new(router_manager)
    } else {
        info!("Single router mode (enable_igw=false)");
        // Create single router with the context
        RouterFactory::create_router(&app_context).await?
    };

    // 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
    );
675

676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
    // Set up concurrency limiter with queue if configured
    let (limiter, processor) = crate::middleware::ConcurrencyLimiter::new(
        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 {
        tokio::spawn(processor.run());
        info!(
            "Started request queue with size: {}, timeout: {}s",
            config.router_config.queue_size, config.router_config.queue_timeout_secs
        );
    }

692
693
694
695
    // Create app state with router and context
    let app_state = Arc::new(AppState {
        router: Arc::from(router),
        context: app_context.clone(),
696
        concurrency_queue_tx: limiter.queue_tx.clone(),
697
    });
698
    let router_arc = Arc::clone(&app_state.router);
699

700
701
702
    // Start the service discovery if enabled
    if let Some(service_discovery_config) = config.service_discovery_config {
        if service_discovery_config.enabled {
703
            match start_service_discovery(service_discovery_config, router_arc).await {
704
                Ok(handle) => {
705
                    info!("Service discovery started");
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
                    // 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");
                }
            }
        }
    }

721
    info!(
722
        "Router ready | workers: {:?}",
723
        app_state.router.get_worker_urls()
724
    );
725

726
727
728
729
730
731
732
733
734
735
    // Configure request ID headers
    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(),
        ]
    });

736
737
738
739
740
741
742
    // Build the application
    let app = build_app(
        app_state,
        config.max_payload_size,
        request_id_headers,
        config.router_config.cors_allowed_origins.clone(),
    );
743

744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
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
    // Create TCP listener - use the configured host
    let addr = format!("{}:{}", config.host, config.port);
    let listener = TcpListener::bind(&addr).await?;

    // Start server with graceful shutdown
    info!("Starting server on {}", addr);

    // Serve the application with graceful shutdown
    axum::serve(listener, app)
        .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))
815
}