[rust] refactor server and router (#1922)

rust/src/server.rs

@@ -1,38 +1,28 @@
-use crate::router::create_router;
 use crate::router::Router;
-use actix_web::http::header::{HeaderValue, CONTENT_TYPE};
 use actix_web::{get, post, web, App, HttpRequest, HttpResponse, HttpServer, Responder};
 use bytes::Bytes;
-use futures_util::StreamExt;
 
 #[derive(Debug)]
 pub struct AppState {
-    router: Box<dyn Router>,
+    router: Router,
     client: reqwest::Client,
 }
 
 impl AppState {
     pub fn new(worker_urls: Vec<String>, policy: String, client: reqwest::Client) -> Self {
         // Create router based on policy
-        let router = create_router(worker_urls, policy);
+        let router = Router::new(worker_urls, policy);
 
         Self { router, client }
     }
 }
 
-#[get("/v1/models")]
-async fn v1_model(data: web::Data<AppState>) -> impl Responder {
-    let worker_url = match data.router.get_first() {
-        Some(url) => url,
-        None => return HttpResponse::InternalServerError().finish(),
-    };
-    // Use the shared client
-    match data
-        .client
-        .get(format!("{}/v1/models", worker_url))
-        .send()
-        .await
-    {
+async fn forward_request(
+    client: &reqwest::Client,
+    worker_url: String,
+    route: String,
+) -> HttpResponse {
+    match client.get(format!("{}{}", worker_url, route)).send().await {
         Ok(res) => {
             let status = actix_web::http::StatusCode::from_u16(res.status().as_u16())
                 .unwrap_or(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR);
@@ -48,85 +38,31 @@ async fn v1_model(data: web::Data<AppState>) -> impl Responder {
     }
 }
 
+#[get("/v1/models")]
+async fn v1_model(data: web::Data<AppState>) -> impl Responder {
+    // TODO: extract forward_to_route
+    let worker_url = match data.router.get_first() {
+        Some(url) => url,
+        None => return HttpResponse::InternalServerError().finish(),
+    };
+
+    forward_request(&data.client, worker_url, "/v1/models".to_string()).await
+}
+
 #[get("/get_model_info")]
 async fn get_model_info(data: web::Data<AppState>) -> impl Responder {
     let worker_url = match data.router.get_first() {
         Some(url) => url,
         None => return HttpResponse::InternalServerError().finish(),
     };
-    // Use the shared client
-    match data
-        .client
-        .get(format!("{}/get_model_info", worker_url))
-        .send()
-        .await
-    {
-        Ok(res) => {
-            let status = actix_web::http::StatusCode::from_u16(res.status().as_u16())
-                .unwrap_or(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR);
 
-            // print the status
-            println!("Worker URL: {}, Status: {}", worker_url, status);
-            match res.bytes().await {
-                Ok(body) => HttpResponse::build(status).body(body.to_vec()),
-                Err(_) => HttpResponse::InternalServerError().finish(),
-            }
-        }
-        Err(_) => HttpResponse::InternalServerError().finish(),
-    }
+    forward_request(&data.client, worker_url, "/get_model_info".to_string()).await
 }
 
 // no deser and ser, just forward and return
 #[post("/generate")]
 async fn generate(req: HttpRequest, body: Bytes, data: web::Data<AppState>) -> impl Responder {
-    // create a router struct
-    // TODO: use router abstraction for different policy
-    let worker_url = match data.router.select() {
-        Some(url) => url,
-        None => return HttpResponse::InternalServerError().finish(),
-    };
-
-    // Check if client requested streaming
-    let is_stream = serde_json::from_slice::<serde_json::Value>(&body)
-        .map(|v| v.get("stream").and_then(|s| s.as_bool()).unwrap_or(false))
-        .unwrap_or(false);
-
-    let res = match data
-        .client
-        .post(format!("{}/generate", worker_url))
-        .header(
-            "Content-Type",
-            req.headers()
-                .get("Content-Type")
-                .and_then(|h| h.to_str().ok())
-                .unwrap_or("application/json"),
-        )
-        .body(body.to_vec())
-        .send()
-        .await
-    {
-        Ok(res) => res,
-        Err(_) => return HttpResponse::InternalServerError().finish(),
-    };
-
-    let status = actix_web::http::StatusCode::from_u16(res.status().as_u16())
-        .unwrap_or(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR);
-
-    if !is_stream {
-        match res.bytes().await {
-            Ok(body) => HttpResponse::build(status).body(body.to_vec()),
-            Err(_) => HttpResponse::InternalServerError().finish(),
-        }
-    } else {
-        HttpResponse::build(status)
-            .insert_header((CONTENT_TYPE, HeaderValue::from_static("text/event-stream")))
-            .streaming(res.bytes_stream().map(|b| match b {
-                Ok(b) => Ok::<_, actix_web::Error>(b),
-                Err(_) => Err(actix_web::error::ErrorInternalServerError(
-                    "Failed to read stream",
-                )),
-            }))
-    }
+    data.router.dispatch(&data.client, req, body).await
 }
 
 pub async fn startup(