wip:milestone 0 fixes

gateway/src/proxy.rs

@@ -0,0 +1,265 @@
+use axum::{
+    body::Body,
+    extract::{Request, State},
+    http::StatusCode,
+    response::Response,
+    routing::get,
+    Router,
+};
+use std::net::SocketAddr;
+use std::sync::Arc;
+use tokio::sync::RwLock;
+use tracing::{error, info};
+
+#[derive(Clone, Debug)]
+struct Upstream {
+    name: String,
+    url: String,
+    healthy: Arc<RwLock<bool>>,
+}
+
+impl Upstream {
+    fn new(name: String, url: String) -> Self {
+        Self {
+            name,
+            url,
+            healthy: Arc::new(RwLock::new(true)),
+        }
+    }
+}
+
+#[derive(Clone)]
+struct ProxyState {
+    control_upstream: Upstream,
+    worker_upstreams: Arc<RwLock<Vec<Upstream>>>,
+    current_worker_index: Arc<RwLock<usize>>,
+}
+
+impl ProxyState {
+    fn new(control_url: String, worker_urls: Vec<String>) -> Self {
+        let worker_upstreams = worker_urls
+            .into_iter()
+            .map(|url| Upstream::new(format!("worker-{}", url), url))
+            .collect();
+
+        Self {
+            control_upstream: Upstream::new("control".to_string(), control_url),
+            worker_upstreams: Arc::new(RwLock::new(worker_upstreams)),
+            current_worker_index: Arc::new(RwLock::new(0)),
+        }
+    }
+
+    async fn get_next_worker(&self) -> Option<Upstream> {
+        let upstreams = self.worker_upstreams.read().await;
+        let current_len = upstreams.len();
+        
+        if current_len == 0 {
+            return None;
+        }
+
+        let mut index = self.current_worker_index.write().await;
+        let selected = upstreams[*index % current_len].clone();
+        *index = (*index + 1) % current_len;
+        
+        Some(selected)
+    }
+
+    async fn get_healthy_worker(&self) -> Option<Upstream> {
+        let upstreams = self.worker_upstreams.read().await;
+        
+        for upstream in upstreams.iter() {
+            let is_healthy = *upstream.healthy.read().await;
+            if is_healthy {
+                return Some(upstream.clone());
+            }
+        }
+        None
+    }
+
+    async fn start_health_check_loop(&self) {
+        let mut interval = tokio::time::interval(std::time::Duration::from_secs(5));
+        info!("Starting proxy health check loop");
+
+        loop {
+            interval.tick().await;
+
+            // Check workers
+            let worker_upstreams = self.worker_upstreams.read().await;
+            for worker in worker_upstreams.iter() {
+                let worker = worker.clone();
+                tokio::spawn(async move {
+                    let client = reqwest::Client::builder()
+                        .timeout(std::time::Duration::from_secs(2))
+                        .build()
+                        .unwrap();
+
+                    let res = client.get(format!("{}/health", worker.url)).send().await;
+                    let is_healthy = res.is_ok() && res.unwrap().status().is_success();
+                    
+                    let mut healthy = worker.healthy.write().await;
+                    if *healthy != is_healthy {
+                        if is_healthy {
+                            info!("Worker {} is now healthy", worker.url);
+                        } else {
+                            error!("Worker {} is now unhealthy", worker.url);
+                        }
+                    }
+                    *healthy = is_healthy;
+                });
+            }
+
+            // Check control plane
+            let control = self.control_upstream.clone();
+            tokio::spawn(async move {
+                let client = reqwest::Client::builder()
+                    .timeout(std::time::Duration::from_secs(2))
+                    .build()
+                    .unwrap();
+
+                let res = client.get(format!("{}/health", control.url)).send().await;
+                let is_healthy = res.is_ok() && res.unwrap().status().is_success();
+                
+                let mut healthy = control.healthy.write().await;
+                if *healthy != is_healthy {
+                    if is_healthy {
+                        info!("Control plane {} is now healthy", control.url);
+                    } else {
+                        error!("Control plane {} is now unhealthy", control.url);
+                    }
+                }
+                *healthy = is_healthy;
+            });
+        }
+    }
+}
+
+async fn proxy_request(
+    State(state): State<ProxyState>,
+    req: Request,
+) -> Result<Response, StatusCode> {
+    let path = req.uri().path();
+    
+    // Route /platform/* to control plane
+    if path.starts_with("/platform") || path.starts_with("/dashboard") || path == "/login" {
+        return forward_request(state.control_upstream.clone(), req).await;
+    }
+
+    // Route /auth/v1, /rest/v1, /storage/v1, /realtime/v1, /functions/v1 to workers
+    if path.starts_with("/auth/v1") 
+        || path.starts_with("/rest/v1") 
+        || path.starts_with("/storage/v1") 
+        || path.starts_with("/realtime/v1")
+        || path.starts_with("/functions/v1") {
+        
+        // Try to get a healthy worker, fall back to round-robin
+        let mut selected_worker = state.get_healthy_worker().await;
+        if selected_worker.is_none() {
+            selected_worker = state.get_next_worker().await;
+        }
+        
+        if let Some(upstream) = selected_worker {
+            forward_request(upstream, req).await
+        } else {
+            Err(StatusCode::SERVICE_UNAVAILABLE)
+        }
+    } else {
+        // Default to control plane
+        forward_request(state.control_upstream.clone(), req).await
+    }
+}
+
+async fn forward_request(upstream: Upstream, req: Request) -> Result<Response, StatusCode> {
+    let client = reqwest::Client::new();
+    
+    // Update the request URI
+    let original_uri = req.uri().clone();
+    let path_and_query = original_uri
+        .path_and_query()
+        .map(|pq| pq.as_str())
+        .unwrap_or("/");
+    
+    let target_url = format!("{}{}", upstream.url, path_and_query);
+    
+    info!("Proxying {} -> {}", original_uri.path(), target_url);
+    
+    // Build the request
+    let request_builder = client
+        .request(req.method().clone(), &target_url)
+        .headers(req.headers().clone());
+    
+    let response = request_builder
+        .send()
+        .await
+        .map_err(|e| {
+            error!("Failed to proxy request to {}: {}", upstream.name, e);
+            StatusCode::BAD_GATEWAY
+        })?;
+
+    let status = StatusCode::from_u16(response.status().as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
+    let headers = response.headers().clone();
+    let body_bytes = response.bytes().await.map_err(|e| {
+        error!("Failed to read response body from {}: {}", upstream.name, e);
+        StatusCode::BAD_GATEWAY
+    })?;
+
+    let mut response_builder = Response::builder().status(status);
+    
+    // Copy relevant headers
+    for (name, value) in headers.iter() {
+        if name != "connection" && name != "transfer-encoding" {
+            response_builder = response_builder.header(name, value);
+        }
+    }
+
+    response_builder
+        .body(Body::from(body_bytes.to_vec()))
+        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
+}
+
+async fn health_check() -> &'static str {
+    "OK"
+}
+
+pub async fn run() -> anyhow::Result<()> {
+    info!("Starting MadBase Proxy...");
+
+    let control_url = std::env::var("CONTROL_UPSTREAM_URL")
+        .unwrap_or_else(|_| "http://control:8001".to_string());
+    
+    let worker_urls_str = std::env::var("WORKER_UPSTREAM_URLS")
+        .unwrap_or_else(|_| "http://worker1:8002".to_string());
+    
+    let worker_urls: Vec<String> = worker_urls_str
+        .split(',')
+        .map(|s| s.trim().to_string())
+        .filter(|s| !s.is_empty())
+        .collect();
+
+    info!("Control upstream: {}", control_url);
+    info!("Worker upstreams: {:?}", worker_urls);
+
+    let state = ProxyState::new(control_url, worker_urls);
+
+    // Start health check loop in background
+    let state_clone = state.clone();
+    tokio::spawn(async move {
+        state_clone.start_health_check_loop().await;
+    });
+
+    let app = Router::new()
+        .route("/health", get(health_check))
+        .fallback(proxy_request)
+        .with_state(state);
+
+    let port = std::env::var("PROXY_PORT")
+        .unwrap_or_else(|_| "8000".to_string())
+        .parse::<u16>()?;
+    
+    let addr = SocketAddr::from(([0, 0, 0, 0], port));
+    info!("Proxy listening on {}", addr);
+
+    let listener = tokio::net::TcpListener::bind(addr).await?;
+    axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>()).await?;
+
+    Ok(())
+}