github.com/bhojpur/cache@v0.0.4/pkg/ioutils/writeflusher.go (about)

     1  package ioutils
     2  
     3  // Copyright (c) 2018 Bhojpur Consulting Private Limited, India. All rights reserved.
     4  
     5  // Permission is hereby granted, free of charge, to any person obtaining a copy
     6  // of this software and associated documentation files (the "Software"), to deal
     7  // in the Software without restriction, including without limitation the rights
     8  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     9  // copies of the Software, and to permit persons to whom the Software is
    10  // furnished to do so, subject to the following conditions:
    11  
    12  // The above copyright notice and this permission notice shall be included in
    13  // all copies or substantial portions of the Software.
    14  
    15  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    16  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    17  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    18  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    19  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    20  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    21  // THE SOFTWARE.
    22  
    23  import (
    24  	"io"
    25  	"sync"
    26  )
    27  
    28  // WriteFlusher wraps the Write and Flush operation ensuring that every write
    29  // is a flush. In addition, the Close method can be called to intercept
    30  // Read/Write calls if the targets lifecycle has already ended.
    31  type WriteFlusher struct {
    32  	w           io.Writer
    33  	flusher     flusher
    34  	flushed     chan struct{}
    35  	flushedOnce sync.Once
    36  	closed      chan struct{}
    37  	closeLock   sync.Mutex
    38  }
    39  
    40  type flusher interface {
    41  	Flush()
    42  }
    43  
    44  var errWriteFlusherClosed = io.EOF
    45  
    46  func (wf *WriteFlusher) Write(b []byte) (n int, err error) {
    47  	select {
    48  	case <-wf.closed:
    49  		return 0, errWriteFlusherClosed
    50  	default:
    51  	}
    52  
    53  	n, err = wf.w.Write(b)
    54  	wf.Flush() // every write is a flush.
    55  	return n, err
    56  }
    57  
    58  // Flush the stream immediately.
    59  func (wf *WriteFlusher) Flush() {
    60  	select {
    61  	case <-wf.closed:
    62  		return
    63  	default:
    64  	}
    65  
    66  	wf.flushedOnce.Do(func() {
    67  		close(wf.flushed)
    68  	})
    69  	wf.flusher.Flush()
    70  }
    71  
    72  // Flushed returns the state of flushed.
    73  // If it's flushed, return true, or else it return false.
    74  func (wf *WriteFlusher) Flushed() bool {
    75  	// BUG(stevvooe): Remove this method. Its use is inherently racy. Seems to
    76  	// be used to detect whether or a response code has been issued or not.
    77  	// Another hook should be used instead.
    78  	var flushed bool
    79  	select {
    80  	case <-wf.flushed:
    81  		flushed = true
    82  	default:
    83  	}
    84  	return flushed
    85  }
    86  
    87  // Close closes the write flusher, disallowing any further writes to the
    88  // target. After the flusher is closed, all calls to write or flush will
    89  // result in an error.
    90  func (wf *WriteFlusher) Close() error {
    91  	wf.closeLock.Lock()
    92  	defer wf.closeLock.Unlock()
    93  
    94  	select {
    95  	case <-wf.closed:
    96  		return errWriteFlusherClosed
    97  	default:
    98  		close(wf.closed)
    99  	}
   100  	return nil
   101  }
   102  
   103  // NewWriteFlusher returns a new WriteFlusher.
   104  func NewWriteFlusher(w io.Writer) *WriteFlusher {
   105  	var fl flusher
   106  	if f, ok := w.(flusher); ok {
   107  		fl = f
   108  	} else {
   109  		fl = &NopFlusher{}
   110  	}
   111  	return &WriteFlusher{w: w, flusher: fl, closed: make(chan struct{}), flushed: make(chan struct{})}
   112  }