- Switch from the flatbuffers to the protobuf

Signed-off-by: Valery Piashchynski <[email protected]>

4 files changed, 58 insertions, 116 deletions

diff --git a/plugins/websockets/executor/executor.go b/plugins/websockets/executor/executor.go
index 69aad7d4..951c9a1a 100644
--- a/plugins/websockets/executor/executor.go
+++ b/plugins/websockets/executor/executor.go
@@ -9,6 +9,7 @@ import (
 	json "github.com/json-iterator/go"
 	"github.com/spiral/errors"
 	"github.com/spiral/roadrunner/v2/pkg/pubsub"
+	"github.com/spiral/roadrunner/v2/pkg/pubsub/message"
 	"github.com/spiral/roadrunner/v2/plugins/logger"
 	"github.com/spiral/roadrunner/v2/plugins/websockets/commands"
 	"github.com/spiral/roadrunner/v2/plugins/websockets/connection"
@@ -63,7 +64,7 @@ func (e *Executor) StartCommandLoop() error { //nolint:gocognit
 			return errors.E(op, err)
 		}
 
-		msg := &pubsub.Message{}
+		msg := &message.Message{}
 
 		err = json.Unmarshal(data, msg)
 		if err != nil {
diff --git a/plugins/websockets/plugin.go b/plugins/websockets/plugin.go
index 4c0edcad..39a4e139 100644
--- a/plugins/websockets/plugin.go
+++ b/plugins/websockets/plugin.go
@@ -25,7 +25,7 @@ import (
 	"github.com/spiral/roadrunner/v2/plugins/websockets/executor"
 	"github.com/spiral/roadrunner/v2/plugins/websockets/pool"
 	"github.com/spiral/roadrunner/v2/plugins/websockets/validator"
-	"github.com/spiral/roadrunner/v2/utils"
+	"google.golang.org/protobuf/proto"
 )
 
 const (
@@ -301,16 +301,21 @@ func (p *Plugin) Publish(m []byte) error {
 	p.Lock()
 	defer p.Unlock()
 
+	msg := &message.Message{}
+	err := proto.Unmarshal(m, msg)
+	if err != nil {
+		return err
+	}
+
 	// Get payload
-	fbsMsg := message.GetRootAsMessage(m, 0)
-	for i := 0; i < fbsMsg.TopicsLength(); i++ {
-		if br, ok := p.pubsubs[utils.AsString(fbsMsg.Broker())]; ok {
-			err := br.Publish(fbsMsg.Table().Bytes)
+	for i := 0; i < len(msg.GetTopics()); i++ {
+		if br, ok := p.pubsubs[msg.GetBroker()]; ok {
+			err := br.Publish(m)
 			if err != nil {
 				return errors.E(err)
 			}
 		} else {
-			p.log.Warn("no such broker", "available", p.pubsubs, "requested", fbsMsg.Broker())
+			p.log.Warn("no such broker", "available", p.pubsubs, "requested", msg.GetBroker())
 		}
 	}
 	return nil
@@ -320,16 +325,21 @@ func (p *Plugin) PublishAsync(m []byte) {
 	go func() {
 		p.Lock()
 		defer p.Unlock()
-		fbsMsg := message.GetRootAsMessage(m, 0)
-		for i := 0; i < fbsMsg.TopicsLength(); i++ {
-			if br, ok := p.pubsubs[utils.AsString(fbsMsg.Broker())]; ok {
-				err := br.Publish(fbsMsg.Table().Bytes)
+		msg := &message.Message{}
+		err := proto.Unmarshal(m, msg)
+		if err != nil {
+			p.log.Error("message unmarshal")
+		}
+
+		// Get payload
+		for i := 0; i < len(msg.GetTopics()); i++ {
+			if br, ok := p.pubsubs[msg.GetBroker()]; ok {
+				err := br.Publish(m)
 				if err != nil {
 					p.log.Error("publish async error", "error", err)
-					return
 				}
 			} else {
-				p.log.Warn("no such broker", "available", p.pubsubs, "requested", fbsMsg.Broker())
+				p.log.Warn("no such broker", "available", p.pubsubs, "requested", msg.GetBroker())
 			}
 		}
 	}()
diff --git a/plugins/websockets/pool/workers_pool.go b/plugins/websockets/pool/workers_pool.go
index 544f3ede..efafb2d3 100644
--- a/plugins/websockets/pool/workers_pool.go
+++ b/plugins/websockets/pool/workers_pool.go
@@ -1,7 +1,6 @@
 package pool
 
 import (
-	"bytes"
 	"sync"
 
 	"github.com/fasthttp/websocket"
@@ -9,14 +8,12 @@ import (
 	"github.com/spiral/roadrunner/v2/pkg/pubsub/message"
 	"github.com/spiral/roadrunner/v2/plugins/logger"
 	"github.com/spiral/roadrunner/v2/plugins/websockets/connection"
-	"github.com/spiral/roadrunner/v2/utils"
 )
 
 type WorkersPool struct {
 	storage     map[string]pubsub.PubSub
 	connections *sync.Map
 	resPool     sync.Pool
-	bPool       sync.Pool
 	log         logger.Logger
 
 	queue chan *message.Message
@@ -36,9 +33,6 @@ func NewWorkersPool(pubsubs map[string]pubsub.PubSub, connections *sync.Map, log
 	wp.resPool.New = func() interface{} {
 		return make(map[string]struct{}, 10)
 	}
-	wp.bPool.New = func() interface{} {
-		return new(bytes.Buffer)
-	}
 
 	// start 10 workers
 	for i := 0; i < 50; i++ {
@@ -73,15 +67,6 @@ func (wp *WorkersPool) get() map[string]struct{} {
 	return wp.resPool.Get().(map[string]struct{})
 }
 
-func (wp *WorkersPool) putBytes(b *bytes.Buffer) {
-	b.Reset()
-	wp.bPool.Put(b)
-}
-
-func (wp *WorkersPool) getBytes() *bytes.Buffer {
-	return wp.bPool.Get().(*bytes.Buffer)
-}
-
 func (wp *WorkersPool) do() { //nolint:gocognit
 	go func() {
 		for {
@@ -94,29 +79,27 @@ func (wp *WorkersPool) do() { //nolint:gocognit
 				if msg == nil {
 					continue
 				}
-				if msg.TopicsLength() == 0 {
+				if len(msg.GetTopics()) == 0 {
 					continue
 				}
 
-				br, ok := wp.storage[utils.AsString(msg.Broker())]
+				br, ok := wp.storage[msg.Broker]
 				if !ok {
-					wp.log.Warn("no such broker", "requested", utils.AsString(msg.Broker()), "available", wp.storage)
+					wp.log.Warn("no such broker", "requested", msg.GetBroker(), "available", wp.storage)
 					continue
 				}
 
 				res := wp.get()
-				bb := wp.getBytes()
 
-				for i := 0; i < msg.TopicsLength(); i++ {
+				for i := 0; i < len(msg.GetTopics()); i++ {
 					// get connections for the particular topic
-					br.Connections(utils.AsString(msg.Topics(i)), res)
+					br.Connections(msg.GetTopics()[i], res)
 				}
 
 				if len(res) == 0 {
-					for i := 0; i < msg.TopicsLength(); i++ {
-						wp.log.Info("no such topic", "topic", utils.AsString(msg.Topics(i)))
+					for i := 0; i < len(msg.GetTopics()); i++ {
+						wp.log.Info("no such topic", "topic", msg.GetTopics()[i])
 					}
-					wp.putBytes(bb)
 					wp.put(res)
 					continue
 				}
@@ -124,8 +107,8 @@ func (wp *WorkersPool) do() { //nolint:gocognit
 				for i := range res {
 					c, ok := wp.connections.Load(i)
 					if !ok {
-						for i := 0; i < msg.TopicsLength(); i++ {
-							wp.log.Warn("the user disconnected connection before the message being written to it", "broker", utils.AsString(msg.Broker()), "topics", utils.AsString(msg.Topics(i)))
+						for i := 0; i < len(msg.GetTopics()); i++ {
+							wp.log.Warn("the user disconnected connection before the message being written to it", "broker", msg.GetBroker(), "topics", msg.GetTopics()[i])
 						}
 						continue
 					}
@@ -133,20 +116,15 @@ func (wp *WorkersPool) do() { //nolint:gocognit
 					conn := c.(*connection.Connection)
 
 					// put data into the bytes buffer
-					for i := 0; i < msg.PayloadLength(); i++ {
-						bb.WriteByte(byte(msg.Payload(i)))
-					}
-					err := conn.Write(websocket.BinaryMessage, bb.Bytes())
+					err := conn.Write(websocket.BinaryMessage, msg.GetPayload())
 					if err != nil {
-						for i := 0; i < msg.TopicsLength(); i++ {
-							wp.log.Error("error sending payload over the connection", "error", err, "broker", utils.AsString(msg.Broker()), "topics", utils.AsString(msg.Topics(i)))
+						for i := 0; i < len(msg.GetTopics()); i++ {
+							wp.log.Error("error sending payload over the connection", "error", err, "broker", msg.GetBroker(), "topics", msg.GetTopics()[i])
 						}
 						continue
 					}
 				}
 
-				// put bytes buffer back
-				wp.putBytes(bb)
 				// put map with results back
 				wp.put(res)
 			case <-wp.exit:
diff --git a/plugins/websockets/rpc.go b/plugins/websockets/rpc.go
index 6c2cacb4..ef44884a 100644
--- a/plugins/websockets/rpc.go
+++ b/plugins/websockets/rpc.go
@@ -1,10 +1,10 @@
 package websockets
 
 import (
-	flatbuffers "github.com/google/flatbuffers/go"
 	"github.com/spiral/errors"
 	"github.com/spiral/roadrunner/v2/pkg/pubsub/message"
 	"github.com/spiral/roadrunner/v2/plugins/logger"
+	"google.golang.org/protobuf/proto"
 )
 
 // rpc collectors struct
@@ -13,39 +13,32 @@ type rpc struct {
 	log    logger.Logger
 }
 
-// Publish ... msg is a flatbuffers decoded payload
+// Publish ... msg is a proto decoded payload
 // see: pkg/pubsub/message.fbs
-func (r *rpc) Publish(msg []byte, ok *bool) error {
+func (r *rpc) Publish(in *message.Messages, ok *bool) error {
 	const op = errors.Op("broadcast_publish")
-	r.log.Debug("message published")
 
 	// just return in case of nil message
-	if msg == nil {
+	if in == nil {
 		*ok = true
 		return nil
 	}
 
-	fbsMsg := message.GetRootAsMessages(msg, 0)
-	tmpMsg := &message.Message{}
+	r.log.Debug("message published", "msg", in.Messages)
 
-	b := flatbuffers.NewBuilder(100)
+	msgLen := len(in.GetMessages())
 
-	for i := 0; i < fbsMsg.MessagesLength(); i++ {
-		// init a message
-		fbsMsg.Messages(tmpMsg, i)
-
-		// overhead HERE
-		orig := serializeMsg(b, tmpMsg)
-		bb := make([]byte, len(orig))
-		copy(bb, orig)
+	for i := 0; i < msgLen; i++ {
+		bb, err := proto.Marshal(in.GetMessages()[i])
+		if err != nil {
+			return errors.E(op, err)
+		}
 
-		err := r.plugin.Publish(bb)
+		err = r.plugin.Publish(bb)
 		if err != nil {
 			*ok = false
-			b.Reset()
 			return errors.E(op, err)
 		}
-		b.Reset()
 	}
 
 	*ok = true
@@ -54,68 +47,28 @@ func (r *rpc) Publish(msg []byte, ok *bool) error {
 
 // PublishAsync ...
 // see: pkg/pubsub/message.fbs
-func (r *rpc) PublishAsync(msg []byte, ok *bool) error {
-	r.log.Debug("message published", "msg", msg)
+func (r *rpc) PublishAsync(in *message.Messages, ok *bool) error {
+	const op = errors.Op("publish_async")
 
 	// just return in case of nil message
-	if msg == nil {
+	if in == nil {
 		*ok = true
 		return nil
 	}
 
-	fbsMsg := message.GetRootAsMessages(msg, 0)
-	tmpMsg := &message.Message{}
-
-	b := flatbuffers.NewBuilder(100)
+	r.log.Debug("message published", "msg", in.Messages)
 
-	for i := 0; i < fbsMsg.MessagesLength(); i++ {
-		// init a message
-		fbsMsg.Messages(tmpMsg, i)
+	msgLen := len(in.GetMessages())
 
-		// overhead HERE
-		orig := serializeMsg(b, tmpMsg)
-		bb := make([]byte, len(orig))
-		copy(bb, orig)
+	for i := 0; i < msgLen; i++ {
+		bb, err := proto.Marshal(in.GetMessages()[i])
+		if err != nil {
+			return errors.E(op, err)
+		}
 
 		r.plugin.PublishAsync(bb)
-		b.Reset()
 	}
 
 	*ok = true
 	return nil
 }
-
-func serializeMsg(b *flatbuffers.Builder, msg *message.Message) []byte {
-	cmdOff := b.CreateByteString(msg.Command())
-	brokerOff := b.CreateByteString(msg.Broker())
-
-	offsets := make([]flatbuffers.UOffsetT, msg.TopicsLength())
-	for j := msg.TopicsLength() - 1; j >= 0; j-- {
-		offsets[j] = b.CreateByteString(msg.Topics(j))
-	}
-
-	message.MessageStartTopicsVector(b, len(offsets))
-
-	for j := len(offsets) - 1; j >= 0; j-- {
-		b.PrependUOffsetT(offsets[j])
-	}
-
-	tOff := b.EndVector(len(offsets))
-	bb := make([]byte, msg.PayloadLength())
-	for i := 0; i < msg.PayloadLength(); i++ {
-		bb[i] = byte(msg.Payload(i))
-	}
-	pOff := b.CreateByteVector(bb)
-
-	message.MessageStart(b)
-
-	message.MessageAddCommand(b, cmdOff)
-	message.MessageAddBroker(b, brokerOff)
-	message.MessageAddTopics(b, tOff)
-	message.MessageAddPayload(b, pOff)
-
-	fOff := message.MessageEnd(b)
-	b.Finish(fOff)
-
-	return b.FinishedBytes()
-}