From 7f386fc0421223a23c21e3c5043d974ec05487c2 Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Sun, 31 Jul 2016 03:26:41 +0800
Subject: [PATCH] utils/pool: use sync.pool to reduce the pressure of garbage
 collection

---
 src/frp/models/msg/process.go | 11 +++++--
 src/frp/utils/pool/pool.go    | 58 +++++++++++++++++++++++++++++++++++
 src/frp/utils/vhost/http.go   |  9 ++++--
 src/frp/utils/vhost/https.go  |  5 ++-
 4 files changed, 77 insertions(+), 6 deletions(-)
 create mode 100644 src/frp/utils/pool/pool.go

diff --git a/src/frp/models/msg/process.go b/src/frp/models/msg/process.go
index 43e34ab..55d99ba 100644
--- a/src/frp/models/msg/process.go
+++ b/src/frp/models/msg/process.go
@@ -26,6 +26,7 @@ import (
 	"frp/utils/conn"
 	"frp/utils/log"
 	"frp/utils/pcrypto"
+	"frp/utils/pool"
 )
 
 // will block until connection close
@@ -115,7 +116,10 @@ func pipeDecrypt(r *conn.Conn, w *conn.Conn, conf config.BaseConf, needRecord bo
 		return fmt.Errorf("Pcrypto Init error: %v", err)
 	}
 
-	buf := make([]byte, 5*1024+4)
+	// get []byte from buffer pool
+	buf := pool.GetBuf(5*1024 + 4)
+	defer pool.PutBuf(buf)
+
 	var left, res []byte
 	var cnt int = -1
 
@@ -198,7 +202,10 @@ func pipeEncrypt(r *conn.Conn, w *conn.Conn, conf config.BaseConf, needRecord bo
 		}()
 	}
 
-	buf := make([]byte, 5*1024)
+	// get []byte from buffer pool
+	buf := pool.GetBuf(5*1024 + 4)
+	defer pool.PutBuf(buf)
+
 	for {
 		n, err := r.Read(buf)
 		if err != nil {
diff --git a/src/frp/utils/pool/pool.go b/src/frp/utils/pool/pool.go
new file mode 100644
index 0000000..5b7c6d2
--- /dev/null
+++ b/src/frp/utils/pool/pool.go
@@ -0,0 +1,58 @@
+// Copyright 2016 fatedier, fatedier@gmail.com
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pool
+
+import "sync"
+
+var (
+	bufPool5k sync.Pool
+	bufPool2k sync.Pool
+	bufPool1k sync.Pool
+	bufPool   sync.Pool
+)
+
+func GetBuf(size int) []byte {
+	var x interface{}
+	if size >= 5*1024 {
+		x = bufPool5k.Get()
+	} else if size >= 2*1024 {
+		x = bufPool2k.Get()
+	} else if size >= 1*1024 {
+		x = bufPool1k.Get()
+	} else {
+		x = bufPool.Get()
+	}
+	if x == nil {
+		return make([]byte, size)
+	}
+	buf := x.([]byte)
+	if cap(buf) < size {
+		return make([]byte, size)
+	}
+	return buf[:size]
+}
+
+func PutBuf(buf []byte) {
+	size := cap(buf)
+	if size >= 5*1024 {
+		bufPool5k.Put(buf)
+	} else if size >= 2*1024 {
+		bufPool2k.Put(buf)
+	} else if size >= 1*1024 {
+		bufPool1k.Put(buf)
+	} else {
+		bufPool.Put(buf)
+	}
+}
diff --git a/src/frp/utils/vhost/http.go b/src/frp/utils/vhost/http.go
index 4bc720c..aa66b6a 100644
--- a/src/frp/utils/vhost/http.go
+++ b/src/frp/utils/vhost/http.go
@@ -26,6 +26,7 @@ import (
 	"time"
 
 	"frp/utils/conn"
+	"frp/utils/pool"
 )
 
 type HttpMuxer struct {
@@ -61,9 +62,11 @@ func HttpHostNameRewrite(c *conn.Conn, rewriteHost string) (_ net.Conn, err erro
 }
 
 func hostNameRewrite(request io.Reader, rewriteHost string) (_ []byte, err error) {
-	buffer := make([]byte, 1024)
-	request.Read(buffer)
-	retBuffer, err := parseRequest(buffer, rewriteHost)
+	buf := pool.GetBuf(1024)
+	defer pool.PutBuf(buf)
+
+	request.Read(buf)
+	retBuffer, err := parseRequest(buf, rewriteHost)
 	return retBuffer, err
 }
 
diff --git a/src/frp/utils/vhost/https.go b/src/frp/utils/vhost/https.go
index eedfab3..7f96365 100644
--- a/src/frp/utils/vhost/https.go
+++ b/src/frp/utils/vhost/https.go
@@ -22,6 +22,7 @@ import (
 	"time"
 
 	"frp/utils/conn"
+	"frp/utils/pool"
 )
 
 const (
@@ -52,7 +53,9 @@ func NewHttpsMuxer(listener *conn.Listener, timeout time.Duration) (*HttpsMuxer,
 }
 
 func readHandshake(rd io.Reader) (host string, err error) {
-	data := make([]byte, 1024)
+	data := pool.GetBuf(1024)
+	origin := data
+	defer pool.PutBuf(origin)
 	length, err := rd.Read(data)
 	if err != nil {
 		return