github.com/tonyhb/nomad@v0.11.8/helper/gated-writer/writer.go (about)

     1  package gatedwriter
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  )
     7  
     8  // Writer is an io.Writer implementation that buffers all of its
     9  // data into an internal buffer until it is told to let data through.
    10  type Writer struct {
    11  	Writer io.Writer
    12  
    13  	buf   [][]byte
    14  	flush bool
    15  	lock  sync.RWMutex
    16  }
    17  
    18  // Flush tells the Writer to flush any buffered data and to stop
    19  // buffering.
    20  func (w *Writer) Flush() {
    21  	w.lock.Lock()
    22  	w.flush = true
    23  	w.lock.Unlock()
    24  
    25  	for _, p := range w.buf {
    26  		w.Write(p)
    27  	}
    28  	w.buf = nil
    29  }
    30  
    31  func (w *Writer) Write(p []byte) (n int, err error) {
    32  	w.lock.RLock()
    33  	defer w.lock.RUnlock()
    34  
    35  	if w.flush {
    36  		return w.Writer.Write(p)
    37  	}
    38  
    39  	p2 := make([]byte, len(p))
    40  	copy(p2, p)
    41  	w.buf = append(w.buf, p2)
    42  	return len(p), nil
    43  }