linewriter.go raw

   1  // Copyright 2018 The gVisor Authors.
   2  //
   3  // Licensed under the Apache License, Version 2.0 (the "License");
   4  // you may not use this file except in compliance with the License.
   5  // You may obtain a copy of the License at
   6  //
   7  //     http://www.apache.org/licenses/LICENSE-2.0
   8  //
   9  // Unless required by applicable law or agreed to in writing, software
  10  // distributed under the License is distributed on an "AS IS" BASIS,
  11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12  // See the License for the specific language governing permissions and
  13  // limitations under the License.
  14  
  15  // Package linewriter provides an io.Writer which calls an emitter on each line.
  16  package linewriter
  17  
  18  import (
  19  	"bytes"
  20  
  21  	"gvisor.dev/gvisor/pkg/sync"
  22  )
  23  
  24  // Writer is an io.Writer which buffers input, flushing
  25  // individual lines through an emitter function.
  26  type Writer struct {
  27  	// the mutex locks buf.
  28  	sync.Mutex
  29  
  30  	// buf holds the data we haven't emitted yet.
  31  	buf bytes.Buffer
  32  
  33  	// emit is used to flush individual lines.
  34  	emit func(p []byte)
  35  }
  36  
  37  // NewWriter creates a Writer which emits using emitter.
  38  // The emitter must not retain p. It may change after emitter returns.
  39  func NewWriter(emitter func(p []byte)) *Writer {
  40  	return &Writer{emit: emitter}
  41  }
  42  
  43  // Write implements io.Writer.Write.
  44  // It calls emit on each line of input, not including the newline.
  45  // Write may be called concurrently.
  46  func (w *Writer) Write(p []byte) (int, error) {
  47  	w.Lock()
  48  	defer w.Unlock()
  49  
  50  	total := 0
  51  	for len(p) > 0 {
  52  		emit := true
  53  		i := bytes.IndexByte(p, '\n')
  54  		if i < 0 {
  55  			// No newline, we will buffer everything.
  56  			i = len(p)
  57  			emit = false
  58  		}
  59  
  60  		n, err := w.buf.Write(p[:i])
  61  		if err != nil {
  62  			return total, err
  63  		}
  64  		total += n
  65  
  66  		p = p[i:]
  67  
  68  		if emit {
  69  			// Skip the newline, but still count it.
  70  			p = p[1:]
  71  			total++
  72  
  73  			w.emit(w.buf.Bytes())
  74  			w.buf.Reset()
  75  		}
  76  	}
  77  
  78  	return total, nil
  79  }
  80