Move types from progressreader and broadcastwriter to broadcaster

progressreader.Broadcaster becomes broadcaster.Buffered and
broadcastwriter.Writer becomes broadcaster.Unbuffered.

The package broadcastwriter is thus renamed to broadcaster.

Signed-off-by: Tibor Vass <tibor@docker.com>

broadcaster/buffered.go

@@ -0,0 +1,167 @@
+package broadcaster
+
+import (
+	"errors"
+	"io"
+	"sync"
+)
+
+// Buffered keeps track of one or more observers watching the progress
+// of an operation. For example, if multiple clients are trying to pull an
+// image, they share a Buffered struct for the download operation.
+type Buffered struct {
+	sync.Mutex
+	// c is a channel that observers block on, waiting for the operation
+	// to finish.
+	c chan struct{}
+	// cond is a condition variable used to wake up observers when there's
+	// new data available.
+	cond *sync.Cond
+	// history is a buffer of the progress output so far, so a new observer
+	// can catch up. The history is stored as a slice of separate byte
+	// slices, so that if the writer is a WriteFlusher, the flushes will
+	// happen in the right places.
+	history [][]byte
+	// wg is a WaitGroup used to wait for all writes to finish on Close
+	wg sync.WaitGroup
+	// result is the argument passed to the first call of Close, and
+	// returned to callers of Wait
+	result error
+}
+
+// NewBuffered returns an initialized Buffered structure.
+func NewBuffered() *Buffered {
+	b := &Buffered{
+		c: make(chan struct{}),
+	}
+	b.cond = sync.NewCond(b)
+	return b
+}
+
+// closed returns true if and only if the broadcaster has been closed
+func (broadcaster *Buffered) closed() bool {
+	select {
+	case <-broadcaster.c:
+		return true
+	default:
+		return false
+	}
+}
+
+// receiveWrites runs as a goroutine so that writes don't block the Write
+// function. It writes the new data in broadcaster.history each time there's
+// activity on the broadcaster.cond condition variable.
+func (broadcaster *Buffered) receiveWrites(observer io.Writer) {
+	n := 0
+
+	broadcaster.Lock()
+
+	// The condition variable wait is at the end of this loop, so that the
+	// first iteration will write the history so far.
+	for {
+		newData := broadcaster.history[n:]
+		// Make a copy of newData so we can release the lock
+		sendData := make([][]byte, len(newData), len(newData))
+		copy(sendData, newData)
+		broadcaster.Unlock()
+
+		for len(sendData) > 0 {
+			_, err := observer.Write(sendData[0])
+			if err != nil {
+				broadcaster.wg.Done()
+				return
+			}
+			n++
+			sendData = sendData[1:]
+		}
+
+		broadcaster.Lock()
+
+		// If we are behind, we need to catch up instead of waiting
+		// or handling a closure.
+		if len(broadcaster.history) != n {
+			continue
+		}
+
+		// detect closure of the broadcast writer
+		if broadcaster.closed() {
+			broadcaster.Unlock()
+			broadcaster.wg.Done()
+			return
+		}
+
+		broadcaster.cond.Wait()
+
+		// Mutex is still locked as the loop continues
+	}
+}
+
+// Write adds data to the history buffer, and also writes it to all current
+// observers.
+func (broadcaster *Buffered) Write(p []byte) (n int, err error) {
+	broadcaster.Lock()
+	defer broadcaster.Unlock()
+
+	// Is the broadcaster closed? If so, the write should fail.
+	if broadcaster.closed() {
+		return 0, errors.New("attempted write to a closed broadcaster.Buffered")
+	}
+
+	// Add message in p to the history slice
+	newEntry := make([]byte, len(p), len(p))
+	copy(newEntry, p)
+	broadcaster.history = append(broadcaster.history, newEntry)
+
+	broadcaster.cond.Broadcast()
+
+	return len(p), nil
+}
+
+// Add adds an observer to the broadcaster. The new observer receives the
+// data from the history buffer, and also all subsequent data.
+func (broadcaster *Buffered) Add(w io.Writer) error {
+	// The lock is acquired here so that Add can't race with Close
+	broadcaster.Lock()
+	defer broadcaster.Unlock()
+
+	if broadcaster.closed() {
+		return errors.New("attempted to add observer to a closed broadcaster.Buffered")
+	}
+
+	broadcaster.wg.Add(1)
+	go broadcaster.receiveWrites(w)
+
+	return nil
+}
+
+// CloseWithError signals to all observers that the operation has finished. Its
+// argument is a result that should be returned to waiters blocking on Wait.
+func (broadcaster *Buffered) CloseWithError(result error) {
+	broadcaster.Lock()
+	if broadcaster.closed() {
+		broadcaster.Unlock()
+		return
+	}
+	broadcaster.result = result
+	close(broadcaster.c)
+	broadcaster.cond.Broadcast()
+	broadcaster.Unlock()
+
+	// Don't return until all writers have caught up.
+	broadcaster.wg.Wait()
+}
+
+// Close signals to all observers that the operation has finished. It causes
+// all calls to Wait to return nil.
+func (broadcaster *Buffered) Close() {
+	broadcaster.CloseWithError(nil)
+}
+
+// Wait blocks until the operation is marked as completed by the Close method,
+// and all writer goroutines have completed. It returns the argument that was
+// passed to Close.
+func (broadcaster *Buffered) Wait() error {
+	<-broadcaster.c
+	broadcaster.wg.Wait()
+	return broadcaster.result
+}

