From d3963eedfd08b88fa373e635fbc1b9aa59d06119 Mon Sep 17 00:00:00 2001
From: Ye Yin <eyniy@qq.com>
Date: Mon, 12 Sep 2016 15:10:11 +0800
Subject: [PATCH] Use sync.pool for decode buffer

---
 connection.go             | 10 ++++++++++
 peer_protocol/protocol.go | 13 ++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/connection.go b/connection.go
index d98b1c27..b25e2662 100644
--- a/connection.go
+++ b/connection.go
@@ -698,9 +698,16 @@ func (c *connection) lastHelpful() (ret time.Time) {
 func (c *connection) mainReadLoop() error {
 	t := c.t
 	cl := t.cl
+	pool := &sync.Pool{
+		New: func() interface{} {
+			return make([]byte, t.chunkSize)
+		},
+	}
+
 	decoder := pp.Decoder{
 		R:         bufio.NewReader(c.rw),
 		MaxLength: 256 * 1024,
+		Pool:      pool,
 	}
 	for {
 		cl.mu.Unlock()
@@ -774,6 +781,9 @@ func (c *connection) mainReadLoop() error {
 			err = c.peerSentHaveNone()
 		case pp.Piece:
 			cl.downloadedChunk(t, c, &msg)
+			if len(msg.Piece) == int(t.chunkSize) {
+				pool.Put(msg.Piece)
+			}
 		case pp.Extended:
 			switch msg.ExtendedID {
 			case pp.HandshakeExtendedID:
diff --git a/peer_protocol/protocol.go b/peer_protocol/protocol.go
index da45d677..1e572bd8 100644
--- a/peer_protocol/protocol.go
+++ b/peer_protocol/protocol.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"sync"
 )
 
 type (
@@ -123,6 +124,7 @@ func (msg Message) MarshalBinary() (data []byte, err error) {
 
 type Decoder struct {
 	R         *bufio.Reader
+	Pool      *sync.Pool
 	MaxLength Integer // TODO: Should this include the length header or not?
 }
 
@@ -197,7 +199,16 @@ func (d *Decoder) Decode(msg *Message) (err error) {
 		if err != nil {
 			break
 		}
-		msg.Piece, err = ioutil.ReadAll(r)
+		//msg.Piece, err = ioutil.ReadAll(r)
+		b := d.Pool.Get().([]byte)
+		n, err := io.ReadFull(r, b)
+		if err != nil {
+			if err != io.ErrUnexpectedEOF || n != int(length-9) {
+				return err
+			}
+			b = b[0:n]
+		}
+		msg.Piece = b
 	case Extended:
 		msg.ExtendedID, err = r.ReadByte()
 		if err != nil {