From 9ddf3e147c3a0da50475a8309978a47e0c9a7232 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Wed, 20 Jun 2018 08:24:17 +0300
Subject: [PATCH] Use bytes.Buffer pools for wrappers

---
 wrappers/buffer_pool.go     | 27 +++++++++++++++++++++++++++
 wrappers/streamcipherrwc.go | 24 ++++++++++++------------
 2 files changed, 39 insertions(+), 12 deletions(-)
 create mode 100644 wrappers/buffer_pool.go

diff --git a/wrappers/buffer_pool.go b/wrappers/buffer_pool.go
new file mode 100644
index 0000000..ead700d
--- /dev/null
+++ b/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{}
+		},
+	}
+}
diff --git a/wrappers/streamcipherrwc.go b/wrappers/streamcipherrwc.go
index 77243b7..1d7d73c 100644
--- a/wrappers/streamcipherrwc.go
+++ b/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.