Merge pull request #867 from buildpacks/fix/865-log-prefix

Fix lifecycle phase output

internal/container/run.go

@@ -32,6 +32,8 @@ func Run(ctx context.Context, docker client.CommonAPIClient, ctrID string, out,
 	copyErr := make(chan error)
 	go func() {
 		_, err := stdcopy.StdCopy(out, errOut, resp.Reader)
+		defer optionallyCloseWriter(out)
+		defer optionallyCloseWriter(errOut)
 
 		copyErr <- err
 	}()
@@ -47,3 +49,11 @@ func Run(ctx context.Context, docker client.CommonAPIClient, ctrID string, out,
 
 	return <-copyErr
 }
+
+func optionallyCloseWriter(writer io.Writer) error {
+	if closer, ok := writer.(io.Closer); ok {
+		return closer.Close()
+	}
+
+	return nil
+}

logging/prefix_writer.go

@@ -9,9 +9,10 @@ import (
 	"github.com/buildpacks/pack/internal/style"
 )
 
-// PrefixWriter will prefix writes
+// PrefixWriter is a buffering writer that prefixes each new line. Close should be called to properly flush the buffer.
 type PrefixWriter struct {
 	out    io.Writer
+	buf    *bytes.Buffer
 	prefix string
 }
 
@@ -20,15 +21,80 @@ func NewPrefixWriter(w io.Writer, prefix string) *PrefixWriter {
 	return &PrefixWriter{
 		out:    w,
 		prefix: fmt.Sprintf("[%s] ", style.Prefix(prefix)),
+		buf:    &bytes.Buffer{},
 	}
 }
 
-// Writes bytes to the embedded log function
-func (w *PrefixWriter) Write(buf []byte) (int, error) {
-	scanner := bufio.NewScanner(bytes.NewReader(buf))
+// Write writes bytes to the embedded log function
+func (w *PrefixWriter) Write(data []byte) (int, error) {
+	scanner := bufio.NewScanner(bytes.NewReader(data))
+	scanner.Split(ScanLinesKeepNewLine)
 	for scanner.Scan() {
-		_, _ = fmt.Fprintln(w.out, w.prefix+scanner.Text())
+		newBits := scanner.Bytes()
+		if newBits[len(newBits)-1] != '\n' { // just append if we don't have a new line
+			_, err := w.buf.Write(newBits)
+			if err != nil {
+				return 0, err
+			}
+		} else { // write our complete message
+			var allBits []byte
+			if w.buf.Len() > 0 {
+				allBits = append(w.buf.Bytes(), newBits...)
+				w.buf.Reset()
+			} else {
+				allBits = newBits
+			}
+
+			err := w.writeWithPrefix(allBits)
+			if err != nil {
+				return 0, err
+			}
+		}
+	}
+
+	return len(data), nil
+}
+
+// Close writes any pending data in the buffer
+func (w *PrefixWriter) Close() error {
+	if w.buf.Len() > 0 {
+		err := w.writeWithPrefix(w.buf.Bytes())
+		if err != nil {
+			return err
+		}
+	}
+
+	w.buf.Reset()
+
+	return nil
+}
+
+func (w *PrefixWriter) writeWithPrefix(bits []byte) error {
+	_, err := fmt.Fprint(w.out, w.prefix+string(bits))
+	return err
+}
+
+// A customized implementation of bufio.ScanLines that preserves new line characters.
+func ScanLinesKeepNewLine(data []byte, atEOF bool) (advance int, token []byte, err error) {
+	if atEOF && len(data) == 0 {
+		return 0, nil, nil
 	}
+	if i := bytes.IndexByte(data, '\n'); i >= 0 {
+		// We have a full newline-terminated line.
+		return i + 1, append(dropCR(data[0:i]), '\n'), nil
+	}
+	// If we're at EOF, we have a final, non-terminated line. Return it.
+	if atEOF {
+		return len(data), dropCR(data), nil
+	}
+	// Request more data.
+	return 0, nil, nil
+}
 
-	return len(buf), nil
+// dropCR drops a terminal \r from the data.
+func dropCR(data []byte) []byte {
+	if len(data) > 0 && data[len(data)-1] == '\r' {
+		return data[0 : len(data)-1]
+	}
+	return data
 }

logging/prefix_writer_test.go

@@ -26,20 +26,30 @@ func testPrefixWriter(t *testing.T, when spec.G, it spec.S) {
 			prefix := "test prefix"
 			writer := logging.NewPrefixWriter(&w, prefix)
 			_, _ = writer.Write([]byte("test"))
-			h.AssertEq(t, w.String(), fmt.Sprintf("[%s] %s", prefix, "test\n"))
+			_ = writer.Close()
+
+			h.AssertEq(t, w.String(), fmt.Sprintf("[%s] %s", prefix, "test"))
 		})
 
 		it("prepends prefix to multi-line string", func() {
 			var w bytes.Buffer
 
 			writer := logging.NewPrefixWriter(&w, "prefix")
 			_, _ = writer.Write([]byte("line 1\nline 2\nline 3"))
-			h.AssertEq(t,
-				w.String(),
-				`[prefix] line 1
-[prefix] line 2
-[prefix] line 3
-`)
+			_ = writer.Close()
+			h.AssertEq(t, w.String(), "[prefix] line 1\n[prefix] line 2\n[prefix] line 3")
+		})
+
+		it("buffers mid-line calls", func() {
+			var buf bytes.Buffer
+
+			writer := logging.NewPrefixWriter(&buf, "prefix")
+			_, _ = writer.Write([]byte("word 1, "))
+			_, _ = writer.Write([]byte("word 2, "))
+			_, _ = writer.Write([]byte("word 3."))
+			_ = writer.Close()
+
+			h.AssertEq(t, buf.String(), "[prefix] word 1, word 2, word 3.")
 		})
 	})
 }