From 655817231b544385b6a6e57bf26c63738da2f87e Mon Sep 17 00:00:00 2001
From: Robert Escriva <robert@trychroma.com>
Date: Mon, 16 Dec 2024 16:23:08 -0800
Subject: [PATCH] [ENH]  Make the error messages out of chroma-load
 decipherable. (#3316)

Staring was failing with a 400 that was really a 404.  This PR changes
that case to a 404 instead of 400 and prints the body in the
chroma-load-start tool.  Also, a tool to stop all workloads.
---
 rust/load/src/bin/chroma-load-start.rs    |  5 +-
 rust/load/src/bin/chroma-load-stop-all.rs | 60 +++++++++++++++++++++++
 rust/load/src/lib.rs                      |  3 +-
 3 files changed, 65 insertions(+), 3 deletions(-)
 create mode 100644 rust/load/src/bin/chroma-load-stop-all.rs

diff --git a/rust/load/src/bin/chroma-load-start.rs b/rust/load/src/bin/chroma-load-start.rs
index c6b54208362..d2d715611b7 100644
--- a/rust/load/src/bin/chroma-load-start.rs
+++ b/rust/load/src/bin/chroma-load-start.rs
@@ -116,10 +116,13 @@ async fn main() {
                 );
             } else {
                 eprintln!(
-                    "Failed to start workload on {}: {}",
+                    "Categorically failed to start workload on {}: {}",
                     args.host,
                     resp.status()
                 );
+                if let Ok(text) = resp.text().await {
+                    eprintln!("{}", text.trim());
+                }
             }
         }
         Err(e) => eprintln!("Failed to start workload on {}: {}", args.host, e),
diff --git a/rust/load/src/bin/chroma-load-stop-all.rs b/rust/load/src/bin/chroma-load-stop-all.rs
new file mode 100644
index 00000000000..68d5e9dc78c
--- /dev/null
+++ b/rust/load/src/bin/chroma-load-stop-all.rs
@@ -0,0 +1,60 @@
+//! Stop all workloads on the chroma-load server.
+//!
+//! If you are looking to stop traffic for a SEV, see chroma-load-inhibit.
+
+use clap::Parser;
+
+use chroma_load::rest::StopRequest;
+
+#[derive(Parser, Debug)]
+struct Args {
+    #[arg(long)]
+    host: String,
+}
+
+#[tokio::main]
+async fn main() {
+    let args = Args::parse();
+    let client = reqwest::Client::new();
+    match client
+        .get(format!("{}/", args.host))
+        .header(reqwest::header::ACCEPT, "application/json")
+        .send()
+        .await
+    {
+        Ok(resp) => {
+            let resp = resp.error_for_status().expect("Failed to get status");
+            let resp = resp
+                .json::<chroma_load::rest::Status>()
+                .await
+                .expect("Failed to parse status");
+            for workload in resp.running {
+                let req = StopRequest {
+                    uuid: workload.uuid,
+                };
+                match client
+                    .post(format!("{}/stop", args.host))
+                    .json(&req)
+                    .send()
+                    .await
+                {
+                    Ok(resp) => {
+                        if resp.status().is_success() {
+                            println!("Stopped workload on {}", args.host);
+                        } else {
+                            eprintln!(
+                                "Failed to stop workload on {}: {}",
+                                args.host,
+                                resp.status()
+                            );
+                        }
+                    }
+                    Err(e) => eprintln!("Failed to stop workload on {}: {}", args.host, e),
+                }
+            }
+        }
+        Err(e) => {
+            eprintln!("Failed to get status: {}", e);
+        }
+    }
+}
diff --git a/rust/load/src/lib.rs b/rust/load/src/lib.rs
index 388fa594c15..907a1e71d7d 100644
--- a/rust/load/src/lib.rs
+++ b/rust/load/src/lib.rs
@@ -478,7 +478,7 @@ impl Workload {
                 if let Some(workload) = workloads.get(name) {
                     *self = workload.clone();
                 } else {
-                    return Err(Error::InvalidRequest(format!("workload not found: {name}")));
+                    return Err(Error::NotFound(format!("workload not found: {name}")));
                 }
             }
             Workload::Get(_) => {}
@@ -1412,7 +1412,6 @@ pub async fn entrypoint() {
                     "http_request",
                     method = ?request.method(),
                     matched_path,
-                    some_other_field = tracing::field::Empty,
                 )
             }),
         )