From 79826063710171012c64be90ef074fcb50486ac0 Mon Sep 17 00:00:00 2001
From: Brian Goff <cpuguy83@gmail.com>
Date: Fri, 18 Dec 2015 16:06:28 -0500
Subject: [PATCH] Cleanup WriteFlusher

Signed-off-by: Brian Goff <cpuguy83@gmail.com>
---
 ioutils/writeflusher.go | 82 ++++++++++++++++++++---------------------
 1 file changed, 41 insertions(+), 41 deletions(-)

diff --git a/ioutils/writeflusher.go b/ioutils/writeflusher.go
index 2b35a26..52a4901 100644
--- a/ioutils/writeflusher.go
+++ b/ioutils/writeflusher.go
@@ -1,9 +1,7 @@
 package ioutils
 
 import (
-	"errors"
 	"io"
-	"net/http"
 	"sync"
 )
 
@@ -11,45 +9,43 @@ import (
 // is a flush. In addition, the Close method can be called to intercept
 // Read/Write calls if the targets lifecycle has already ended.
 type WriteFlusher struct {
-	mu      sync.Mutex
-	w       io.Writer
-	flusher http.Flusher
-	flushed bool
-	closed  error
-
-	// TODO(stevvooe): Use channel for closed instead, remove mutex. Using a
-	// channel will allow one to properly order the operations.
+	w           io.Writer
+	flusher     flusher
+	flushed     chan struct{}
+	flushedOnce sync.Once
+	closed      chan struct{}
+	closeLock   sync.Mutex
 }
 
-var errWriteFlusherClosed = errors.New("writeflusher: closed")
+type flusher interface {
+	Flush()
+}
+
+var errWriteFlusherClosed = io.EOF
 
 func (wf *WriteFlusher) Write(b []byte) (n int, err error) {
-	wf.mu.Lock()
-	defer wf.mu.Unlock()
-	if wf.closed != nil {
-		return 0, wf.closed
+	select {
+	case <-wf.closed:
+		return 0, errWriteFlusherClosed
+	default:
 	}
 
 	n, err = wf.w.Write(b)
-	wf.flush() // every write is a flush.
+	wf.Flush() // every write is a flush.
 	return n, err
 }
 
 // Flush the stream immediately.
 func (wf *WriteFlusher) Flush() {
-	wf.mu.Lock()
-	defer wf.mu.Unlock()
-
-	wf.flush()
-}
-
-// flush the stream immediately without taking a lock. Used internally.
-func (wf *WriteFlusher) flush() {
-	if wf.closed != nil {
+	select {
+	case <-wf.closed:
 		return
+	default:
 	}
 
-	wf.flushed = true
+	wf.flushedOnce.Do(func() {
+		close(wf.flushed)
+	})
 	wf.flusher.Flush()
 }
 
@@ -59,34 +55,38 @@ func (wf *WriteFlusher) Flushed() bool {
 	// BUG(stevvooe): Remove this method. Its use is inherently racy. Seems to
 	// be used to detect whether or a response code has been issued or not.
 	// Another hook should be used instead.
-	wf.mu.Lock()
-	defer wf.mu.Unlock()
-
-	return wf.flushed
+	var flushed bool
+	select {
+	case <-wf.flushed:
+		flushed = true
+	default:
+	}
+	return flushed
 }
 
 // Close closes the write flusher, disallowing any further writes to the
 // target. After the flusher is closed, all calls to write or flush will
 // result in an error.
 func (wf *WriteFlusher) Close() error {
-	wf.mu.Lock()
-	defer wf.mu.Unlock()
+	wf.closeLock.Lock()
+	defer wf.closeLock.Unlock()
 
-	if wf.closed != nil {
-		return wf.closed
+	select {
+	case <-wf.closed:
+		return errWriteFlusherClosed
+	default:
+		close(wf.closed)
 	}
-
-	wf.closed = errWriteFlusherClosed
 	return nil
 }
 
 // NewWriteFlusher returns a new WriteFlusher.
 func NewWriteFlusher(w io.Writer) *WriteFlusher {
-	var flusher http.Flusher
-	if f, ok := w.(http.Flusher); ok {
-		flusher = f
+	var fl flusher
+	if f, ok := w.(flusher); ok {
+		fl = f
 	} else {
-		flusher = &NopFlusher{}
+		fl = &NopFlusher{}
 	}
-	return &WriteFlusher{w: w, flusher: flusher}
+	return &WriteFlusher{w: w, flusher: fl, closed: make(chan struct{}), flushed: make(chan struct{})}
 }