broadcaster/unbuffered.go

@@ -0,0 +1,49 @@
+package broadcaster
+
+import (
+	"io"
+	"sync"
+)
+
+// Unbuffered accumulates multiple io.WriteCloser by stream.
+type Unbuffered struct {
+	mu      sync.Mutex
+	writers []io.WriteCloser
+}
+
+// Add adds new io.WriteCloser.
+func (w *Unbuffered) Add(writer io.WriteCloser) {
+	w.mu.Lock()
+	w.writers = append(w.writers, writer)
+	w.mu.Unlock()
+}
+
+// Write writes bytes to all writers. Failed writers will be evicted during
+// this call.
+func (w *Unbuffered) Write(p []byte) (n int, err error) {
+	w.mu.Lock()
+	var evict []int
+	for i, sw := range w.writers {
+		if n, err := sw.Write(p); err != nil || n != len(p) {
+			// On error, evict the writer
+			evict = append(evict, i)
+		}
+	}
+	for n, i := range evict {
+		w.writers = append(w.writers[:i-n], w.writers[i-n+1:]...)
+	}
+	w.mu.Unlock()
+	return len(p), nil
+}
+
+// Clean closes and removes all writers. Last non-eol-terminated part of data
+// will be saved.
+func (w *Unbuffered) Clean() error {
+	w.mu.Lock()
+	for _, sw := range w.writers {
+		sw.Close()
+	}
+	w.writers = nil
+	w.mu.Unlock()
+	return nil
+}

broadcaster/unbuffered_test.go

