github.com/hernad/nomad@v1.6.112/helper/gated-writer/writer.go (about)

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