Use bytes.Buffer pools for wrappers

wrappers/buffer_pool.go

@@ -0,0 +1,27 @@
+package wrappers
+
+import (
+	"bytes"
+	"sync"
+)
+
+var bufPool sync.Pool
+
+func getBuffer() *bytes.Buffer {
+	buf := bufPool.Get().(*bytes.Buffer)
+	buf.Reset()
+
+	return buf
+}
+
+func putBuffer(buf *bytes.Buffer) {
+	bufPool.Put(buf)
+}
+
+func init() {
+	bufPool = sync.Pool{
+		New: func() interface{} {
+			return &bytes.Buffer{}
+		},
+	}
+}

wrappers/streamcipherrwc.go

@@ -22,20 +22,20 @@ func (c *StreamCipherReadWriteCloser) Read(p []byte) (n int, err error) {
 
 // Write writes into connection.
 func (c *StreamCipherReadWriteCloser) Write(p []byte) (int, error) {
-	encrypted := make([]byte, len(p))
+	// This is to decrease an amount of allocations. Unfortunately, escape
+	// analysis in (at least Golang 1.10) is absolutely not perfect. For
+	// example, it understands that we want to have a slice locally, right?
+	// But since slice is effectively 2 ints + uintptr to [number]byte, the
+	// most heavyweight part is placed in heap.
+	buf := getBuffer()
+	defer putBuffer(buf)
+	buf.Grow(len(p))
+	buf.Write(p)
+
+	encrypted := buf.Bytes()
 	c.encryptor.XORKeyStream(encrypted, p)
-	allWritten := 0
-
-	for len(encrypted) > 0 {
-		n, err := c.conn.Write(encrypted)
-		allWritten += n
-		if err != nil {
-			return allWritten, err
-		}
-		encrypted = encrypted[n:]
-	}
 
-	return allWritten, nil
+	return c.conn.Write(encrypted)
 }
 
 // Close closes underlying connection.