TBR: gosh: fix buffered_pipe's WriteTo implementation
The initial implementation failed to wait for all but the first write to
the pipe -- as soon as a Write happened, WriteTo would return. Instead,
WriteTo ought to wait for the pipe to be Closed, and write the data from
all the Writes to the pipe that happen up until the Close.
Change-Id: I1e09261084217df5e1fd82c4309ad479bbb423f6
diff --git a/gosh/buffered_pipe.go b/gosh/buffered_pipe.go
index 6bc9d30..87c36cf 100644
--- a/gosh/buffered_pipe.go
+++ b/gosh/buffered_pipe.go
@@ -47,16 +47,19 @@
// WriteTo implements the io.WriterTo method; it is the fast version of Read
// used by io.Copy.
+// Unlike Read, which returns io.EOF to signal that all data has been read,
+// WriteTo blocks until all data has been written to w, and never returns
+// io.EOF.
func (p *bufferedPipe) WriteTo(w io.Writer) (int64, error) {
p.cond.L.Lock()
defer p.cond.L.Unlock()
+ var written int64
for {
- // Read any remaining data before checking whether the pipe is closed.
- if p.buf.Len() > 0 {
- return p.buf.WriteTo(w)
- }
- if p.closed {
- return 0, io.EOF
+ // Keep writing data until the pipe is closed.
+ n, err := p.buf.WriteTo(w)
+ written += n
+ if p.closed || err != nil {
+ return written, err
}
p.cond.Wait()
}
diff --git a/gosh/buffered_pipe_test.go b/gosh/buffered_pipe_test.go
index a4fef76..b3d45b9 100644
--- a/gosh/buffered_pipe_test.go
+++ b/gosh/buffered_pipe_test.go
@@ -40,22 +40,61 @@
func TestReadFromWriteTo(t *testing.T) {
p, buf := newBufferedPipe(), new(bytes.Buffer)
if n, err := p.(io.ReaderFrom).ReadFrom(strings.NewReader("foobarbaz")); n != 9 || err != nil {
- t.Errorf("readfrom got (%v,%v) want (9,nil)", n, err)
+ t.Errorf("ReadFrom got (%v, %v), want (9, <nil>)", n, err)
}
- if n, err := p.(io.WriterTo).WriteTo(buf); n != 9 || err != nil {
- t.Errorf("writeto got (%v,%v) want (9,nil)", n, err)
- }
- if got, want := buf.String(), "foobarbaz"; got != want {
- t.Errorf("writeto got %v want %v", got, want)
- }
- buf.Reset()
+ nCh, errCh := make(chan int64, 1), make(chan error, 1)
+ go func() {
+ n, err := p.(io.WriterTo).WriteTo(buf)
+ nCh <- n
+ errCh <- err
+ }()
if n, err := p.(io.ReaderFrom).ReadFrom(strings.NewReader("foobarbaz")); n != 9 || err != nil {
- t.Errorf("readfrom got (%v,%v) want (9,nil)", n, err)
+ t.Errorf("ReadFrom got (%v, %v), want (9, <nil>)", n, err)
}
if err := p.Close(); err != nil {
- t.Errorf("close failed: %v", err)
+ t.Errorf("Close failed: %v", err)
}
- if n, err := p.(io.WriterTo).WriteTo(buf); n != 9 || err != nil {
- t.Errorf("writeto got (%v,%v) want (9,nil)", n, err)
+ if n, err := <-nCh, <-errCh; n != 18 || err != nil {
+ t.Errorf("WriteTo got (%v, %v), want (18, <nil>)", n, err)
+ }
+ if got, want := buf.String(), "foobarbazfoobarbaz"; got != want {
+ t.Errorf("WriteTo got %v, want %v", got, want)
+ }
+}
+
+func TestWriteToMany(t *testing.T) {
+ p := newBufferedPipe()
+ pR, pW := io.Pipe()
+ nCh, errCh := make(chan int64, 1), make(chan error, 1)
+ go func() {
+ n, err := p.(io.WriterTo).WriteTo(pW)
+ nCh <- n
+ errCh <- err
+ }()
+ var nTotal int64
+ for _, m := range []string{
+ "mary had",
+ "a little lamb",
+ "three helpings of corn",
+ "two baked potatoes",
+ "and extra bread",
+ } {
+ if n, err := p.Write([]byte(m)); n != len(m) || err != nil {
+ t.Errorf("Write(%v) got (%v, %v), want (%v, <nil>)", m, n, err, len(m))
+ }
+
+ nTotal += int64(len(m))
+ for i := 0; i < len(m); i++ {
+ b := make([]byte, 1)
+ if n, err := pR.Read(b); n != 1 || err != nil || b[0] != m[i] {
+ t.Errorf("Read got (%v, %v, %v), want (1, <nil>, %v)", n, err, b[0], m[i])
+ }
+ }
+ }
+ if err := p.Close(); err != nil {
+ t.Errorf("Close failed: %v", err)
+ }
+ if n, err := <-nCh, <-errCh; n != nTotal || err != nil {
+ t.Errorf("WriteTo got (%v, %v), want (%v, <nil>)", n, err, nTotal)
}
}
