better 503 handling

jamsocket · Feb 10, 2025 · c3a1aa3 · c3a1aa3
1 parent 78e2846
commit c3a1aa3
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 15 deletions.
diff --git a/src/main.rs b/src/main.rs
@@ -120,7 +120,10 @@ async fn upload(
 }
 
 async fn wait_for_reload(State(state): State<ArcServerState>) -> (StatusCode, String) {
-    state.wrapped_server.wait_for_reload().await;
+    if !state.wrapped_server.running() {
+        state.wrapped_server.wait_for_reload().await;
+    }
+
     (StatusCode::OK, "OK".into())
 }
 

diff --git a/src/proxy.rs b/src/proxy.rs
@@ -29,33 +29,44 @@ fn translate_request(mut req: Request<Body>, port: u16) -> http::Request<SimpleB
     Request::from_parts(parts, body)
 }
 
+pub async fn render_error(
+    State(server_state): State<Arc<ServerState>>,
+    path: String,
+) -> Response<SimpleBody> {
+    let stdout = server_state.wrapped_server.stdout();
+    let handlebars = Handlebars::new();
+    let await_url = format!("{path}_frag/await");
+
+    let rendered = handlebars
+        .render_template(
+            ERROR_TEMPLATE,
+            &serde_json::json!({ "stdout": stdout, "await_url": await_url }),
+        )
+        .unwrap();
+
+    let body = to_simple_body(rendered);
+
+    return Response::builder().status(503).body(body).unwrap();
+}
+
 pub async fn proxy_request(
     State(server_state): State<Arc<ServerState>>,
     VerifiedPath(path): VerifiedPath,
     req: Request<Body>,
 ) -> Response<SimpleBody> {
     if !server_state.wrapped_server.running() {
-        let stdout = server_state.wrapped_server.stdout();
-        let handlebars = Handlebars::new();
-        let await_url = format!("{path}_frag/await");
-
-        let rendered = handlebars
-            .render_template(
-                ERROR_TEMPLATE,
-                &serde_json::json!({ "stdout": stdout, "await_url": await_url }),
-            )
-            .unwrap();
-
-        let body = to_simple_body(rendered);
-
-        return Response::builder().status(503).body(body).unwrap();
+        return render_error(State(server_state.clone()), path).await;
     }
 
     let client = ProxyClient::new();
     let req = translate_request(req, server_state.subprocess_port);
 
     let (response, handler) = client.request(req).await.expect("Infallable");
 
+    if !response.status().is_success() {
+        return render_error(State(server_state.clone()), path).await;
+    }
+
     if let Some(handler) = handler {
         // proxy websocket connection
         tracing::info!("Proxying websocket connection");