connection.go | 10 ++++++++++ peer_protocol/protocol.go | 13 ++++++++++++- diff --git a/connection.go b/connection.go index d98b1c273001f0e32e697637485c94f493bcde10..b25e26629ce345b1cb6b7ce2240019f81c7c2008 100644 --- a/connection.go +++ b/connection.go @@ -698,9 +698,16 @@ // and exit. Returning will end the connection. 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 @@ case pp.HaveNone: 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 da45d6772dc60f612e3a15b390aae7f4f6cb1aef..1e572bd89354bee443a703ff07e3aac501d5f5d7 100644 --- a/peer_protocol/protocol.go +++ b/peer_protocol/protocol.go @@ -8,6 +8,7 @@ "errors" "fmt" "io" "io/ioutil" + "sync" ) type ( @@ -123,6 +124,7 @@ } type Decoder struct { R *bufio.Reader + Pool *sync.Pool MaxLength Integer // TODO: Should this include the length header or not? } @@ -197,7 +199,16 @@ } 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 {