Add buffer pooling for message data

gor.go

@@ -11,7 +11,6 @@ import (
 	"os"
 	"os/signal"
 	"runtime"
-	"runtime/debug"
 	"runtime/pprof"
 	"syscall"
 	"time"

tcp/tcp_message.go

@@ -12,6 +12,95 @@ import (
 	"github.com/buger/goreplay/size"
 )
 
+var bufferPool = NewBufferPool(1000, 1)
+
+type buf struct {
+	b       []byte
+	created time.Time
+	gc      bool
+}
+
+type bufPool struct {
+	buffers chan *buf
+	ttl     int
+}
+
+func NewBufferPool(max int, ttl int) *bufPool {
+	pool := &bufPool{
+		buffers: make(chan *buf, max),
+		ttl:     ttl,
+	}
+
+	// Ensure that memory released over time
+	go func() {
+		var released int
+		// GC
+		for {
+			for i := 0; i < 100; i++ {
+				select {
+				case c := <-pool.buffers:
+					if now.Sub(c.created) < time.Duration(ttl)*time.Second {
+						select {
+						case pool.buffers <- c:
+						default:
+							c.b = nil
+							c.gc = true
+							released++
+						}
+					} else {
+						// Else GC
+						c.b = nil
+						c.gc = true
+						released++
+					}
+				default:
+					break
+				}
+			}
+
+			time.Sleep(1000 * time.Millisecond)
+		}
+	}()
+
+	return pool
+}
+
+// Borrow a Client from the pool.
+func (p *bufPool) Get() *buf {
+	var c *buf
+	select {
+	case c = <-p.buffers:
+	default:
+		c = new(buf)
+		c.b = make([]byte, 1024)
+		c.created = now
+
+		// Use this technique to find if pool leaks, and objects get GCd
+		//
+		// runtime.SetFinalizer(c, func(p *buf) {
+		// 	if !p.gc {
+		// 		panic("Pool leak")
+		// 	}
+		// })
+	}
+	return c
+}
+
+// Return returns a Client to the pool.
+func (p *bufPool) Put(c *buf) {
+	select {
+	case p.buffers <- c:
+	default:
+		c.gc = true
+		c.b = nil
+		// if pool overloaded, let it go
+	}
+}
+
+func (p *bufPool) Len() int {
+	return len(p.buffers)
+}
+
 // Stats every message carry its own stats object
 type Stats struct {
 	LostData  int
@@ -31,6 +120,7 @@ type Message struct {
 	packets  []*Packet
 	parser   *MessageParser
 	feedback interface{}
+	dataBuf  *buf
 	Stats
 }
 
@@ -128,18 +218,20 @@ func (m *Message) PacketData() [][]byte {
 
 // Data returns data in this message
 func (m *Message) Data() []byte {
-	var totalLen int
-	for _, p := range m.packets {
-		totalLen += len(p.Payload)
-	}
-	tmp := make([]byte, totalLen)
+	m.dataBuf = bufferPool.Get()
 
-	var i int
-	for _, p := range m.packets {
-		i += copy(tmp[i:], p.Payload)
+	// var totalLen int
+	// for _, p := range m.packets {
+	// 	totalLen += len(p.Payload)
+	// }
+	// tmp := make([]byte, totalLen)
+	var n int
+	if m.dataBuf == nil {
+		panic("asdsd")
 	}
+	m.dataBuf.b, n = copySlice(m.dataBuf.b, m.PacketData()...)
 
-	return tmp
+	return m.dataBuf.b[:n]
 }
 
 // SetProtocolState set feedback/data that can be used later, e.g with End or Start hint
@@ -162,6 +254,10 @@ func (m *Message) Finalize() {
 	for _, p := range m.packets {
 		packetPool.Put(p)
 	}
+
+	if m.dataBuf != nil {
+		bufferPool.Put(m.dataBuf)
+	}
 }
 
 // Emitter message handler

tcp/tcp_packet.go

@@ -11,13 +11,23 @@ import (
 	"github.com/google/gopacket"
 )
 
-func copySlice(b, a []byte) []byte {
-	if cap(b) < len(a) {
-		diff := (cap(b) - len(b)) + len(a)
-		b = append(b, make([]byte, diff)...)
+func copySlice(to []byte, from ...[]byte) ([]byte, int) {
+	var totalLen int
+	for _, s := range from {
+		totalLen += len(s)
 	}
-	copy(b, a)
-	return b
+
+	if cap(to) < totalLen {
+		diff := (cap(to) - len(to)) + totalLen
+		to = append(to, make([]byte, diff)...)
+	}
+
+	var i int
+	for _, s := range from {
+		i += copy(to[i:], s)
+	}
+
+	return to, i
 }
 
 var now time.Time
@@ -32,17 +42,17 @@ func init() {
 	}()
 }
 
-var packetPool = NewPool(10000, 1)
+var packetPool = NewPacketPool(10000, 1)
 
 // Pool holds Clients.
-type Pool struct {
+type pktPool struct {
 	packets chan *Packet
 	ttl     int
 }
 
 // NewPool creates a new pool of Clients.
-func NewPool(max int, ttl int) *Pool {
-	pool := &Pool{
+func NewPacketPool(max int, ttl int) *pktPool {
+	pool := &pktPool{
 		packets: make(chan *Packet, max),
 		ttl:     ttl,
 	}
@@ -87,7 +97,7 @@ func NewPool(max int, ttl int) *Pool {
 }
 
 // Borrow a Client from the pool.
-func (p *Pool) Get() *Packet {
+func (p *pktPool) Get() *Packet {
 	var c *Packet
 	select {
 	case c = <-p.packets:
@@ -107,7 +117,7 @@ func (p *Pool) Get() *Packet {
 }
 
 // Return returns a Client to the pool.
-func (p *Pool) Put(c *Packet) {
+func (p *pktPool) Put(c *Packet) {
 	select {
 	case p.packets <- c:
 	default:
@@ -117,7 +127,7 @@ func (p *Pool) Put(c *Packet) {
 	}
 }
 
-func (p *Pool) Len() int {
+func (p *pktPool) Len() int {
 	return len(p.packets)
 }
 
@@ -262,7 +272,7 @@ func (pckt *Packet) parse(data []byte, lType, lTypeLen int, cp *gopacket.Capture
 	pckt.RST = transLayer[13]&0x04 != 0
 	pckt.ACK = transLayer[13]&0x10 != 0
 	pckt.Lost = uint32(cp.Length - cp.CaptureLength)
-	pckt.buf = copySlice(pckt.buf, ndata[dOf:])
+	pckt.buf, _ = copySlice(pckt.buf, ndata[dOf:])
 	pckt.Payload = pckt.buf[:len(ndata[dOf:])]
 
 	return nil