Change to use a pool of byte buffer to encode data

Author: daichirata

buffer.go

@@ -5,44 +5,43 @@ import (
 )
 
 type buffer struct {
-	buf   []byte
+	buf   []*Message
 	mu    sync.Mutex
 	Dirty chan struct{}
 }
 
 func newBuffer() buffer {
 	return buffer{
-		buf:   []byte{},
 		Dirty: make(chan struct{}),
 	}
 }
 
-func (buffer *buffer) Add(raw []byte) {
+func (buffer *buffer) Add(message *Message) {
 	buffer.mu.Lock()
 	defer buffer.mu.Unlock()
 
-	buffer.buf = append(buffer.buf, raw...)
+	buffer.buf = append(buffer.buf, message)
 	go func() {
 		buffer.Dirty <- struct{}{}
 	}()
 }
 
-func (buffer *buffer) Remove() []byte {
+func (buffer *buffer) Remove() []*Message {
 	buffer.mu.Lock()
 	defer buffer.mu.Unlock()
 
 	if len(buffer.buf) == 0 {
 		return nil
 	}
 
-	data := buffer.buf
+	m := buffer.buf
 	buffer.buf = buffer.buf[:0]
-	return data
+	return m
 }
 
-func (buffer *buffer) Back(raw []byte) {
+func (buffer *buffer) Back(messages []*Message) {
 	buffer.mu.Lock()
 	defer buffer.mu.Unlock()
 
-	buffer.buf = append(buffer.buf, raw...)
+	buffer.buf = append(buffer.buf, messages...)
 }

logger.go

@@ -7,8 +7,6 @@ import (
 	"sync"
 	"time"
 
-	"bytes"
-
 	"gopkg.in/vmihailenco/msgpack.v2"
 )
 
@@ -66,12 +64,12 @@ func (logger *Logger) PostWithTime(tag string, t time.Time, obj interface{}) err
 		obj,
 	}
 
-	buf := bytes.NewBuffer([]byte{})
-	enc := msgpack.NewEncoder(buf)
+	m := getMessage()
+	enc := msgpack.NewEncoder(m.buf)
 	if err := enc.Encode(record); err != nil {
 		return err
 	}
-	logger.buf.Add(buf.Bytes())
+	logger.buf.Add(m)
 	return nil
 }
 
@@ -119,24 +117,31 @@ func (logger *Logger) disconnect() error {
 const maxWriteAttempts = 3
 
 func (logger *Logger) send() error {
-	data := logger.buf.Remove()
-	if len(data) == 0 {
+	messages := logger.buf.Remove()
+	if len(messages) == 0 {
 		return nil
 	}
+	var data []byte
+	for _, m := range messages {
+		data = append(data, m.buf.Bytes()...)
+	}
 
 	var err error
 	for i := 0; i < maxWriteAttempts; i++ {
 		err = logger.connect()
 		if err == nil {
 			_, err := logger.conn.Write(data)
 			if err == nil {
+				for _, m := range messages {
+					putMessage(m)
+				}
 				break
 			}
 		}
 		logger.disconnect()
 	}
 	if err != nil {
-		logger.buf.Back(data)
+		logger.buf.Back(messages)
 	}
 	return err
 }
@@ -149,12 +154,18 @@ func (logger *Logger) start() {
 		for {
 			select {
 			case <-logger.done:
-				logger.send()
+				err := logger.send()
+				if err != nil {
+					panic(err)
+				}
 				return
 			case <-logger.buf.Dirty:
 			case <-ticker.C:
 			}
-			logger.send()
+			err := logger.send()
+			if err != nil {
+				panic(err)
+			}
 		}
 	}()
 }

message.go

@@ -0,0 +1,32 @@
+package fluent
+
+import (
+	"bytes"
+	"sync"
+)
+
+type Message struct {
+	buf *bytes.Buffer
+}
+
+var messagePool = sync.Pool{
+	New: func() interface{} {
+		return newMessage()
+	},
+}
+
+func getMessage() *Message {
+	message := messagePool.Get().(*Message)
+	message.buf.Reset()
+	return message
+}
+
+func putMessage(message *Message) {
+	messagePool.Put(message)
+}
+
+func newMessage() *Message {
+	return &Message{
+		buf: bytes.NewBuffer([]byte{}),
+	}
+}