mirror of
https://github.com/oarkflow/mq.git
synced 2025-10-04 15:42:49 +08:00
92 lines
2.0 KiB
Go
92 lines
2.0 KiB
Go
package codec
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"net"
|
|
"sync"
|
|
|
|
"github.com/vmihailenco/msgpack/v5"
|
|
|
|
"github.com/oarkflow/mq/consts"
|
|
)
|
|
|
|
// Message represents the structure of our message.
|
|
type Message struct {
|
|
Headers map[string]string `msgpack:"h"`
|
|
Queue string `msgpack:"q"`
|
|
Command consts.CMD `msgpack:"c"`
|
|
Payload []byte `msgpack:"p"` // Using []byte instead of json.RawMessage
|
|
m sync.RWMutex
|
|
}
|
|
|
|
// NewMessage creates a new Message instance.
|
|
func NewMessage(cmd consts.CMD, payload []byte, queue string, headers map[string]string) *Message {
|
|
return &Message{
|
|
Headers: headers,
|
|
Queue: queue,
|
|
Command: cmd,
|
|
Payload: payload,
|
|
}
|
|
}
|
|
|
|
// Serialize encodes the Message to a byte slice using MessagePack.
|
|
func (m *Message) Serialize() ([]byte, error) {
|
|
m.m.RLock()
|
|
defer m.m.RUnlock()
|
|
|
|
data, err := msgpack.Marshal(m)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return data, nil
|
|
}
|
|
|
|
// Deserialize decodes a byte slice to a Message instance using MessagePack.
|
|
func Deserialize(data []byte) (*Message, error) {
|
|
var msg Message
|
|
if err := msgpack.Unmarshal(data, &msg); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &msg, nil
|
|
}
|
|
|
|
// SendMessage sends a Message over a net.Conn.
|
|
func SendMessage(conn net.Conn, msg *Message) error {
|
|
data, err := msg.Serialize()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Send the length of the data followed by the data itself
|
|
length := make([]byte, 4)
|
|
binary.BigEndian.PutUint32(length, uint32(len(data)))
|
|
|
|
if _, err := conn.Write(length); err != nil {
|
|
return err
|
|
}
|
|
if _, err := conn.Write(data); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ReadMessage receives a Message from a net.Conn.
|
|
func ReadMessage(conn net.Conn) (*Message, error) {
|
|
// Read the length of the incoming message
|
|
lengthBytes := make([]byte, 4)
|
|
if _, err := conn.Read(lengthBytes); err != nil {
|
|
return nil, err
|
|
}
|
|
length := binary.BigEndian.Uint32(lengthBytes)
|
|
|
|
// Read the actual message data
|
|
data := make([]byte, length)
|
|
if _, err := conn.Read(data); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return Deserialize(data)
|
|
}
|