diff --git a/internal/container/run.go b/internal/container/run.go index 04bedc0652..3fe096a750 100644 --- a/internal/container/run.go +++ b/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 +} diff --git a/logging/prefix_writer.go b/logging/prefix_writer.go index bdd01f4e6c..6c3f999e67 100644 --- a/logging/prefix_writer.go +++ b/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 } diff --git a/logging/prefix_writer_test.go b/logging/prefix_writer_test.go index bf8b3c3551..6cedc21552 100644 --- a/logging/prefix_writer_test.go +++ b/logging/prefix_writer_test.go @@ -26,7 +26,9 @@ 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() { @@ -34,12 +36,20 @@ func testPrefixWriter(t *testing.T, when spec.G, it spec.S) { 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.") }) }) }