pathway/broadcast.go

package pathway

import (
	"io"
)

type broadcastWriter struct {
	writers []io.Writer
}

func (t *broadcastWriter) Close() (err error) {
	for i := 0; i < len(t.writers); i++ {
		if w, ok := t.writers[i].(io.WriteCloser); ok {
			if e := w.Close(); e != nil {
				err = e
			}
		}
	}
	return
}

func (t *broadcastWriter) Write(p []byte) (n int, err error) {
	for i := 0; i < len(t.writers); i++ {
		n, err = t.writers[i].Write(p)

		if err == nil && n != len(p) {
			err = io.ErrShortWrite
		}

		if err != nil {
			// if not at the end, move last to here and redo this
			if i < len(t.writers)-1 {
				// close the writer if it implements WriteCloser
				if w, ok := t.writers[i].(io.WriteCloser); ok {
					defer w.Close()
				}
				t.writers[i] = t.writers[len(t.writers)-1]
				i = i - 1
			}
			// always shorten list by one to drop the last one
			t.writers = t.writers[:len(t.writers)-1]
		}
	}
	if len(t.writers) <= 0 {
		return 0, err
	}

	return len(p), nil
}

// BroadcastWriter creates a writer that duplicates its writes to all the
// provided writers, similar to the Unix tee(1) command.
//
// Each write is written to each listed writer, one at a time.
// If a listed writer returns an error, that overall write operation
// continues and the offending writer is dropped from the list.
//
// Only if all writers are dropped, the last error is returned at the end.
func BroadcastWriter(writers ...io.Writer) io.WriteCloser {
	allWriters := make([]io.Writer, 0, len(writers))
	for _, w := range writers {
		if bw, ok := w.(*broadcastWriter); ok {
			allWriters = append(allWriters, bw.writers...)
		} else {
			allWriters = append(allWriters, w)
		}
	}
	return &broadcastWriter{allWriters}
}