feat: Add connection

codec/codec.go

@@ -1,6 +1,8 @@
 package codec
 
 import (
+	"bufio"
+	"context"
 	"encoding/binary"
 	"net"
 	"sync"
@@ -12,11 +14,13 @@ type Message struct {
 	Headers map[string]string `msgpack:"h"`
 	Queue   string            `msgpack:"q"`
 	Payload []byte            `msgpack:"p"`
-	m       sync.RWMutex
-	Command consts.CMD `msgpack:"c"`
+	Command consts.CMD        `msgpack:"c"`
 }
 
 func NewMessage(cmd consts.CMD, payload []byte, queue string, headers map[string]string) *Message {
+	if headers == nil {
+		headers = make(map[string]string)
+	}
 	return &Message{
 		Headers: headers,
 		Queue:   queue,
@@ -26,8 +30,6 @@ func NewMessage(cmd consts.CMD, payload []byte, queue string, headers map[string
 }
 
 func (m *Message) Serialize() ([]byte, error) {
-	m.m.RLock()
-	defer m.m.RUnlock()
 	data, err := Marshal(m)
 	if err != nil {
 		return nil, err
@@ -40,37 +42,66 @@ func Deserialize(data []byte) (*Message, error) {
 	if err := Unmarshal(data, &msg); err != nil {
 		return nil, err
 	}
-
 	return &msg, nil
 }
 
-func SendMessage(conn net.Conn, msg *Message) error {
+var byteBufferPool = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, 0, 4096)
+	},
+}
+
+func SendMessage(ctx context.Context, conn net.Conn, msg *Message) error {
 	data, err := msg.Serialize()
 	if err != nil {
 		return err
 	}
 
-	length := make([]byte, 4)
-	binary.BigEndian.PutUint32(length, uint32(len(data)))
+	totalLength := 4 + len(data)
+	buffer := byteBufferPool.Get().([]byte)
+	if cap(buffer) < totalLength {
+		buffer = make([]byte, totalLength)
+	} else {
+		buffer = buffer[:totalLength]
+	}
+	defer byteBufferPool.Put(buffer)
 
-	if _, err := conn.Write(length); err != nil {
-		return err
+	binary.BigEndian.PutUint32(buffer[:4], uint32(len(data)))
+	copy(buffer[4:], data)
+
+	writer := bufio.NewWriter(conn)
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+		if _, err := writer.Write(buffer); err != nil {
+			return err
+		}
 	}
-	if _, err := conn.Write(data); err != nil {
-		return err
-	}
-	return nil
+
+	return writer.Flush()
 }
 
-func ReadMessage(conn net.Conn) (*Message, error) {
+func ReadMessage(ctx context.Context, conn net.Conn) (*Message, error) {
 	lengthBytes := make([]byte, 4)
 	if _, err := conn.Read(lengthBytes); err != nil {
 		return nil, err
 	}
 	length := binary.BigEndian.Uint32(lengthBytes)
-	data := make([]byte, length)
-	if _, err := conn.Read(data); err != nil {
-		return nil, err
+	data := byteBufferPool.Get().([]byte)[:length]
+	defer byteBufferPool.Put(data)
+	totalRead := 0
+	for totalRead < int(length) {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+			n, err := conn.Read(data[totalRead:])
+			if err != nil {
+				return nil, err
+			}
+			totalRead += n
+		}
 	}
-	return Deserialize(data)
+	return Deserialize(data[:length])
 }