From 950e59fa3d45aac222cf91b19c55b5d507cc6892 Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Tue, 26 Jul 2016 15:36:28 +0200
Subject: [PATCH] add a sync.Pool of byte buffers with maximum packet size as
 cap

ref #217
---
 buffer_pool.go      | 23 +++++++++++++++++++++++
 buffer_pool_test.go | 26 ++++++++++++++++++++++++++
 packet_packer.go    |  2 +-
 server.go           |  3 ++-
 session.go          |  2 +-
 session_test.go     |  4 +++-
 6 files changed, 56 insertions(+), 4 deletions(-)
 create mode 100644 buffer_pool.go
 create mode 100644 buffer_pool_test.go

diff --git a/buffer_pool.go b/buffer_pool.go
new file mode 100644
index 00000000..1c81308e
--- /dev/null
+++ b/buffer_pool.go
@@ -0,0 +1,23 @@
+package quic
+
+import (
+	"sync"
+
+	"github.com/lucas-clemente/quic-go/protocol"
+)
+
+var bufferPool sync.Pool
+
+func getPacketBuffer() []byte {
+	return bufferPool.Get().([]byte)
+}
+
+func putPacketBuffer(buf []byte) {
+	bufferPool.Put(buf[:0])
+}
+
+func init() {
+	bufferPool.New = func() interface{} {
+		return make([]byte, 0, protocol.MaxPacketSize)
+	}
+}
diff --git a/buffer_pool_test.go b/buffer_pool_test.go
new file mode 100644
index 00000000..9fa1a53f
--- /dev/null
+++ b/buffer_pool_test.go
@@ -0,0 +1,26 @@
+package quic
+
+import (
+	"github.com/lucas-clemente/quic-go/protocol"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("Buffer Pool", func() {
+	It("returns buffers of correct len and cap", func() {
+		buf := getPacketBuffer()
+		Expect(buf).To(HaveLen(0))
+		Expect(buf).To(HaveCap(int(protocol.MaxPacketSize)))
+	})
+
+	It("zeroes put buffers' length", func() {
+		for i := 0; i < 1000; i++ {
+			buf := getPacketBuffer()
+			putPacketBuffer(buf[0:10])
+			buf = getPacketBuffer()
+			Expect(buf).To(HaveLen(0))
+			Expect(buf).To(HaveCap(int(protocol.MaxPacketSize)))
+		}
+	})
+})
diff --git a/packet_packer.go b/packet_packer.go
index 946fe91e..69363632 100644
--- a/packet_packer.go
+++ b/packet_packer.go
@@ -110,7 +110,7 @@ func (p *packetPacker) packPacket(stopWaitingFrame *frames.StopWaitingFrame, con
 		}
 	}
 
-	raw := make([]byte, 0, protocol.MaxPacketSize)
+	raw := getPacketBuffer()
 	buffer := bytes.NewBuffer(raw)
 
 	if responsePublicHeader.WritePublicHeader(buffer, p.version) != nil {
diff --git a/server.go b/server.go
index 2652dc83..9fc6cb77 100644
--- a/server.go
+++ b/server.go
@@ -86,7 +86,8 @@ func (s *Server) Serve(conn *net.UDPConn) error {
 	s.connMutex.Unlock()
 
 	for {
-		data := make([]byte, protocol.MaxPacketSize)
+		data := getPacketBuffer()
+		data = data[:protocol.MaxPacketSize]
 		n, remoteAddr, err := conn.ReadFromUDP(data)
 		if err != nil {
 			if strings.HasSuffix(err.Error(), "use of closed network connection") {
diff --git a/session.go b/session.go
index 434dc6d4..a7d25b31 100644
--- a/session.go
+++ b/session.go
@@ -531,10 +531,10 @@ func (s *Session) sendPacket() error {
 		s.delayedAckOriginTime = time.Time{}
 
 		err = s.conn.write(packet.raw)
+		putPacketBuffer(packet.raw)
 		if err != nil {
 			return err
 		}
-
 	}
 }
 
diff --git a/session_test.go b/session_test.go
index 6d255416..41561e14 100644
--- a/session_test.go
+++ b/session_test.go
@@ -27,7 +27,9 @@ type mockConnection struct {
 }
 
 func (m *mockConnection) write(p []byte) error {
-	m.written = append(m.written, p)
+	b := make([]byte, len(p))
+	copy(b, p)
+	m.written = append(m.written, b)
 	return nil
 }