From 757fd11af50e220b7bd5bb0ca66b31b9b80aed6f Mon Sep 17 00:00:00 2001
From: Jason <xjasonlyu@gmail.com>
Date: Mon, 12 Aug 2019 15:15:35 +0800
Subject: [PATCH] fix buffer pool

---
 core/buffer_pool.go         | 23 +++++++++++++++++++++++
 core/output_export.go       |  4 ++--
 core/tcp_callback_export.go |  4 ++--
 core/udp_callback_export.go |  4 ++--
 4 files changed, 29 insertions(+), 6 deletions(-)
 create mode 100644 core/buffer_pool.go

diff --git a/core/buffer_pool.go b/core/buffer_pool.go
new file mode 100644
index 0000000..c77d3dc
--- /dev/null
+++ b/core/buffer_pool.go
@@ -0,0 +1,23 @@
+package core
+
+import (
+	"sync"
+)
+
+const defaultBufferSize = 2 * 1024
+
+var bufPool = sync.Pool{New: func() interface{} { return make([]byte, defaultBufferSize) }}
+
+func newBytes(size int) []byte {
+	if size <= defaultBufferSize {
+		return bufPool.Get().([]byte)
+	} else {
+		return make([]byte, size)
+	}
+}
+
+func freeBytes(b []byte) {
+	if len(b) >= defaultBufferSize {
+		bufPool.Put(b)
+	}
+}
diff --git a/core/output_export.go b/core/output_export.go
index 79ae18a..749cbee 100644
--- a/core/output_export.go
+++ b/core/output_export.go
@@ -19,10 +19,10 @@ func output(p *C.struct_pbuf) C.err_t {
 		buf := (*[1 << 30]byte)(unsafe.Pointer(p.payload))[:totlen:totlen]
 		OutputFn(buf[:totlen])
 	} else {
-		buf := NewBytes(totlen)
+		buf := newBytes(totlen)
 		C.pbuf_copy_partial(p, unsafe.Pointer(&buf[0]), p.tot_len, 0) // data copy here!
 		OutputFn(buf[:totlen])
-		FreeBytes(buf)
+		freeBytes(buf)
 	}
 	return C.ERR_OK
 }
diff --git a/core/tcp_callback_export.go b/core/tcp_callback_export.go
index 0b75f99..b3bea00 100644
--- a/core/tcp_callback_export.go
+++ b/core/tcp_callback_export.go
@@ -81,8 +81,8 @@ func tcpRecvFn(arg unsafe.Pointer, tpcb *C.struct_tcp_pcb, p *C.struct_pbuf, err
 	if p.tot_len == p.len {
 		buf = (*[1 << 30]byte)(unsafe.Pointer(p.payload))[:totlen:totlen]
 	} else {
-		buf = NewBytes(totlen)
-		defer FreeBytes(buf)
+		buf = newBytes(totlen)
+		defer freeBytes(buf)
 		C.pbuf_copy_partial(p, unsafe.Pointer(&buf[0]), p.tot_len, 0)
 	}
 
diff --git a/core/udp_callback_export.go b/core/udp_callback_export.go
index db885a7..4ff0167 100644
--- a/core/udp_callback_export.go
+++ b/core/udp_callback_export.go
@@ -53,8 +53,8 @@ func udpRecvFn(arg unsafe.Pointer, pcb *C.struct_udp_pcb, p *C.struct_pbuf, addr
 	if p.tot_len == p.len {
 		buf = (*[1 << 30]byte)(unsafe.Pointer(p.payload))[:totlen:totlen]
 	} else {
-		buf = NewBytes(totlen)
-		defer FreeBytes(buf)
+		buf = newBytes(totlen)
+		defer freeBytes(buf)
 		C.pbuf_copy_partial(p, unsafe.Pointer(&buf[0]), p.tot_len, 0)
 	}