@@ -0,0 +1,162 @@
+package broadcaster
+
+import (
+	"bytes"
+	"errors"
+	"strings"
+
+	"testing"
+)
+
+type dummyWriter struct {
+	buffer      bytes.Buffer
+	failOnWrite bool
+}
+
+func (dw *dummyWriter) Write(p []byte) (n int, err error) {
+	if dw.failOnWrite {
+		return 0, errors.New("Fake fail")
+	}
+	return dw.buffer.Write(p)
+}
+
+func (dw *dummyWriter) String() string {
+	return dw.buffer.String()
+}
+
+func (dw *dummyWriter) Close() error {
+	return nil
+}
+
+func TestUnbuffered(t *testing.T) {
+	writer := new(Unbuffered)
+
+	// Test 1: Both bufferA and bufferB should contain "foo"
+	bufferA := &dummyWriter{}
+	writer.Add(bufferA)
+	bufferB := &dummyWriter{}
+	writer.Add(bufferB)
+	writer.Write([]byte("foo"))
+
+	if bufferA.String() != "foo" {
+		t.Errorf("Buffer contains %v", bufferA.String())
+	}
+
+	if bufferB.String() != "foo" {
+		t.Errorf("Buffer contains %v", bufferB.String())
+	}
+
+	// Test2: bufferA and bufferB should contain "foobar",
+	// while bufferC should only contain "bar"
+	bufferC := &dummyWriter{}
+	writer.Add(bufferC)
+	writer.Write([]byte("bar"))
+
+	if bufferA.String() != "foobar" {
+		t.Errorf("Buffer contains %v", bufferA.String())
+	}
+
+	if bufferB.String() != "foobar" {
+		t.Errorf("Buffer contains %v", bufferB.String())
+	}
+
+	if bufferC.String() != "bar" {
+		t.Errorf("Buffer contains %v", bufferC.String())
+	}
+
+	// Test3: Test eviction on failure
+	bufferA.failOnWrite = true
+	writer.Write([]byte("fail"))
+	if bufferA.String() != "foobar" {
+		t.Errorf("Buffer contains %v", bufferA.String())
+	}
+	if bufferC.String() != "barfail" {
+		t.Errorf("Buffer contains %v", bufferC.String())
+	}
+	// Even though we reset the flag, no more writes should go in there
+	bufferA.failOnWrite = false
+	writer.Write([]byte("test"))
+	if bufferA.String() != "foobar" {
+		t.Errorf("Buffer contains %v", bufferA.String())
+	}
+	if bufferC.String() != "barfailtest" {
+		t.Errorf("Buffer contains %v", bufferC.String())
+	}
+
+	// Test4: Test eviction on multiple simultaneous failures
+	bufferB.failOnWrite = true
+	bufferC.failOnWrite = true
+	bufferD := &dummyWriter{}
+	writer.Add(bufferD)
+	writer.Write([]byte("yo"))
+	writer.Write([]byte("ink"))
+	if strings.Contains(bufferB.String(), "yoink") {
+		t.Errorf("bufferB received write. contents: %q", bufferB)
+	}
+	if strings.Contains(bufferC.String(), "yoink") {
+		t.Errorf("bufferC received write. contents: %q", bufferC)
+	}
+	if g, w := bufferD.String(), "yoink"; g != w {
+		t.Errorf("bufferD = %q, want %q", g, w)
+	}
+
+	writer.Clean()
+}
+
+type devNullCloser int
+
+func (d devNullCloser) Close() error {
+	return nil
+}
+
+func (d devNullCloser) Write(buf []byte) (int, error) {
+	return len(buf), nil
+}
+
+// This test checks for races. It is only useful when run with the race detector.
+func TestRaceUnbuffered(t *testing.T) {
+	writer := new(Unbuffered)
+	c := make(chan bool)
+	go func() {
+		writer.Add(devNullCloser(0))
+		c <- true
+	}()
+	writer.Write([]byte("hello"))
+	<-c
+}
+
+func BenchmarkUnbuffered(b *testing.B) {
+	writer := new(Unbuffered)
+	setUpWriter := func() {
+		for i := 0; i < 100; i++ {
+			writer.Add(devNullCloser(0))
+			writer.Add(devNullCloser(0))
+			writer.Add(devNullCloser(0))
+		}
+	}
+	testLine := "Line that thinks that it is log line from docker"
+	var buf bytes.Buffer
+	for i := 0; i < 100; i++ {
+		buf.Write([]byte(testLine + "\n"))
+	}
+	// line without eol
+	buf.Write([]byte(testLine))
+	testText := buf.Bytes()
+	b.SetBytes(int64(5 * len(testText)))
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		setUpWriter()
+		b.StartTimer()
+
+		for j := 0; j < 5; j++ {
+			if _, err := writer.Write(testText); err != nil {
+				b.Fatal(err)
+			}
+		}
+
+		b.StopTimer()
+		writer.Clean()
+		b.StartTimer()
+	}
+}