Added klines, liquidations, tickers and trades to be recorded. Bundled the stats as cli argument.

websocket.go

@@ -5,33 +5,30 @@ import (
 	"encoding/json"
 	"fmt"
 	"log"
-	"strconv"
+	"strings"
 	"time"
 
 	"nhooyr.io/websocket"
 )
 
-// Ingestor connects to the Bybit V5 WebSocket, parses trade messages,
-// sends ticks to the writer channel, and feeds ticks to the aggregator.
+// Ingestor connects to the Bybit V5 WebSocket, receives raw JSON frames,
+// extracts the topic prefix, and routes messages to registered MessageHandlers.
 type Ingestor struct {
-	cfg        Config
-	tickCh     chan<- Tick
-	aggregator *Aggregator
+	cfg      Config
+	handlers map[string]MessageHandler // topic prefix -> handler
 }
 
-// NewIngestor creates an Ingestor wired to the tick channel and aggregator.
-func NewIngestor(cfg Config, tickCh chan<- Tick, agg *Aggregator) *Ingestor {
+// NewIngestor creates an Ingestor with registered stream handlers.
+func NewIngestor(cfg Config, handlers map[string]MessageHandler) *Ingestor {
 	return &Ingestor{
-		cfg:        cfg,
-		tickCh:     tickCh,
-		aggregator: agg,
+		cfg:      cfg,
+		handlers: handlers,
 	}
 }
 
-// Run connects to the WebSocket and processes messages until ctx is cancelled.
-// It automatically reconnects on connection failures.
+// Run connects to the WebSocket and routes messages until ctx is cancelled.
 func (ing *Ingestor) Run(ctx context.Context) {
-	log.Println("[ingestor] Starting WebSocket ingestor...")
+	log.Println("[ingestor] Starting multi-stream WebSocket ingestor...")
 	defer log.Println("[ingestor] Ingestor stopped.")
 
 	for {
@@ -43,7 +40,7 @@ func (ing *Ingestor) Run(ctx context.Context) {
 
 		if err := ing.connectAndConsume(ctx); err != nil {
 			if ctx.Err() != nil {
-				return // Context cancelled, clean exit
+				return // Context cancelled
 			}
 			log.Printf("[ingestor] Connection error: %v. Reconnecting in 5s...", err)
 			select {
@@ -65,21 +62,47 @@ func (ing *Ingestor) connectAndConsume(ctx context.Context) error {
 	}
 	defer conn.CloseNow()
 
-	// Set a generous read limit for large trade batches (up to 1024 trades per message)
-	conn.SetReadLimit(1 << 20) // 1 MB
+	// 2 MB read limit for large orderbook snapshots or trade batches
+	conn.SetReadLimit(2 << 20)
+
+	// Collect all topics from registered handlers
+	var allTopics []string
+	for _, h := range ing.handlers {
+		allTopics = append(allTopics, h.Topics()...)
+	}
+
+	if len(allTopics) == 0 {
+		return fmt.Errorf("no topics registered to subscribe")
+	}
 
-	// Subscribe to public trades
-	topic := fmt.Sprintf("publicTrade.%s", ing.cfg.Symbol)
 	subMsg := map[string]interface{}{
 		"op":   "subscribe",
-		"args": []string{topic},
+		"args": allTopics,
 	}
 	subJSON, _ := json.Marshal(subMsg)
 
 	if err := conn.Write(ctx, websocket.MessageText, subJSON); err != nil {
 		return fmt.Errorf("subscribe: %w", err)
 	}
-	log.Printf("[ingestor] Subscribed to %s", topic)
+	log.Printf("[ingestor] Subscribed to topics: %v", allTopics)
+
+	// Ping ticker to keep Bybit connection alive
+	pingTicker := time.NewTicker(20 * time.Second)
+	defer pingTicker.Stop()
+
+	go func() {
+		for {
+			select {
+			case <-pingTicker.C:
+				pingMsg := []byte(`{"op":"ping"}`)
+				if err := conn.Write(ctx, websocket.MessageText, pingMsg); err != nil {
+					return
+				}
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
 
 	// Read loop
 	for {
@@ -99,64 +122,27 @@ func (ing *Ingestor) connectAndConsume(ctx context.Context) error {
 	}
 }
 
+type topicHeader struct {
+	Topic string `json:"topic"`
+}
+
 func (ing *Ingestor) handleMessage(data []byte) {
-	var msg BybitWSMessage
-
-	if err := json.Unmarshal(data, &msg); err != nil {
-		// subscription confirmations, pings, etc.
-		return
+	var header topicHeader
+	if err := json.Unmarshal(data, &header); err != nil || header.Topic == "" {
+		return // pong, sub ack, or malformed message
 	}
 
-	if len(msg.Data) == 0 {
-		return
-	}
-
-	recvTS := time.Now().UnixMilli()
-
-	for _, raw := range msg.Data {
-
-		price, err := strconv.ParseFloat(raw.P, 64)
-		if err != nil {
-			log.Printf("[ingestor] bad price %q: %v", raw.P, err)
-			continue
-		}
-
-		volume, err := strconv.ParseFloat(raw.V, 64)
-		if err != nil {
-			log.Printf("[ingestor] bad volume %q: %v", raw.V, err)
-			continue
-		}
-
-		tick := Tick{
-			// IDs
-			TradeID: raw.I,
-			Seq:     raw.Seq,
-
-			// Timing
-			TradeTS:   raw.T,
-			MessageTS: msg.TS,
-			RecvTS:    recvTS,
-
-			// Trade data
-			Symbol: raw.S,
-			Side:   raw.SD,
-			Price:  price,
-			Volume: volume,
-
-			// Metadata
-			TickDir:    raw.L,
-			BlockTrade: raw.BT,
-			RPI:        raw.RPI,
-		}
-
-		// Feed feature generator
-		ing.aggregator.ProcessTick(tick)
-
-		// Feed writer
-		select {
-		case ing.tickCh <- tick:
-		default:
-			log.Println("[ingestor] WARNING: tick channel full, dropping tick")
-		}
+	prefix := extractTopicPrefix(header.Topic)
+	if handler, ok := ing.handlers[prefix]; ok {
+		handler.HandleMessage(data)
 	}
 }
+
+// extractTopicPrefix extracts the topic family name (e.g. "publicTrade", "tickers", "kline", "orderbook", "allLiquidation").
+func extractTopicPrefix(topic string) string {
+	parts := strings.Split(topic, ".")
+	if len(parts) > 0 {
+		return parts[0]
+	}
+	return topic
+}