github.com/grailbio/base@v0.0.11/iofmt/linewriter.go (about)

     1  // Copyright 2021 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package iofmt
     6  
     7  import (
     8  	"bytes"
     9  	"io"
    10  )
    11  
    12  type lineWriter struct {
    13  	w   io.Writer
    14  	buf []byte
    15  }
    16  
    17  // LineWriter returns an io.WriteCloser that only calls w.Write with
    18  // complete lines.  This can be used to make it less likely (without
    19  // locks) for lines to interleave, for example if you are concurrently
    20  // writing lines of text to os.Stdout. This is particularly useful when
    21  // composed with PrefixWriter.
    22  //
    23  //  // Full lines will be written to os.Stdout, so they will be less likely to
    24  //  // be interleaved with other output.
    25  //  linew := LineWriter(os.Stdout)
    26  //  defer func() {
    27  //		_ = linew.Close() // Handle the possible error.
    28  //  }()
    29  //  w := PrefixWriter(linew, "my-prefix: ")
    30  //
    31  // Close will write any remaining partial line to the underlying writer.
    32  func LineWriter(w io.Writer) io.WriteCloser {
    33  	return &lineWriter{w: w}
    34  }
    35  
    36  func (w *lineWriter) Write(p []byte) (int, error) {
    37  	var n int
    38  	for {
    39  		i := bytes.Index(p, newline)
    40  		// TODO(jcharumilind): Limit buffer size.
    41  		switch i {
    42  		case -1:
    43  			w.buf = append(w.buf, p...)
    44  			return n + len(p), nil
    45  		default:
    46  			var err error
    47  			if len(w.buf) > 0 {
    48  				w.buf = append(w.buf, p[:i+1]...)
    49  				_, err = w.w.Write(w.buf)
    50  				w.buf = w.buf[:0]
    51  			} else {
    52  				_, err = w.w.Write(p[:i+1])
    53  			}
    54  			n += i + 1
    55  			if err != nil {
    56  				return n, err
    57  			}
    58  			p = p[i+1:]
    59  		}
    60  	}
    61  }
    62  
    63  func (w *lineWriter) Close() error {
    64  	if len(w.buf) == 0 {
    65  		return nil
    66  	}
    67  	_, err := w.w.Write(w.buf)
    68  	w.buf = nil
    69  	return err
    70  }