From 0bb72e0f6088c2eb616ec08803b5fc7d84eff5a4 Mon Sep 17 00:00:00 2001 From: Leonid Bugaev Date: Wed, 7 Jul 2021 02:58:00 +0300 Subject: [PATCH] Add buffer pooling for message data --- gor.go | 1 - tcp/tcp_message.go | 114 +++++++++++++++++++++++++++++++++++++++++---- tcp/tcp_packet.go | 38 +++++++++------ 3 files changed, 129 insertions(+), 24 deletions(-) diff --git a/gor.go b/gor.go index b2af768..c40d624 100644 --- a/gor.go +++ b/gor.go @@ -11,7 +11,6 @@ import ( "os" "os/signal" "runtime" - "runtime/debug" "runtime/pprof" "syscall" "time" diff --git a/tcp/tcp_message.go b/tcp/tcp_message.go index 1e447da..de7a18c 100644 --- a/tcp/tcp_message.go +++ b/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 diff --git a/tcp/tcp_packet.go b/tcp/tcp_packet.go index 4d7a251..37be350 100644 --- a/tcp/tcp_packet.go +++ b/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