[rust] refactor server and router (#1922)

.github/workflows/pr-test-rust.yml

@@ -4,11 +4,11 @@ on:
   push:
     branches: [ main ]
     paths:
-      - "rust/*"
+      - "rust/**"
   pull_request:
     branches: [ main ]
     paths:
-      - "rust/*"
+      - "rust/**"
   workflow_dispatch:
 
 concurrency:

rust/src/router.rs

-// src/router.rs
-
+use actix_web::http::header::{HeaderValue, CONTENT_TYPE};
+use actix_web::{HttpRequest, HttpResponse};
+use bytes::Bytes;
+use futures_util::TryStreamExt;
 use std::fmt::Debug;
 
-/// Generic Router trait that can be implemented with different policies
-pub trait Router: Send + Sync + Debug {
-    /// Select a worker URL based on the implementation's policy
-    /// Returns None if no worker is available
-    fn select(&self) -> Option<String>;
-
-    // get first worker
-    fn get_first(&self) -> Option<String>;
-}
-
-// Round Robin Router
 #[derive(Debug)]
-pub struct RoundRobinRouter {
+pub enum Router {
+    RoundRobin {
         worker_urls: Vec<String>,
-    current_index: std::sync::atomic::AtomicUsize, // AtomicUsize is a thread-safe integer
+        current_index: std::sync::atomic::AtomicUsize,
+    },
+    Random {
+        worker_urls: Vec<String>,
+    },
 }
 
-impl RoundRobinRouter {
-    pub fn new(worker_urls: Vec<String>) -> Self {
-        Self {
+impl Router {
+    pub fn new(worker_urls: Vec<String>, policy: String) -> Self {
+        match policy.to_lowercase().as_str() {
+            "random" => Router::Random { worker_urls },
+            "round_robin" => Router::RoundRobin {
                 worker_urls,
                 current_index: std::sync::atomic::AtomicUsize::new(0),
+            },
+            _ => panic!(
+                "Unknown routing policy: {}. The available policies are 'random' and 'round_robin'",
+                policy
+            ),
         }
     }
-}
 
-impl Router for RoundRobinRouter {
-    fn select(&self) -> Option<String> {
-        if self.worker_urls.is_empty() {
-            return None;
+    pub fn get_first(&self) -> Option<String> {
+        match self {
+            Router::RoundRobin { worker_urls, .. } | Router::Random { worker_urls } => {
+                if worker_urls.is_empty() {
+                    None
+                } else {
+                    Some(worker_urls[0].clone())
                 }
-        // Use relaxed because operation order doesn't matter in round robin
-        let index = self
-            .current_index
-            .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
-            % self.worker_urls.len();
-        Some(self.worker_urls[index].clone())
             }
-
-    fn get_first(&self) -> Option<String> {
-        if self.worker_urls.is_empty() {
-            return None;
         }
-        Some(self.worker_urls[0].clone())
     }
-}
 
-// Random Router
-#[derive(Debug)]
-pub struct RandomRouter {
-    worker_urls: Vec<String>,
-}
+    pub async fn dispatch(
+        &self,
+        client: &reqwest::Client,
+        req: HttpRequest,
+        body: Bytes,
+    ) -> HttpResponse {
+        let worker_url = match self {
+            Router::RoundRobin {
+                worker_urls,
+                current_index,
+            } => {
+                current_index
+                    .fetch_update(
+                        std::sync::atomic::Ordering::SeqCst,
+                        std::sync::atomic::Ordering::SeqCst,
+                        |x| Some((x + 1) % worker_urls.len()),
+                    )
+                    .expect_err("Error updating index in round robin");
 
-impl RandomRouter {
-    pub fn new(worker_urls: Vec<String>) -> Self {
-        Self { worker_urls }
+                &worker_urls[current_index.load(std::sync::atomic::Ordering::SeqCst)]
             }
-}
+            Router::Random { worker_urls } => {
+                &worker_urls[rand::random::<usize>() % worker_urls.len()]
+            }
+        };
 
-impl Router for RandomRouter {
-    fn select(&self) -> Option<String> {
-        use rand::seq::SliceRandom;
+        // 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);
 
-        if self.worker_urls.is_empty() {
-            return None;
-        }
+        let res = match 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(),
+        };
 
-        self.worker_urls.choose(&mut rand::thread_rng()).cloned()
-    }
+        let status = actix_web::http::StatusCode::from_u16(res.status().as_u16())
+            .unwrap_or(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR);
 
-    fn get_first(&self) -> Option<String> {
-        if self.worker_urls.is_empty() {
-            return None;
+        if !is_stream {
+            match res.bytes().await {
+                Ok(body) => HttpResponse::build(status).body(body.to_vec()),
+                Err(_) => HttpResponse::InternalServerError().finish(),
             }
-        Some(self.worker_urls[0].clone())
+        } else {
+            HttpResponse::build(status)
+                .insert_header((CONTENT_TYPE, HeaderValue::from_static("text/event-stream")))
+                .streaming(res.bytes_stream().map_err(|_| {
+                    actix_web::error::ErrorInternalServerError("Failed to read string")
+                }))
         }
-}
-
-// create a router based on routing policy
-pub fn create_router(worker_urls: Vec<String>, policy: String) -> Box<dyn Router> {
-    match policy.to_lowercase().as_str() {
-        "random" => Box::new(RandomRouter::new(worker_urls)),
-        "round_robin" => Box::new(RoundRobinRouter::new(worker_urls)),
-        _ => panic!(
-            "Unknown routing policy: {}. The available policies are 'random' and 'round_robin'",
-            policy
-        ),
     }
 }

rust/src/server.rs

-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("/get_model_info")]
-async fn get_model_info(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(),
     };
-    // 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, "/v1/models".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() {
+#[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(),
     };
 
-    // 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);
+    forward_request(&data.client, worker_url, "/get_model_info".to_string()).await
+}
 
-    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",
-                )),
-            }))
-    }
+// no deser and ser, just forward and return
+#[post("/generate")]
+async fn generate(req: HttpRequest, body: Bytes, data: web::Data<AppState>) -> impl Responder {
+    data.router.dispatch(&data.client, req, body).await
 }
 
 pub async fn startup(