From 57f8bd53d48ee63fb0a74bb3899b71d25cacad92 Mon Sep 17 00:00:00 2001
From: Chris Olszewski <chris.olszewski@vercel.com>
Date: Tue, 26 Sep 2023 12:42:50 -0700
Subject: [PATCH] fix: prefixed writer no longer includes prefix length in
 bytes written (#6032)

### Description

When hooking up the prefixed writer as part of task graph execution I
encountered:
```
[0 olszewski@chriss-mbp] /tmp/pnpm-test $ turbo_dev build --experimental-rust-codepath
docs:build: next build
web:build: next build
docs:build:
thread 'tokio-runtime-worker' panicked at /rustc/3223b0b5e8dadda3f76c3fd1a8d6c5addc09599e/library/std/src/io/mod.rs:1581:36:
range start index 22 out of range for slice of length 1
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
web:build:
thread 'tokio-runtime-worker' panicked at /rustc/3223b0b5e8dadda3f76c3fd1a8d6c5addc09599e/library/std/src/io/mod.rs:1581:36:
range start index 21 out of range for slice of length 1
thread 'main' panicked at crates/turborepo-lib/src/task_graph/visitor.rs:259:20:
task executor panicked: JoinError::Panic(Id(28), ...)
```

After a bit of digging, I found that this was caused by us violating
part of the `Write` trait, specifically our `write` return value:
> If the return value is Ok(n) then n must satisfy n <= buf.len()

Where we would return `prefix.len() + buf.len()` which upstream would
cause and index out of bounds panic.

This PR:
- Now returns `buf.len()` in `write` so we return values that callers
expect
- Refactor `PrefixedWriter` to eagerly apply the styled writer instead
of applying it on every write call.

### Testing Instructions

Added quick unit tests. Real test for the panic fix is use in the
process manager hookup PR:
https://github.com/vercel/turbo/tree/olszewski/hook_up_process_manager


Closes TURBO-1363

Co-authored-by: Chris Olszewski <Chris Olszewski>
---
 crates/turborepo-ui/src/prefixed.rs | 31 ++++++++++++++++++++---------
 1 file changed, 22 insertions(+), 9 deletions(-)

diff --git a/crates/turborepo-ui/src/prefixed.rs b/crates/turborepo-ui/src/prefixed.rs
index 82a4b00dfb8e9..fdaa0fcb52d40 100644
--- a/crates/turborepo-ui/src/prefixed.rs
+++ b/crates/turborepo-ui/src/prefixed.rs
@@ -94,33 +94,33 @@ enum Command {
 
 /// Wraps a writer with a prefix before the actual message.
 pub struct PrefixedWriter<W> {
-    prefix: StyledObject<String>,
+    prefix: String,
     writer: W,
-    ui: UI,
 }
 
 impl<W> Debug for PrefixedWriter<W> {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("PrefixedWriter")
             .field("prefix", &self.prefix)
-            .field("ui", &self.ui)
             .finish()
     }
 }
 
 impl<W: Write> PrefixedWriter<W> {
-    pub fn new(ui: UI, prefix: StyledObject<String>, writer: W) -> Self {
-        Self { ui, prefix, writer }
+    pub fn new(ui: UI, prefix: StyledObject<impl Display>, writer: W) -> Self {
+        let prefix = ui.apply(prefix).to_string();
+        Self { prefix, writer }
     }
 }
 
 impl<W: Write> Write for PrefixedWriter<W> {
     fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
-        let prefix = self.prefix.clone();
-        let prefix = self.ui.apply(prefix);
-        let prefix_bytes_written = self.writer.write(prefix.to_string().as_bytes())?;
+        self.writer.write_all(self.prefix.as_bytes())?;
 
-        Ok(prefix_bytes_written + self.writer.write(buf)?)
+        // We do end up writing more bytes than this to the underlying writer, but we
+        // cannot report this to the callers as the amount of bytes we report
+        // written must be less than or equal to the number of bytes in the buffer.
+        self.writer.write(buf)
     }
 
     fn flush(&mut self) -> std::io::Result<()> {
@@ -166,4 +166,17 @@ mod test {
         };
         assert_eq!(String::from_utf8(buffer).unwrap(), expected);
     }
+
+    #[test_case(true, "foo#build: cool!")]
+    #[test_case(false, "\u{1b}[1mfoo#build: \u{1b}[0mcool!")]
+    fn test_prefixed_writer(strip_ansi: bool, expected: &str) {
+        let mut buffer = Vec::new();
+        let mut writer = PrefixedWriter::new(
+            UI::new(strip_ansi),
+            crate::BOLD.apply_to("foo#build: "),
+            &mut buffer,
+        );
+        writer.write_all(b"cool!").unwrap();
+        assert_eq!(String::from_utf8(buffer).unwrap(), expected);
+    }
 }