From ce0e78724b3ddd3b7902348231901648faa30ea0 Mon Sep 17 00:00:00 2001
From: Erik Dubbelboer <erik@dubbelboer.com>
Date: Tue, 29 Nov 2016 08:34:23 +0000
Subject: [PATCH] Add ReadCloser method to jwriter and buffer

When marshalling something as body for a http.Request we can only use
BuildBytes() which always allocates a buffer for our whole body. This
pull request adds a ReadCloser() method which returns an io.ReadCloser
that can be passed as body into http.NewRequest. This ReadCloser will
read from the existing buffs and putBuf them when they are not needed
anymore.
---
 buffer/pool.go      | 54 +++++++++++++++++++++++++++++++++++++++++++++
 buffer/pool_test.go | 28 +++++++++++++++++++++++
 jwriter/writer.go   | 10 +++++++++
 3 files changed, 92 insertions(+)

diff --git a/buffer/pool.go b/buffer/pool.go
index 4de4a51..f1024e1 100644
--- a/buffer/pool.go
+++ b/buffer/pool.go
@@ -205,3 +205,57 @@ func (b *Buffer) BuildBytes() []byte {
 
 	return ret
 }
+
+type readCloser struct {
+	offset int
+	bufs   [][]byte
+}
+
+func (r *readCloser) Read(p []byte) (n int, err error) {
+	for _, buf := range r.bufs {
+		// Copy as much as we can.
+		x := copy(p[n:], buf[r.offset:])
+		n += x // Increment how much we filled.
+
+		// Did we empty the whole buffer?
+		if r.offset+x == len(buf) {
+			// On to the next buffer.
+			r.offset = 0
+			r.bufs = r.bufs[1:]
+
+			// We can release this buffer.
+			putBuf(buf)
+		} else {
+			r.offset += x
+		}
+
+		if n == len(p) {
+			break
+		}
+	}
+	// No buffers left or nothing read?
+	if len(r.bufs) == 0 {
+		err = io.EOF
+	}
+	return
+}
+
+func (r *readCloser) Close() error {
+	// Release all remainig buffers.
+	for _, buf := range r.bufs {
+		putBuf(buf)
+	}
+
+	return nil
+}
+
+// ReadCloser creates an io.ReadCloser with all the contents of the buffer.
+func (b *Buffer) ReadCloser() io.ReadCloser {
+	ret := &readCloser{0, append(b.bufs, b.Buf)}
+
+	b.bufs = nil
+	b.toPool = nil
+	b.Buf = nil
+
+	return ret
+}
diff --git a/buffer/pool_test.go b/buffer/pool_test.go
index 254ca6f..680623a 100644
--- a/buffer/pool_test.go
+++ b/buffer/pool_test.go
@@ -77,3 +77,31 @@ func TestDumpTo(t *testing.T) {
 		t.Errorf("DumpTo() = %v; want %v", n, len(want))
 	}
 }
+
+func TestReadCloser(t *testing.T) {
+	var b Buffer
+	var want []byte
+
+	s := "test"
+	for i := 0; i < 1000; i++ {
+		b.AppendBytes([]byte(s))
+		want = append(want, s...)
+	}
+
+	out := &bytes.Buffer{}
+	rc := b.ReadCloser()
+	n, err := out.ReadFrom(rc)
+	if err != nil {
+		t.Errorf("ReadCloser() error: %v", err)
+	}
+	rc.Close() // Will always return nil
+
+	got := out.Bytes()
+	if !bytes.Equal(got, want) {
+		t.Errorf("DumpTo(): got %v; want %v", got, want)
+	}
+
+	if n != int64(len(want)) {
+		t.Errorf("DumpTo() = %v; want %v", n, len(want))
+	}
+}
diff --git a/jwriter/writer.go b/jwriter/writer.go
index a3ef534..9952a87 100644
--- a/jwriter/writer.go
+++ b/jwriter/writer.go
@@ -46,6 +46,16 @@ func (w *Writer) BuildBytes() ([]byte, error) {
 	return w.Buffer.BuildBytes(), nil
 }
 
+// ReadCloser returns an io.ReadCloser that can be used to read the data.
+// ReadCloser also resets the buffer.
+func (w *Writer) ReadCloser() (io.ReadCloser, error) {
+	if w.Error != nil {
+		return nil, w.Error
+	}
+
+	return w.Buffer.ReadCloser(), nil
+}
+
 // RawByte appends raw binary data to the buffer.
 func (w *Writer) RawByte(c byte) {
 	w.Buffer.AppendByte(c)