Merge pull request #14 from libp2p/fix/memory-improvements

combine writes and avoid a few more allocations

Author: Steven Allen

chan.go

@@ -38,9 +38,6 @@ func (s *Chan) ReadFromWithPool(r io.Reader, p *pool.BufferPool) {
 // ReadFrom wraps the given io.Reader with a msgio.Reader, reads all
 // messages, ands sends them down the channel.
 func (s *Chan) readFrom(mr Reader) {
-	// single reader, no need for Mutex
-	mr.(*reader).lock = new(nullLocker)
-
 Loop:
 	for {
 		buf, err := mr.ReadMsg()
@@ -74,8 +71,6 @@ func (s *Chan) WriteTo(w io.Writer) {
 	// if bottleneck, cycle around a set of buffers
 	mw := NewWriter(w)
 
-	// single writer, no need for Mutex
-	mw.(*writer).lock = new(nullLocker)
 Loop:
 	for {
 		select {

msgio.go

@@ -76,13 +76,20 @@ type ReadWriteCloser interface {
 type writer struct {
 	W io.Writer
 
-	lock sync.Locker
+	pool *pool.BufferPool
+	lock sync.Mutex
 }
 
 // NewWriter wraps an io.Writer with a msgio framed writer. The msgio.Writer
 // will write the length prefix of every message written.
 func NewWriter(w io.Writer) WriteCloser {
-	return &writer{W: w, lock: new(sync.Mutex)}
+	return NewWriterWithPool(w, pool.GlobalPool)
+}
+
+// NewWriterWithPool is identical to NewWriter but allows the user to pass a
+// custom buffer pool.
+func NewWriterWithPool(w io.Writer, p *pool.BufferPool) WriteCloser {
+	return &writer{W: w, pool: p}
 }
 
 func (s *writer) Write(msg []byte) (int, error) {
@@ -96,10 +103,13 @@ func (s *writer) Write(msg []byte) (int, error) {
 func (s *writer) WriteMsg(msg []byte) (err error) {
 	s.lock.Lock()
 	defer s.lock.Unlock()
-	if err := WriteLen(s.W, len(msg)); err != nil {
-		return err
-	}
-	_, err = s.W.Write(msg)
+
+	buf := s.pool.Get(len(msg) + lengthSize)
+	NBO.PutUint32(buf, uint32(len(msg)))
+	copy(buf[lengthSize:], msg)
+	_, err = s.W.Write(buf)
+	s.pool.Put(buf)
+
 	return err
 }
 
@@ -114,10 +124,10 @@ func (s *writer) Close() error {
 type reader struct {
 	R io.Reader
 
-	lbuf []byte
+	lbuf [lengthSize]byte
 	next int
 	pool *pool.BufferPool
-	lock sync.Locker
+	lock sync.Mutex
 	max  int // the maximal message size (in bytes) this reader handles
 }
 
@@ -137,10 +147,8 @@ func NewReaderWithPool(r io.Reader, p *pool.BufferPool) ReadCloser {
 	}
 	return &reader{
 		R:    r,
-		lbuf: make([]byte, lengthSize),
 		next: -1,
 		pool: p,
-		lock: new(sync.Mutex),
 		max:  defaultMaxSize,
 	}
 }
@@ -156,7 +164,7 @@ func (s *reader) NextMsgLen() (int, error) {
 
 func (s *reader) nextMsgLen() (int, error) {
 	if s.next == -1 {
-		n, err := ReadLen(s.R, s.lbuf)
+		n, err := ReadLen(s.R, s.lbuf[:])
 		if err != nil {
 			return 0, err
 		}

varint.go

@@ -12,16 +12,21 @@ import (
 type varintWriter struct {
 	W io.Writer
 
-	lbuf [binary.MaxVarintLen64]byte // for encoding varints
-	lock sync.Mutex                  // for threadsafe writes
+	pool *pool.BufferPool
+	lock sync.Mutex // for threadsafe writes
 }
 
 // NewVarintWriter wraps an io.Writer with a varint msgio framed writer.
 // The msgio.Writer will write the length prefix of every message written
 // as a varint, using https://golang.org/pkg/encoding/binary/#PutUvarint
 func NewVarintWriter(w io.Writer) WriteCloser {
+	return NewVarintWriterWithPool(w, pool.GlobalPool)
+}
+
+func NewVarintWriterWithPool(w io.Writer, p *pool.BufferPool) WriteCloser {
 	return &varintWriter{
-		W: w,
+		pool: p,
+		W:    w,
 	}
 }
 
@@ -37,12 +42,12 @@ func (s *varintWriter) WriteMsg(msg []byte) error {
 	s.lock.Lock()
 	defer s.lock.Unlock()
 
-	length := uint64(len(msg))
-	n := binary.PutUvarint(s.lbuf[:], length)
-	if _, err := s.W.Write(s.lbuf[:n]); err != nil {
-		return err
-	}
-	_, err := s.W.Write(msg)
+	buf := s.pool.Get(len(msg) + binary.MaxVarintLen64)
+	n := binary.PutUvarint(buf, uint64(len(msg)))
+	n += copy(buf[n:], msg)
+	_, err := s.W.Write(buf[:n])
+	s.pool.Put(buf)
+
 	return err
 }