github.com/billybanfield/evergreen@v0.0.0-20170525200750-eeee692790f7/util/line_buffer_writer.go (about)

     1  package util
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"sync"
     7  )
     8  
     9  //LineBufferingWriter is an implementation of io.Writer that sends the newline
    10  //delimited subslices of its inputs to a wrapped io.Writer, buffering as needed.
    11  type LineBufferingWriter struct {
    12  	buf           []byte
    13  	wrappedWriter io.Writer
    14  	bufferLock    sync.Mutex
    15  }
    16  
    17  func NewLineBufferingWriter(wrapped io.Writer) *LineBufferingWriter {
    18  	return &LineBufferingWriter{
    19  		buf:           []byte{},
    20  		wrappedWriter: wrapped,
    21  	}
    22  }
    23  
    24  // writes whatever is in the buffer out using the wrapped io.Writer
    25  func (lbw *LineBufferingWriter) Flush() error {
    26  	lbw.bufferLock.Lock()
    27  	defer lbw.bufferLock.Unlock()
    28  	if len(lbw.buf) == 0 {
    29  		return nil
    30  	}
    31  
    32  	_, err := lbw.wrappedWriter.Write(lbw.buf)
    33  	if err != nil {
    34  		return err
    35  	}
    36  	lbw.buf = []byte{}
    37  	return nil
    38  }
    39  
    40  // use the wrapped io.Writer to write anything that is delimited with a newline
    41  func (lbw *LineBufferingWriter) Write(p []byte) (n int, err error) {
    42  	lbw.bufferLock.Lock()
    43  	defer lbw.bufferLock.Unlock()
    44  
    45  	// Check to see if sum of two buffers is greater than 4K bytes
    46  	if len(p)+len(lbw.buf) >= (4*1024) && len(lbw.buf) > 0 {
    47  		_, err := lbw.wrappedWriter.Write(lbw.buf)
    48  		if err != nil {
    49  			return 0, err
    50  		}
    51  		lbw.buf = []byte{}
    52  	}
    53  
    54  	fullBuffer := append(lbw.buf, p...)
    55  	lines := bytes.Split(fullBuffer, []byte{'\n'})
    56  	for idx, val := range lines {
    57  		if idx == len(lines)-1 {
    58  			lbw.buf = val
    59  		} else {
    60  			_, err := lbw.wrappedWriter.Write(val)
    61  			if err != nil {
    62  				return 0, err
    63  			}
    64  		}
    65  	}
    66  
    67  	return len(p), nil
    68  }