Merge pull request #608 from pjbgf/optimise-zlib-reader

Optimise zlib reader and consolidate sync.Pools

6 files changed, 303 insertions, 0 deletions

diff --git a/utils/sync/bufio.go b/utils/sync/bufio.go
new file mode 100644
index 0000000..5009ea8
--- /dev/null
+++ b/utils/sync/bufio.go
@@ -0,0 +1,29 @@
+package sync
+
+import (
+	"bufio"
+	"io"
+	"sync"
+)
+
+var bufioReader = sync.Pool{
+	New: func() interface{} {
+		return bufio.NewReader(nil)
+	},
+}
+
+// GetBufioReader returns a *bufio.Reader that is managed by a sync.Pool.
+// Returns a bufio.Reader that is resetted with reader and ready for use.
+//
+// After use, the *bufio.Reader should be put back into the sync.Pool
+// by calling PutBufioReader.
+func GetBufioReader(reader io.Reader) *bufio.Reader {
+	r := bufioReader.Get().(*bufio.Reader)
+	r.Reset(reader)
+	return r
+}
+
+// PutBufioReader puts reader back into its sync.Pool.
+func PutBufioReader(reader *bufio.Reader) {
+	bufioReader.Put(reader)
+}
diff --git a/utils/sync/bufio_test.go b/utils/sync/bufio_test.go
new file mode 100644
index 0000000..e70f3d8
--- /dev/null
+++ b/utils/sync/bufio_test.go
@@ -0,0 +1,26 @@
+package sync
+
+import (
+	"io"
+	"strings"
+	"testing"
+)
+
+func TestGetAndPutBufioReader(t *testing.T) {
+	wanted := "someinput"
+	r := GetBufioReader(strings.NewReader(wanted))
+	if r == nil {
+		t.Error("nil was not expected")
+	}
+
+	got, err := r.ReadString(0)
+	if err != nil && err != io.EOF {
+		t.Errorf("unexpected error reading string: %v", err)
+	}
+
+	if wanted != got {
+		t.Errorf("wanted %q got %q", wanted, got)
+	}
+
+	PutBufioReader(r)
+}
diff --git a/utils/sync/bytes.go b/utils/sync/bytes.go
new file mode 100644
index 0000000..dd06fc0
--- /dev/null
+++ b/utils/sync/bytes.go
@@ -0,0 +1,51 @@
+package sync
+
+import (
+	"bytes"
+	"sync"
+)
+
+var (
+	byteSlice = sync.Pool{
+		New: func() interface{} {
+			b := make([]byte, 16*1024)
+			return &b
+		},
+	}
+	bytesBuffer = sync.Pool{
+		New: func() interface{} {
+			return bytes.NewBuffer(nil)
+		},
+	}
+)
+
+// GetByteSlice returns a *[]byte that is managed by a sync.Pool.
+// The initial slice length will be 16384 (16kb).
+//
+// After use, the *[]byte should be put back into the sync.Pool
+// by calling PutByteSlice.
+func GetByteSlice() *[]byte {
+	buf := byteSlice.Get().(*[]byte)
+	return buf
+}
+
+// PutByteSlice puts buf back into its sync.Pool.
+func PutByteSlice(buf *[]byte) {
+	byteSlice.Put(buf)
+}
+
+// GetBytesBuffer returns a *bytes.Buffer that is managed by a sync.Pool.
+// Returns a buffer that is resetted and ready for use.
+//
+// After use, the *bytes.Buffer should be put back into the sync.Pool
+// by calling PutBytesBuffer.
+func GetBytesBuffer() *bytes.Buffer {
+	buf := bytesBuffer.Get().(*bytes.Buffer)
+	buf.Reset()
+	return buf
+}
+
+// PutBytesBuffer puts buf back into its sync.Pool.
+func PutBytesBuffer(buf *bytes.Buffer) {
+	bytesBuffer.Put(buf)
+}
diff --git a/utils/sync/bytes_test.go b/utils/sync/bytes_test.go
new file mode 100644
index 0000000..b233429
--- /dev/null
+++ b/utils/sync/bytes_test.go
@@ -0,0 +1,49 @@
+package sync
+
+import (
+	"testing"
+)
+
+func TestGetAndPutBytesBuffer(t *testing.T) {
+	buf := GetBytesBuffer()
+	if buf == nil {
+		t.Error("nil was not expected")
+	}
+
+	initialLen := buf.Len()
+	buf.Grow(initialLen * 2)
+	grownLen := buf.Len()
+
+	PutBytesBuffer(buf)
+
+	buf = GetBytesBuffer()
+	if buf.Len() != grownLen {
+		t.Error("bytes buffer was not reused")
+	}
+
+	buf2 := GetBytesBuffer()
+	if buf2.Len() != initialLen {
+		t.Errorf("new bytes buffer length: wanted %d got %d", initialLen, buf2.Len())
+	}
+}
+
+func TestGetAndPutByteSlice(t *testing.T) {
+	slice := GetByteSlice()
+	if slice == nil {
+		t.Error("nil was not expected")
+	}
+
+	wanted := 16 * 1024
+	got := len(*slice)
+	if wanted != got {
+		t.Errorf("byte slice length: wanted %d got %d", wanted, got)
+	}
+
+	newByteSlice := make([]byte, wanted*2)
+	PutByteSlice(&newByteSlice)
+
+	newSlice := GetByteSlice()
+	if len(*newSlice) != len(newByteSlice) {
+		t.Error("byte slice was not reused")
+	}
+}
diff --git a/utils/sync/zlib.go b/utils/sync/zlib.go
new file mode 100644
index 0000000..c613885
--- /dev/null
+++ b/utils/sync/zlib.go
@@ -0,0 +1,74 @@
+package sync
+
+import (
+	"bytes"
+	"compress/zlib"
+	"io"
+	"sync"
+)
+
+var (
+	zlibInitBytes = []byte{0x78, 0x9c, 0x01, 0x00, 0x00, 0xff, 0xff, 0x00, 0x00, 0x00, 0x01}
+	zlibReader    = sync.Pool{
+		New: func() interface{} {
+			r, _ := zlib.NewReader(bytes.NewReader(zlibInitBytes))
+			return ZLibReader{
+				Reader: r.(zlibReadCloser),
+			}
+		},
+	}
+	zlibWriter = sync.Pool{
+		New: func() interface{} {
+			return zlib.NewWriter(nil)
+		},
+	}
+)
+
+type zlibReadCloser interface {
+	io.ReadCloser
+	zlib.Resetter
+}
+
+type ZLibReader struct {
+	dict   *[]byte
+	Reader zlibReadCloser
+}
+
+// GetZlibReader returns a ZLibReader that is managed by a sync.Pool.
+// Returns a ZLibReader that is resetted using a dictionary that is
+// also managed by a sync.Pool.
+//
+// After use, the ZLibReader should be put back into the sync.Pool
+// by calling PutZlibReader.
+func GetZlibReader(r io.Reader) (ZLibReader, error) {
+	z := zlibReader.Get().(ZLibReader)
+	z.dict = GetByteSlice()
+
+	err := z.Reader.Reset(r, *z.dict)
+
+	return z, err
+}
+
+// PutZlibReader puts z back into its sync.Pool, first closing the reader.
+// The Byte slice dictionary is also put back into its sync.Pool.
+func PutZlibReader(z ZLibReader) {
+	z.Reader.Close()
+	PutByteSlice(z.dict)
+	zlibReader.Put(z)
+}
+
+// GetZlibWriter returns a *zlib.Writer that is managed by a sync.Pool.
+// Returns a writer that is resetted with w and ready for use.
+//
+// After use, the *zlib.Writer should be put back into the sync.Pool
+// by calling PutZlibWriter.
+func GetZlibWriter(w io.Writer) *zlib.Writer {
+	z := zlibWriter.Get().(*zlib.Writer)
+	z.Reset(w)
+	return z
+}
+
+// PutZlibWriter puts w back into its sync.Pool.
+func PutZlibWriter(w *zlib.Writer) {
+	zlibWriter.Put(w)
+}
diff --git a/utils/sync/zlib_test.go b/utils/sync/zlib_test.go
new file mode 100644
index 0000000..b736fb2
--- /dev/null
+++ b/utils/sync/zlib_test.go
@@ -0,0 +1,74 @@
+package sync
+
+import (
+	"bytes"
+	"compress/zlib"
+	"io"
+	"testing"
+)
+
+func TestGetAndPutZlibReader(t *testing.T) {
+	_, err := GetZlibReader(bytes.NewReader(zlibInitBytes))
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	dict := &[]byte{}
+	reader := FakeZLibReader{}
+	PutZlibReader(ZLibReader{dict: dict, Reader: &reader})
+
+	if !reader.wasClosed {
+		t.Errorf("reader was not closed")
+	}
+
+	z2, err := GetZlibReader(bytes.NewReader(zlibInitBytes))
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	if dict != z2.dict {
+		t.Errorf("zlib dictionary was not reused")
+	}
+
+	if &reader != z2.Reader {
+		t.Errorf("zlib reader was not reused")
+	}
+
+	if !reader.wasReset {
+		t.Errorf("reader was not reset")
+	}
+}
+
+func TestGetAndPutZlibWriter(t *testing.T) {
+	w := GetZlibWriter(nil)
+	if w == nil {
+		t.Errorf("nil was not expected")
+	}
+
+	newW := zlib.NewWriter(nil)
+	PutZlibWriter(newW)
+
+	w2 := GetZlibWriter(nil)
+	if w2 != newW {
+		t.Errorf("zlib writer was not reused")
+	}
+}
+
+type FakeZLibReader struct {
+	wasClosed bool
+	wasReset  bool
+}
+
+func (f *FakeZLibReader) Reset(r io.Reader, dict []byte) error {
+	f.wasReset = true
+	return nil
+}
+
+func (f *FakeZLibReader) Read(p []byte) (n int, err error) {
+	return 0, nil
+}
+
+func (f *FakeZLibReader) Close() error {
+	f.wasClosed = true
+	return nil
+}