io: add WriteString support to MultiWriter

Fixes #11805

Change-Id: I081e16b869dc706bd847ee645bb902bc671c123f
Reviewed-on: https://go-review.googlesource.com/12485
Reviewed-by: Josh Bleecher Snyder <josharian@gmail.com>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

diff --git a/src/io/multi.go b/src/io/multi.go
index e26cc53e9eeb7cc8477a00247b177ffd9cea3a36..16860aa361f13ab62f6afc33857a2d308a7ac7c5 100644 (file)
--- a/src/io/multi.go
+++ b/src/io/multi.go
@@ -52,6 +52,30 @@ func (t *multiWriter) Write(p []byte) (n int, err error) {
        return len(p), nil
 }
 
+var _ stringWriter = (*multiWriter)(nil)
+
+func (t *multiWriter) WriteString(s string) (n int, err error) {
+       var p []byte // lazily initialized if/when needed
+       for _, w := range t.writers {
+               if sw, ok := w.(stringWriter); ok {
+                       n, err = sw.WriteString(s)
+               } else {
+                       if p == nil {
+                               p = []byte(s)
+                       }
+                       n, err = w.Write(p)
+               }
+               if err != nil {
+                       return
+               }
+               if n != len(s) {
+                       err = ErrShortWrite
+                       return
+               }
+       }
+       return len(s), nil
+}
+
 // MultiWriter creates a writer that duplicates its writes to all the
 // provided writers, similar to the Unix tee(1) command.
 func MultiWriter(writers ...Writer) Writer {

diff --git a/src/io/multi_test.go b/src/io/multi_test.go
index 56c6769a9eb598cc8e1671c88848eb5eb2f85245..e80592e0f505ef9ae1b7f1aee693e72ed9a3e018 100644 (file)
--- a/src/io/multi_test.go
+++ b/src/io/multi_test.go
@@ -62,8 +62,59 @@ func TestMultiReader(t *testing.T) {
 }
 
 func TestMultiWriter(t *testing.T) {
-       sha1 := sha1.New()
        sink := new(bytes.Buffer)
+       // Hide bytes.Buffer's WriteString method:
+       testMultiWriter(t, struct {
+               Writer
+               fmt.Stringer
+       }{sink, sink})
+}
+
+func TestMultiWriter_String(t *testing.T) {
+       testMultiWriter(t, new(bytes.Buffer))
+}
+
+// test that a multiWriter.WriteString calls results in at most 1 allocation,
+// even if multiple targets don't support WriteString.
+func TestMultiWriter_WriteStringSingleAlloc(t *testing.T) {
+       var sink1, sink2 bytes.Buffer
+       type simpleWriter struct { // hide bytes.Buffer's WriteString
+               Writer
+       }
+       mw := MultiWriter(simpleWriter{&sink1}, simpleWriter{&sink2})
+       allocs := int(testing.AllocsPerRun(1000, func() {
+               WriteString(mw, "foo")
+       }))
+       if allocs != 1 {
+               t.Errorf("num allocations = %d; want 1", allocs)
+       }
+}
+
+type writeStringChecker struct{ called bool }
+
+func (c *writeStringChecker) WriteString(s string) (n int, err error) {
+       c.called = true
+       return len(s), nil
+}
+
+func (c *writeStringChecker) Write(p []byte) (n int, err error) {
+       return len(p), nil
+}
+
+func TestMultiWriter_StringCheckCall(t *testing.T) {
+       var c writeStringChecker
+       mw := MultiWriter(&c)
+       WriteString(mw, "foo")
+       if !c.called {
+               t.Error("did not see WriteString call to writeStringChecker")
+       }
+}
+
+func testMultiWriter(t *testing.T, sink interface {
+       Writer
+       fmt.Stringer
+}) {
+       sha1 := sha1.New()
        mw := MultiWriter(sha1, sink)
 
        sourceString := "My input text."