github.com/mutagen-io/mutagen@v0.18.0-rc1/pkg/stream/valve_writer.go (about)

     1  package stream
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  )
     7  
     8  // ValveWriter is an io.Writer that wraps another io.Writer and forwards writes
     9  // to it until the ValveWriter's internal valve is shut, after which writes will
    10  // continue to succeed but not actually be written to the underlying writer.
    11  type ValveWriter struct {
    12  	// writerLock serializes access to the underlying writer.
    13  	writerLock sync.Mutex
    14  	// writer is the underlying writer.
    15  	writer io.Writer
    16  }
    17  
    18  // NewValveWriter creates a new ValveWriter instance using the specified writer.
    19  // The writer may be nil, in which case the writer will start pre-shut.
    20  func NewValveWriter(writer io.Writer) *ValveWriter {
    21  	return &ValveWriter{writer: writer}
    22  }
    23  
    24  // Write implements io.Writer.Write.
    25  func (w *ValveWriter) Write(buffer []byte) (int, error) {
    26  	// Lock the writer and defer its release.
    27  	w.writerLock.Lock()
    28  	defer w.writerLock.Unlock()
    29  
    30  	// If there's no writer, then just pretend that we wrote all of the data.
    31  	if w.writer == nil {
    32  		return len(buffer), nil
    33  	}
    34  
    35  	// Otherwise write to the underlying writer.
    36  	return w.writer.Write(buffer)
    37  }
    38  
    39  // Shut closes the valve and prevents future writes to the underlying writer. It
    40  // is safe to call Shut concurrently with Write, but doing so will not preempt
    41  // or unblock pending calls to Write. Calling Shut will release the reference to
    42  // the underlying writer.
    43  func (w *ValveWriter) Shut() {
    44  	// Lock the writer and defer its release.
    45  	w.writerLock.Lock()
    46  	defer w.writerLock.Unlock()
    47  
    48  	// Nil out the writer to stop any future writes to it.
    49  	w.writer = nil
    50  }