github.com/ttys3/engine@v17.12.1-ce-rc2+incompatible/pkg/ioutils/writeflusher.go (about)

     1  package ioutils
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  )
     7  
     8  // WriteFlusher wraps the Write and Flush operation ensuring that every write
     9  // is a flush. In addition, the Close method can be called to intercept
    10  // Read/Write calls if the targets lifecycle has already ended.
    11  type WriteFlusher struct {
    12  	w           io.Writer
    13  	flusher     flusher
    14  	flushed     chan struct{}
    15  	flushedOnce sync.Once
    16  	closed      chan struct{}
    17  	closeLock   sync.Mutex
    18  }
    19  
    20  type flusher interface {
    21  	Flush()
    22  }
    23  
    24  var errWriteFlusherClosed = io.EOF
    25  
    26  func (wf *WriteFlusher) Write(b []byte) (n int, err error) {
    27  	select {
    28  	case <-wf.closed:
    29  		return 0, errWriteFlusherClosed
    30  	default:
    31  	}
    32  
    33  	n, err = wf.w.Write(b)
    34  	wf.Flush() // every write is a flush.
    35  	return n, err
    36  }
    37  
    38  // Flush the stream immediately.
    39  func (wf *WriteFlusher) Flush() {
    40  	select {
    41  	case <-wf.closed:
    42  		return
    43  	default:
    44  	}
    45  
    46  	wf.flushedOnce.Do(func() {
    47  		close(wf.flushed)
    48  	})
    49  	wf.flusher.Flush()
    50  }
    51  
    52  // Flushed returns the state of flushed.
    53  // If it's flushed, return true, or else it return false.
    54  func (wf *WriteFlusher) Flushed() bool {
    55  	// BUG(stevvooe): Remove this method. Its use is inherently racy. Seems to
    56  	// be used to detect whether or a response code has been issued or not.
    57  	// Another hook should be used instead.
    58  	var flushed bool
    59  	select {
    60  	case <-wf.flushed:
    61  		flushed = true
    62  	default:
    63  	}
    64  	return flushed
    65  }
    66  
    67  // Close closes the write flusher, disallowing any further writes to the
    68  // target. After the flusher is closed, all calls to write or flush will
    69  // result in an error.
    70  func (wf *WriteFlusher) Close() error {
    71  	wf.closeLock.Lock()
    72  	defer wf.closeLock.Unlock()
    73  
    74  	select {
    75  	case <-wf.closed:
    76  		return errWriteFlusherClosed
    77  	default:
    78  		close(wf.closed)
    79  	}
    80  	return nil
    81  }
    82  
    83  // NewWriteFlusher returns a new WriteFlusher.
    84  func NewWriteFlusher(w io.Writer) *WriteFlusher {
    85  	var fl flusher
    86  	if f, ok := w.(flusher); ok {
    87  		fl = f
    88  	} else {
    89  		fl = &NopFlusher{}
    90  	}
    91  	return &WriteFlusher{w: w, flusher: fl, closed: make(chan struct{}), flushed: make(chan struct{})}
    92  }