factory.rs 3.92 KB
Newer Older
1
2
3
//! Factory for creating router instances

use super::{pd_router::PDRouter, router::Router, RouterTrait};
4
use crate::config::{PolicyConfig, RoutingMode};
5
use crate::policies::PolicyFactory;
6
7
use crate::server::AppContext;
use std::sync::Arc;
8
9
10
11
12

/// Factory for creating router instances based on configuration
pub struct RouterFactory;

impl RouterFactory {
13
    /// Create a router instance from application context
14
    pub async fn create_router(ctx: &Arc<AppContext>) -> Result<Box<dyn RouterTrait>, String> {
15
16
17
18
19
20
        // Check if IGW mode is enabled
        if ctx.router_config.enable_igw {
            return Self::create_igw_router(ctx).await;
        }

        // Default to proxy mode
21
        match &ctx.router_config.mode {
22
            RoutingMode::Regular { worker_urls } => {
23
                Self::create_regular_router(worker_urls, &ctx.router_config.policy, ctx).await
24
25
26
27
            }
            RoutingMode::PrefillDecode {
                prefill_urls,
                decode_urls,
28
29
                prefill_policy,
                decode_policy,
30
31
32
33
34
35
36
37
38
39
40
            } => {
                Self::create_pd_router(
                    prefill_urls,
                    decode_urls,
                    prefill_policy.as_ref(),
                    decode_policy.as_ref(),
                    &ctx.router_config.policy,
                    ctx,
                )
                .await
            }
41
42
43
44
        }
    }

    /// Create a regular router with injected policy
45
    async fn create_regular_router(
46
47
        worker_urls: &[String],
        policy_config: &PolicyConfig,
48
        ctx: &Arc<AppContext>,
49
50
51
52
    ) -> Result<Box<dyn RouterTrait>, String> {
        // Create policy
        let policy = PolicyFactory::create_from_config(policy_config);

53
        // Create regular router with injected policy and client
54
55
56
        let router = Router::new(
            worker_urls.to_vec(),
            policy,
57
58
59
60
61
            ctx.client.clone(),
            ctx.router_config.worker_startup_timeout_secs,
            ctx.router_config.worker_startup_check_interval_secs,
            ctx.router_config.dp_aware,
            ctx.router_config.api_key.clone(),
62
            ctx.router_config.retry.clone(),
63
            ctx.router_config.circuit_breaker.clone(),
64
            ctx.router_config.health_check.clone(),
65
66
        )
        .await?;
67
68
69
70
71

        Ok(Box::new(router))
    }

    /// Create a PD router with injected policy
72
    async fn create_pd_router(
73
74
        prefill_urls: &[(String, Option<u16>)],
        decode_urls: &[String],
75
76
77
        prefill_policy_config: Option<&PolicyConfig>,
        decode_policy_config: Option<&PolicyConfig>,
        main_policy_config: &PolicyConfig,
78
        ctx: &Arc<AppContext>,
79
    ) -> Result<Box<dyn RouterTrait>, String> {
80
81
82
83
84
        // Create policies - use specific policies if provided, otherwise fall back to main policy
        let prefill_policy =
            PolicyFactory::create_from_config(prefill_policy_config.unwrap_or(main_policy_config));
        let decode_policy =
            PolicyFactory::create_from_config(decode_policy_config.unwrap_or(main_policy_config));
85

86
        // Create PD router with separate policies and client
87
88
89
        let router = PDRouter::new(
            prefill_urls.to_vec(),
            decode_urls.to_vec(),
90
91
            prefill_policy,
            decode_policy,
92
93
94
            ctx.client.clone(),
            ctx.router_config.worker_startup_timeout_secs,
            ctx.router_config.worker_startup_check_interval_secs,
95
            ctx.router_config.retry.clone(),
96
            ctx.router_config.circuit_breaker.clone(),
97
            ctx.router_config.health_check.clone(),
98
99
        )
        .await?;
100
101
102

        Ok(Box::new(router))
    }
103
104
105
106
107
108

    /// Create an IGW router (placeholder for future implementation)
    async fn create_igw_router(_ctx: &Arc<AppContext>) -> Result<Box<dyn RouterTrait>, String> {
        // For now, return an error indicating IGW is not yet implemented
        Err("IGW mode is not yet implemented".to_string())
    }
109
}