[router] Add remove worker api (#2380)

rust/src/router.rs

@@ -396,4 +396,23 @@ impl Router {
             }
         }
     }
+
+    pub fn remove_worker(&self, worker_url: String) {
+        match self {
+            Router::RoundRobin { worker_urls, .. }
+            | Router::Random { worker_urls }
+            | Router::CacheAware { worker_urls, .. } => {
+                let mut urls = worker_urls.write().unwrap();
+                let index = urls.iter().position(|url| url == &worker_url).unwrap();
+                urls.remove(index);
+                info!("Removed worker: {}", worker_url);
+            }
+        }
+
+        // if cache aware, remove the worker from the tree
+        if let Router::CacheAware { tree, .. } = self {
+            tree.lock().unwrap().remove_tenant(&worker_url);
+            info!("Removed worker from tree: {}", worker_url);
+        }
+    }
 }

rust/src/server.rs

@@ -145,6 +145,19 @@ async fn add_worker(
     HttpResponse::Ok().finish()
 }
 
+#[post("/remove_worker")]
+async fn remove_worker(
+    query: web::Query<HashMap<String, String>>,
+    data: web::Data<AppState>,
+) -> impl Responder {
+    let worker_url = match query.get("url") {
+        Some(url) => url.to_string(),
+        None => return HttpResponse::BadRequest().finish(),
+    };
+    data.router.remove_worker(worker_url);
+    HttpResponse::Ok().finish()
+}
+
 pub struct ServerConfig {
     pub host: String,
     pub port: u16,
@@ -201,6 +214,7 @@ pub async fn startup(config: ServerConfig) -> std::io::Result<()> {
             .service(health_generate)
             .service(get_server_info)
             .service(add_worker)
+            .service(remove_worker)
     })
     .bind((config.host, config.port))?
     .run()