diff --git a/buffer/pool.go b/buffer/pool.go index 42d9886..0655b4b 100644 --- a/buffer/pool.go +++ b/buffer/pool.go @@ -212,3 +212,57 @@ func (b *Buffer) BuildBytes(reuse ...[]byte) []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 d4c009d..6e78466 100644 --- a/jwriter/writer.go +++ b/jwriter/writer.go @@ -48,6 +48,16 @@ func (w *Writer) BuildBytes(reuse ...[]byte) ([]byte, error) { return w.Buffer.BuildBytes(reuse...), 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)