Add Retry and ACK for Broadcast Flow (#117)

* add retry and ack for broadcast flow

* fix lint

* update unit test

* update unit test

v2/internal/server/broadcast_test.go

@@ -0,0 +1,59 @@
+package server
+
+import (
+	"log"
+	"net/http"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/KyberNetwork/tradinglib/pkg/httpsign"
+	"github.com/gorilla/websocket"
+)
+
+func TestSubscribe(t *testing.T) {
+	t.Skip("set env to run test")
+	wsURL := os.Getenv("URL") // Replace with actual WebSocket URL
+	key := os.Getenv("KEY")
+	secret := os.Getenv("SECRET")
+
+	request, err := http.NewRequest(http.MethodGet, wsURL+"?id=cscv9ubrk77vgbjftu5g", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = httpsign.Sign(request, key, []byte(secret))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// local test
+	// request, err := http.NewRequest("GET", "ws://localhost:8082/eventlogws", nil)
+	// wsURL := "ws://localhost:8082/eventlogws"
+
+	// Establish a connection with the WebSocket server
+	conn, _, err := websocket.DefaultDialer.Dial(wsURL+"?id=cscv9ubrk77vgbjftu5g", request.Header)
+	if err != nil {
+		t.Fatal("Error connecting to WebSocket server:", err)
+	}
+	defer conn.Close()
+
+	go func() {
+		ticker := time.NewTicker(time.Minute)
+		defer ticker.Stop()
+		for ; ; <-ticker.C {
+			if err = conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+				log.Println("Ping error:", err)
+				return
+			}
+		}
+	}()
+
+	for {
+		// cscv9ubrk77vgbjftu5g
+		_, message, err := conn.ReadMessage()
+		if err != nil {
+			t.Fatal("Error reading from websocket:", err)
+		}
+		log.Printf("Received: %s", message)
+	}
+}

v2/internal/worker/broadcast_client.go

@@ -5,10 +5,11 @@ import (
 	"encoding/json"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/IBM/sarama"
-	"github.com/KyberNetwork/tradelogs/pkg/storage"
 	"github.com/KyberNetwork/tradelogs/v2/pkg/kafka"
+	"github.com/KyberNetwork/tradelogs/v2/pkg/storage/tradelogs/types"
 	"github.com/gorilla/websocket"
 	"go.uber.org/zap"
 )
@@ -28,34 +29,48 @@ func (c *Client) run(ctx context.Context, cfg *kafka.Config, topic string) error
 	}
 
 	tradeLogsChan := make(chan *sarama.ConsumerMessage, 100)
+	ackChan := make(chan bool)
 
 	go func() {
-		err = consumer.Consume(ctx, c.l, topic, tradeLogsChan)
+		err = consumer.Consume(ctx, c.l, topic, tradeLogsChan, ackChan)
 		if err != nil {
 			panic(fmt.Errorf("failed to consume trade logs: %w", err))
 		}
 	}()
 
-	go func() {
-		for msg := range tradeLogsChan {
-			c.broadcast(msg)
-		}
-	}()
+	retry := 3
+	delay := time.Second * 5
+	i := 0
 
-	return nil
+	for msg := range tradeLogsChan {
+		for i = 0; i < retry; i++ {
+			err = c.broadcast(msg)
+			if err != nil {
+				time.Sleep(delay)
+				continue
+			}
+			ackChan <- true
+			break
+		}
+		if i == retry {
+			ackChan <- false
+			break
+		}
+	}
+	return fmt.Errorf("broadcast trade logs to client %s failed", c.id)
 }
 
