diff --git a/alloc.go b/alloc.go new file mode 100644 index 0000000..57db6e6 --- /dev/null +++ b/alloc.go @@ -0,0 +1,65 @@ +package smux + +import ( + "sync" + + "github.com/pkg/errors" +) + +var defaultAllocator *Allocator + +func init() { + defaultAllocator = NewAllocator() +} + +// allocator for incoming frames +// to prevent write after zeroing +type Allocator struct { + buffers []sync.Pool +} + +func NewAllocator() *Allocator { + alloc := new(Allocator) + alloc.buffers = make([]sync.Pool, 17) // 1B -> 64K + for k := range alloc.buffers { + i := k + alloc.buffers[k].New = func() interface{} { + return make([]byte, 1< 65536 { + return nil + } + + bits := msb(size) + if size <= 1< 65536 || cap(buf) != 1<>= 1 + if size == 0 { + return pos + } + pos++ + } +} diff --git a/alloc_test.go b/alloc_test.go new file mode 100644 index 0000000..3e1a2b3 --- /dev/null +++ b/alloc_test.go @@ -0,0 +1,72 @@ +package smux + +import "testing" + +func TestAllocGet(t *testing.T) { + alloc := NewAllocator() + if alloc.Get(0) != nil { + t.Fatal(0) + } + if len(alloc.Get(1)) != 1 { + t.Fatal(1) + } + if len(alloc.Get(2)) != 2 { + t.Fatal(2) + } + if len(alloc.Get(3)) != 3 || cap(alloc.Get(3)) != 4 { + t.Fatal(3) + } + if len(alloc.Get(4)) != 4 { + t.Fatal(4) + } + if len(alloc.Get(1023)) != 1023 || cap(alloc.Get(1023)) != 1024 { + t.Fatal(1023) + } + if len(alloc.Get(1024)) != 1024 { + t.Fatal(1024) + } + if len(alloc.Get(65536)) != 65536 { + t.Fatal(65536) + } + if alloc.Get(65537) != nil { + t.Fatal(65537) + } +} + +func TestAllocPut(t *testing.T) { + alloc := NewAllocator() + if err := alloc.Put(nil); err == nil { + t.Fatal("put nil misbehavior") + } + if err := alloc.Put(make([]byte, 3, 3)); err == nil { + t.Fatal("put elem:3 []bytes misbehavior") + } + if err := alloc.Put(make([]byte, 4, 4)); err != nil { + t.Fatal("put elem:4 []bytes misbehavior") + } + if err := alloc.Put(make([]byte, 1023, 1024)); err != nil { + t.Fatal("put elem:1024 []bytes misbehavior") + } + if err := alloc.Put(make([]byte, 65536, 65536)); err != nil { + t.Fatal("put elem:65536 []bytes misbehavior") + } + if err := alloc.Put(make([]byte, 65537, 65537)); err == nil { + t.Fatal("put elem:65537 []bytes misbehavior") + } +} + +func TestAllocPutThenGet(t *testing.T) { + alloc := NewAllocator() + data := alloc.Get(4) + for k := range data { + data[k] = 99 + } + alloc.Put(data) + + newData := alloc.Get(4) + for k := range newData { + if newData[k] != 99 { + t.Fatal("cannot fetch written []bytes from pool") + } + } +} diff --git a/session.go b/session.go index 24f469f..14c434e 100644 --- a/session.go +++ b/session.go @@ -303,7 +303,7 @@ func (s *Session) recvLoop() { s.streamLock.Unlock() case cmdPSH: if hdr.Length() > 0 { - newbuf := make([]byte, hdr.Length()) + newbuf := defaultAllocator.Get(int(hdr.Length())) if written, err := io.ReadFull(s.conn, newbuf); err == nil { s.streamLock.Lock() if stream, ok := s.streams[sid]; ok { diff --git a/stream.go b/stream.go index 5515ee5..6ccc864 100644 --- a/stream.go +++ b/stream.go @@ -12,9 +12,12 @@ import ( // Stream implements net.Conn type Stream struct { - id uint32 - sess *Session - buffers [][]byte + id uint32 + sess *Session + + buffers [][]byte + heads [][]byte // slice heads kept for recycle + bufferLock sync.Mutex frameSize int @@ -65,6 +68,9 @@ func (s *Stream) Read(b []byte) (n int, err error) { if len(s.buffers[0]) == 0 { s.buffers[0] = nil s.buffers = s.buffers[1:] + // full recycle + defaultAllocator.Put(s.heads[0]) + s.heads = s.heads[1:] } } s.bufferLock.Unlock() @@ -217,6 +223,7 @@ func (s *Stream) RemoteAddr() net.Addr { func (s *Stream) pushBytes(buf []byte) (written int, err error) { s.bufferLock.Lock() s.buffers = append(s.buffers, buf) + s.heads = append(s.heads, buf) s.bufferLock.Unlock() return } @@ -226,8 +233,10 @@ func (s *Stream) recycleTokens() (n int) { s.bufferLock.Lock() for k := range s.buffers { n += len(s.buffers[k]) + defaultAllocator.Put(s.heads[k]) } s.buffers = nil + s.heads = nil s.bufferLock.Unlock() return }