From 079a027d27a9eed18c99c0c6a6e2fc70f9dd07b7 Mon Sep 17 00:00:00 2001
From: Jorropo <jorropo.pgm@gmail.com>
Date: Mon, 7 Mar 2022 12:56:01 +0000
Subject: [PATCH] io: add WriterTo to MultiReader
MIME-Version: 1.0
Content-Type: text/plain; charset=utf8
Content-Transfer-Encoding: 8bit

This patch allows to zerocopy using MultiReader.
This is done by MultiReader implementing WriterTo.

Each sub reader is copied using usual io copy helper and thus use
WriterTo or ReadFrom with reflection.

There is a special case for when a subreader is a MultiReader.
Instead of using copyBuffer which would call multiReader.WriteTo,
multiReader.writeToWithBuffer is used instead, the difference
is that the temporary copy buffer is passed along, saving
allocations for nested MultiReaders.

The workflow looks like this:
- multiReader.WriteTo (allocates 32k buffer)
  - multiReader.writeToWithBuffer
    - for each subReader:
      - is instance of multiReader ?
        - yes, call multiReader.writeToWithBuffer
        - no, call copyBuffer(writer, currentReader, buffer)
          - does currentReader implements WriterTo ?
           - yes, use use currentReader.WriteTo
           - no, does writer implement ReadFrom ?
             - yes, use writer.ReadFrom
             - no, copy using Read / Write with buffer

This can be improved by lazy allocating the 32k buffer.
For example a MultiReader of such types:
  MultiReader(
    bytes.Reader, // WriterTo-able
    bytes.Reader, // WriterTo-able
    bytes.Reader, // WriterTo-able
  )

Doesn't need any allocation, all copy can be done using bytes.Reader's
internal data slice. However currently we still allocate a 32k buffer
for nothing.

This optimisation has been omitted for a future patch because of high
complexity costs for a non obvious performance cost (it needs a benchmark).
This patch at least is on par with the previous MultiReader.Read
workflow allocation wise.

Fixes #50842

Change-Id: Ib070c8f36337d9dd86090df8a703c5df97a773ae
GitHub-Last-Rev: 8ebe60ceacec6bd52b63d9bdc05cd5b4ada57a6e
GitHub-Pull-Request: golang/go#51502
Reviewed-on: https://go-review.googlesource.com/c/go/+/390215
Run-TryBot: Ian Lance Taylor <iant@golang.org>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@golang.org>
Reviewed-by: Daniel Martí <mvdan@mvdan.cc>
Trust: Daniel Martí <mvdan@mvdan.cc>
---
 src/io/multi.go      | 25 +++++++++++++++++++++++++
 src/io/multi_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/src/io/multi.go b/src/io/multi.go
index 24ee71e4ca..909b7e4523 100644
--- a/src/io/multi.go
+++ b/src/io/multi.go
@@ -41,6 +41,31 @@ func (mr *multiReader) Read(p []byte) (n int, err error) {
 	return 0, EOF
 }
 
+func (mr *multiReader) WriteTo(w Writer) (sum int64, err error) {
+	return mr.writeToWithBuffer(w, make([]byte, 1024 * 32))
+}
+
+func (mr *multiReader) writeToWithBuffer(w Writer, buf []byte) (sum int64, err error) {
+	for i, r := range mr.readers {
+		var n int64
+		if subMr, ok := r.(*multiReader); ok { // reuse buffer with nested multiReaders
+			n, err = subMr.writeToWithBuffer(w, buf)
+		} else {
+			n, err = copyBuffer(w, r, buf)
+		}
+		sum += n
+		if err != nil {
+			mr.readers = mr.readers[i:] // permit resume / retry after error
+			return sum, err
+		}
+		mr.readers[i] = nil // permit early GC
+	}
+	mr.readers = nil
+	return sum, nil
+}
+
+var _ WriterTo = (*multiReader)(nil)
+
 // MultiReader returns a Reader that's the logical concatenation of
 // the provided input readers. They're read sequentially. Once all
 // inputs have returned EOF, Read will return EOF.  If any of the readers
diff --git a/src/io/multi_test.go b/src/io/multi_test.go
index e877e54571..679312c23b 100644
--- a/src/io/multi_test.go
+++ b/src/io/multi_test.go
@@ -63,6 +63,31 @@ func TestMultiReader(t *testing.T) {
 	})
 }
 
+func TestMultiReaderAsWriterTo(t *testing.T) {
+	mr := MultiReader(
+		strings.NewReader("foo "),
+		MultiReader( // Tickle the buffer reusing codepath
+			strings.NewReader(""),
+			strings.NewReader("bar"),
+		),
+	)
+	mrAsWriterTo, ok := mr.(WriterTo)
+	if !ok {
+		t.Fatalf("expected cast to WriterTo to succeed")
+	}
+	sink := &strings.Builder{}
+	n, err := mrAsWriterTo.WriteTo(sink)
+	if err != nil {
+		t.Fatalf("expected no error; got %v", err)
+	}
+	if n != 7 {
+		t.Errorf("expected read 7 bytes; got %d", n)
+	}
+	if result := sink.String(); result != "foo bar" {
+		t.Errorf(`expected "foo bar"; got %q`, result)
+	}
+}
+
 func TestMultiWriter(t *testing.T) {
 	sink := new(bytes.Buffer)
 	// Hide bytes.Buffer's WriteString method:
-- 
2.50.0