-func (c *Client) broadcast(msg *sarama.ConsumerMessage) {
+func (c *Client) broadcast(msg *sarama.ConsumerMessage) error {
 	var newMsg kafka.Message
 	err := json.Unmarshal(msg.Value, &newMsg)
 	if err != nil {
 		c.l.Errorw("error when unmarshal message", "err", err, "msg", string(msg.Value))
-		return
+		return err
 	}
 	dataBytes, err := json.Marshal(newMsg.Data)
 	if err != nil {
 		c.l.Errorw("error when marshal message data", "err", err, "data", newMsg.Data)
-		return
+		return err
 	}
 
 	switch newMsg.Type {
@@ -64,32 +79,35 @@ func (c *Client) broadcast(msg *sarama.ConsumerMessage) {
 		err = json.Unmarshal(dataBytes, &blocks)
 		if err != nil {
 			c.l.Errorw("error when unmarshal reverted blocks", "err", err, "data", string(dataBytes))
-			return
+			return err
 		}
 		newMsg.Data = blocks
 		if err = c.ws.WriteJSON(newMsg); err != nil {
 			c.l.Errorw("error when send msg", "err", err)
+			return err
 		}
 		c.l.Infow("broadcast revert message", "message", newMsg)
 
 	case kafka.MessageTypeTradeLog:
-		var tradelog storage.TradeLog
-		err = json.Unmarshal(dataBytes, &tradelog)
+		var tradeLog types.TradeLog
+		err = json.Unmarshal(dataBytes, &tradeLog)
 		if err != nil {
 			c.l.Errorw("error when unmarshal trade log", "err", err, "data", string(dataBytes))
-			return
+			return err
 		}
-		newMsg.Data = tradelog
-		if c.match(tradelog) {
+		newMsg.Data = tradeLog
+		if c.match(tradeLog) {
 			if err = c.ws.WriteJSON(newMsg); err != nil {
 				c.l.Errorw("error when send msg", "err", err)
+				return err
 			}
+			c.l.Infow("broadcast trade log message", "message", newMsg)
 		}
-		c.l.Infow("broadcast trade log message", "message", newMsg)
 	}
+	return nil
 }
 
-func (c *Client) match(log storage.TradeLog) bool {
+func (c *Client) match(log types.TradeLog) bool {
 	params := c.params
 	if len(params.EventHash) != 0 && !strings.EqualFold(params.EventHash, log.EventHash) {
 		return false

v2/internal/worker/broadcaster.go

@@ -65,11 +65,14 @@ func (b *Broadcaster) NewConn(req RegisterRequest, conn *websocket.Conn) error {
 
 	b.addClient(conn, req.ID, req)
 
-	err := b.clients[req.ID].run(ctx, b.config, b.topic)
-	if err != nil {
-		b.removeClient(cancel, conn, req.ID)
-		return fmt.Errorf("cannot run client: %w", err)
-	}
+	go func() {
+		err := b.clients[req.ID].run(ctx, b.config, b.topic)
+		if err != nil {
+			b.l.Errorw("error when run client", "id", req.ID, "err", err)
+			b.removeClient(cancel, conn, req.ID)
+		}
+	}()
+
 	return nil
 }
 

v2/pkg/kafka/consumer.go

@@ -32,9 +32,9 @@ func NewConsumer(config *Config, consumerGroup string) (*SaramaConsumer, error)
 	}, nil
 }
 
-func (c *SaramaConsumer) Consume(ctx context.Context, l *zap.SugaredLogger, topic string, ch chan<- *sarama.ConsumerMessage) error {
-	messageHandler := newConsumerGroupHandler(ch)
-	defer close(ch)
+func (c *SaramaConsumer) Consume(ctx context.Context, l *zap.SugaredLogger, topic string, msgCh chan<- *sarama.ConsumerMessage, ackCh chan bool) error {
+	messageHandler := newConsumerGroupHandler(msgCh, ackCh)
+	defer close(msgCh)
 	for {
 		select {
 		case <-ctx.Done():
@@ -52,11 +52,13 @@ func (c *SaramaConsumer) Consume(ctx context.Context, l *zap.SugaredLogger, topi
 
 type consumerGroupHandler struct {
 	msgChan chan<- *sarama.ConsumerMessage
+	ackChan chan bool
 }
 
-func newConsumerGroupHandler(msgChan chan<- *sarama.ConsumerMessage) *consumerGroupHandler {
+func newConsumerGroupHandler(msgChan chan<- *sarama.ConsumerMessage, ackCh chan bool) *consumerGroupHandler {
 	return &consumerGroupHandler{
 		msgChan: msgChan,
+		ackChan: ackCh,
 	}
 }
 
@@ -70,8 +72,13 @@ func (h *consumerGroupHandler) ConsumeClaim(session sarama.ConsumerGroupSession,
 		// Push the message to the channel
 		h.msgChan <- msg
 
-		// Acknowledge the message by marking it as processed (commit the offset)
-		session.MarkMessage(msg, "")
+		ack := <-h.ackChan
+		if ack {
+			// Acknowledge the message by marking it as processed (commit the offset)
+			session.MarkMessage(msg, "")
+		} else {
+			break // stop consume when fail to broadcast message
+		}
 	}
 	return nil
 }