allow customizing log prefixes for wasmtime serve command (#9821)

* allow customizing log prefixes for wasmtime serve command Signed-off-by: Joseph Zhang <jzhang@absolute.com> * pr feedback - use simple boolean flag instead --------- Signed-off-by: Joseph Zhang <jzhang@absolute.com>
bytecodealliance · Dec 16, 2024 · d3f05ee · d3f05ee
1 parent 392df4a
commit d3f05ee
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 10 deletions.
diff --git a/src/commands/serve.rs b/src/commands/serve.rs
@@ -88,9 +88,14 @@ pub struct ServeCommand {
     run: RunCommon,
 
     /// Socket address for the web server to bind to.
-    #[arg(long = "addr", value_name = "SOCKADDR", default_value_t = DEFAULT_ADDR )]
+    #[arg(long = "addr", value_name = "SOCKADDR", default_value_t = DEFAULT_ADDR)]
     addr: SocketAddr,
 
+    /// Disable log prefixes of wasi-http handlers.
+    /// if unspecified, logs will be prefixed with 'stdout|stderr [{req_id}] :: '
+    #[arg(long = "no-logging-prefix")]
+    no_logging_prefix: bool,
+
     /// The WebAssembly component to run.
     #[arg(value_name = "WASM", required = true)]
     component: PathBuf,
@@ -153,15 +158,17 @@ impl ServeCommand {
 
         builder.env("REQUEST_ID", req_id.to_string());
 
-        builder.stdout(LogStream::new(
-            format!("stdout [{req_id}] :: "),
-            Output::Stdout,
-        ));
-
-        builder.stderr(LogStream::new(
-            format!("stderr [{req_id}] :: "),
-            Output::Stderr,
-        ));
+        let stdout_prefix: String;
+        let stderr_prefix: String;
+        if self.no_logging_prefix {
+            stdout_prefix = "".to_string();
+            stderr_prefix = "".to_string();
+        } else {
+            stdout_prefix = format!("stdout [{req_id}] :: ");
+            stderr_prefix = format!("stderr [{req_id}] :: ");
+        }
+        builder.stdout(LogStream::new(stdout_prefix, Output::Stdout));
+        builder.stderr(LogStream::new(stderr_prefix, Output::Stderr));
 
         let mut host = Host {
             table: wasmtime::component::ResourceTable::new(),

diff --git a/tests/all/cli_tests.rs b/tests/all/cli_tests.rs
@@ -1909,6 +1909,52 @@ stderr [1] :: after empty
         Ok(())
     }
 
+    #[tokio::test]
+    async fn cli_serve_with_print_no_prefix() -> Result<()> {
+        let server = WasmtimeServe::new(CLI_SERVE_WITH_PRINT_COMPONENT, |cmd| {
+            cmd.arg("-Scli");
+            cmd.arg("--no-logging-prefix");
+        })?;
+
+        for _ in 0..2 {
+            let resp = server
+                .send_request(
+                    hyper::Request::builder()
+                        .uri("http://localhost/")
+                        .body(String::new())
+                        .context("failed to make request")?,
+                )
+                .await?;
+            assert!(resp.status().is_success());
+        }
+
+        let (out, err) = server.finish()?;
+        assert_eq!(
+            out,
+            "\
+this is half a print to stdout
+\n\
+after empty
+this is half a print to stdout
+\n\
+after empty
+"
+        );
+        assert_eq!(
+            err,
+            "\
+this is half a print to stderr
+\n\
+after empty
+this is half a print to stderr
+\n\
+after empty
+"
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn cli_serve_authority_and_scheme() -> Result<()> {
         let server = WasmtimeServe::new(CLI_SERVE_AUTHORITY_AND_SCHEME_COMPONENT, |cmd| {