github.com/buildpacks/pack@v0.33.3-0.20240516162812-884dd1837311/pkg/logging/prefix_writer.go (about)

     1  package logging
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"fmt"
     7  	"io"
     8  
     9  	"github.com/buildpacks/pack/internal/style"
    10  )
    11  
    12  // PrefixWriter is a buffering writer that prefixes each new line. Close should be called to properly flush the buffer.
    13  type PrefixWriter struct {
    14  	out           io.Writer
    15  	buf           *bytes.Buffer
    16  	prefix        string
    17  	readerFactory func(data []byte) io.Reader
    18  }
    19  
    20  type PrefixWriterOption func(c *PrefixWriter)
    21  
    22  func WithReaderFactory(factory func(data []byte) io.Reader) PrefixWriterOption {
    23  	return func(writer *PrefixWriter) {
    24  		writer.readerFactory = factory
    25  	}
    26  }
    27  
    28  // NewPrefixWriter writes by w will be prefixed
    29  func NewPrefixWriter(w io.Writer, prefix string, opts ...PrefixWriterOption) *PrefixWriter {
    30  	writer := &PrefixWriter{
    31  		out:    w,
    32  		prefix: fmt.Sprintf("[%s] ", style.Prefix(prefix)),
    33  		buf:    &bytes.Buffer{},
    34  		readerFactory: func(data []byte) io.Reader {
    35  			return bytes.NewReader(data)
    36  		},
    37  	}
    38  
    39  	for _, opt := range opts {
    40  		opt(writer)
    41  	}
    42  
    43  	return writer
    44  }
    45  
    46  // Write writes bytes to the embedded log function
    47  func (w *PrefixWriter) Write(data []byte) (int, error) {
    48  	scanner := bufio.NewScanner(w.readerFactory(data))
    49  	scanner.Split(ScanLinesKeepNewLine)
    50  	for scanner.Scan() {
    51  		newBits := scanner.Bytes()
    52  		if len(newBits) > 0 && newBits[len(newBits)-1] != '\n' { // just append if we don't have a new line
    53  			_, err := w.buf.Write(newBits)
    54  			if err != nil {
    55  				return 0, err
    56  			}
    57  		} else { // write our complete message
    58  			_, err := w.buf.Write(bytes.TrimRight(newBits, "\n"))
    59  			if err != nil {
    60  				return 0, err
    61  			}
    62  
    63  			err = w.flush()
    64  			if err != nil {
    65  				return 0, err
    66  			}
    67  		}
    68  	}
    69  
    70  	if err := scanner.Err(); err != nil {
    71  		return 0, err
    72  	}
    73  
    74  	return len(data), nil
    75  }
    76  
    77  // Close writes any pending data in the buffer
    78  func (w *PrefixWriter) Close() error {
    79  	if w.buf.Len() > 0 {
    80  		return w.flush()
    81  	}
    82  
    83  	return nil
    84  }
    85  
    86  func (w *PrefixWriter) flush() error {
    87  	bits := w.buf.Bytes()
    88  	w.buf.Reset()
    89  
    90  	// process any CR in message
    91  	if i := bytes.LastIndexByte(bits, '\r'); i >= 0 {
    92  		bits = bits[i+1:]
    93  	}
    94  
    95  	_, err := fmt.Fprint(w.out, w.prefix+string(bits)+"\n")
    96  	return err
    97  }
    98  
    99  // A customized implementation of bufio.ScanLines that preserves new line characters.
   100  func ScanLinesKeepNewLine(data []byte, atEOF bool) (advance int, token []byte, err error) {
   101  	if atEOF && len(data) == 0 {
   102  		return 0, nil, nil
   103  	}
   104  
   105  	// first we'll split by LF (\n)
   106  	// then remove any preceding CR (\r) [due to CR+LF]
   107  	if i := bytes.IndexByte(data, '\n'); i >= 0 {
   108  		// We have a full newline-terminated line.
   109  		return i + 1, append(dropCR(data[0:i]), '\n'), nil
   110  	}
   111  
   112  	// If we're at EOF, we have a final, non-terminated line. Return it.
   113  	if atEOF {
   114  		return len(data), data, nil
   115  	}
   116  	// Request more data.
   117  	return 0, nil, nil
   118  }
   119  
   120  // dropCR drops a terminal \r from the data.
   121  func dropCR(data []byte) []byte {
   122  	if len(data) > 0 && data[len(data)-1] == '\r' {
   123  		return data[0 : len(data)-1]
   124  	}
   125  	return data
   126  }