From 2c938d284894a257a26e0800e42286e876cb975e Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Wed, 20 Jun 2018 09:33:57 +0300
Subject: [PATCH] Reuse buffers on socket pumping

---
 proxy/copy_pool.go | 16 ++++++++++++++++
 proxy/server.go    | 21 +++++++++++++--------
 2 files changed, 29 insertions(+), 8 deletions(-)
 create mode 100644 proxy/copy_pool.go

diff --git a/proxy/copy_pool.go b/proxy/copy_pool.go
new file mode 100644
index 0000000..b6eb784
--- /dev/null
+++ b/proxy/copy_pool.go
@@ -0,0 +1,16 @@
+package proxy
+
+import "sync"
+
+const copyBufferSize = 30 * 1024
+
+var copyPool sync.Pool
+
+func init() {
+	copyPool = sync.Pool{
+		New: func() interface{} {
+			data := make([]byte, copyBufferSize)
+			return &data
+		},
+	}
+}
diff --git a/proxy/server.go b/proxy/server.go
index bd5564d..4e2d335 100644
--- a/proxy/server.go
+++ b/proxy/server.go
@@ -83,14 +83,10 @@ func (s *Server) accept(conn net.Conn) {
 
 	wait := &sync.WaitGroup{}
 	wait.Add(2)
-	go func() {
-		defer wait.Done()
-		io.Copy(clientConn, tgConn) // nolint: errcheck
-	}()
-	go func() {
-		defer wait.Done()
-		io.Copy(tgConn, clientConn) // nolint: errcheck
-	}()
+
+	go s.pipe(clientConn, tgConn, wait)
+	go s.pipe(tgConn, clientConn, wait)
+
 	<-ctx.Done()
 	wait.Wait()
 
@@ -131,6 +127,15 @@ func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFun
 	return conn, nil
 }
 
+func (s *Server) pipe(dst io.Writer, src io.Reader, wait *sync.WaitGroup) {
+	defer wait.Done()
+
+	buf := copyPool.Get().(*[]byte)
+	defer copyPool.Put(buf)
+
+	io.CopyBuffer(dst, src, *buf) // nolint: errcheck
+}
+
 // NewServer creates new instance of MTPROTO proxy.
 func NewServer(conf *config.Config, logger *zap.SugaredLogger, stat *Stats) *Server {
 	return &Server{