149 lines
3.3 KiB
Go
149 lines
3.3 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"strings"
|
|
"time"
|
|
|
|
"nhooyr.io/websocket"
|
|
)
|
|
|
|
// 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
|
|
handlers map[string]MessageHandler // topic prefix -> handler
|
|
}
|
|
|
|
// NewIngestor creates an Ingestor with registered stream handlers.
|
|
func NewIngestor(cfg Config, handlers map[string]MessageHandler) *Ingestor {
|
|
return &Ingestor{
|
|
cfg: cfg,
|
|
handlers: handlers,
|
|
}
|
|
}
|
|
|
|
// Run connects to the WebSocket and routes messages until ctx is cancelled.
|
|
func (ing *Ingestor) Run(ctx context.Context) {
|
|
log.Println("[ingestor] Starting multi-stream WebSocket ingestor...")
|
|
defer log.Println("[ingestor] Ingestor stopped.")
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
default:
|
|
}
|
|
|
|
if err := ing.connectAndConsume(ctx); err != nil {
|
|
if ctx.Err() != nil {
|
|
return // Context cancelled
|
|
}
|
|
log.Printf("[ingestor] Connection error: %v. Reconnecting in 5s...", err)
|
|
select {
|
|
case <-time.After(5 * time.Second):
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (ing *Ingestor) connectAndConsume(ctx context.Context) error {
|
|
url := ing.cfg.WebSocketURL
|
|
log.Printf("[ingestor] Connecting to %s ...", url)
|
|
|
|
conn, _, err := websocket.Dial(ctx, url, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("dial: %w", err)
|
|
}
|
|
defer conn.CloseNow()
|
|
|
|
// 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")
|
|
}
|
|
|
|
subMsg := map[string]interface{}{
|
|
"op": "subscribe",
|
|
"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 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 {
|
|
select {
|
|
case <-ctx.Done():
|
|
conn.Close(websocket.StatusNormalClosure, "shutting down")
|
|
return nil
|
|
default:
|
|
}
|
|
|
|
_, data, err := conn.Read(ctx)
|
|
if err != nil {
|
|
return fmt.Errorf("read: %w", err)
|
|
}
|
|
|
|
ing.handleMessage(data)
|
|
}
|
|
}
|
|
|
|
type topicHeader struct {
|
|
Topic string `json:"topic"`
|
|
}
|
|
|
|
func (ing *Ingestor) handleMessage(data []byte) {
|
|
var header topicHeader
|
|
if err := json.Unmarshal(data, &header); err != nil || header.Topic == "" {
|
|
return // pong, sub ack, or malformed message
|
|
}
|
|
|
|
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
|
|
}
|