github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/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  	defer w.lock.Unlock()
    23  
    24  	w.flush = true
    25  
    26  	for _, p := range w.buf {
    27  		_, _ = w.Writer.Write(p)
    28  	}
    29  
    30  	w.buf = nil
    31  }
    32  
    33  func (w *Writer) Write(p []byte) (n int, err error) {
    34  	w.lock.Lock()
    35  	defer w.lock.Unlock()
    36  
    37  	if w.flush {
    38  		return w.Writer.Write(p)
    39  	}
    40  
    41  	p2 := make([]byte, len(p))
    42  	copy(p2, p)
    43  	w.buf = append(w.buf, p2)
    44  	return len(p), nil
    45  }