exchange-monitor-go/exchange/connector.go

package exchange

import (
	"log"
	"strconv"
	"time"

	"github.com/gorilla/websocket"
)

// PriceConnector is a reusable WebSocket reconnector with backoff and keepalive.
type PriceConnector struct {
	URL           string
	Name          string
	ReadTimeout   time.Duration
	ReconnectBase time.Duration
	PingInterval  time.Duration // 0 = no client-side pings

	// TextPing: if true, sends text message "ping" instead of WebSocket PingMessage frame.
	// Used by Bitget v2 WS which expects a JSON/text "ping" string.
	TextPing bool

	OnConnect func()
	OnMessage func([]byte)
	OnError   func(error)

	conn *websocket.Conn
	done chan struct{}
}

func NewPriceConnector(url, name string, readTimeout, reconnectBase time.Duration) *PriceConnector {
	return &PriceConnector{
		URL:           url,
		Name:          name,
		ReadTimeout:   readTimeout,
		ReconnectBase: reconnectBase,
		done:          make(chan struct{}),
	}
}

func (pc *PriceConnector) Run() error {
	backoff := time.Second
	attempt := 0

	for {
		select {
		case <-pc.done:
			return nil
		default:
		}

		log.Printf("[%s WS] Connecting... (attempt %d)", pc.Name, attempt+1)

		c, _, err := websocket.DefaultDialer.Dial(pc.URL, nil)
		if err != nil {
			log.Printf("[%s WS] Dial error: %v (retry in %v)", pc.Name, err, backoff)
			if pc.OnError != nil {
				pc.OnError(err)
			}
			time.Sleep(backoff)
			backoff *= 2
			if backoff > 30*time.Second {
				backoff = 30 * time.Second
			}
			attempt++
			continue
		}

			backoff = time.Second
			attempt = 0
			pc.conn = c

			// ── Ping/Pong keepalive ──
			c.SetReadDeadline(time.Now().Add(pc.ReadTimeout))

			// Respond to server pings with pongs (extend deadline)
			c.SetPongHandler(func(appData string) error {
				c.SetReadDeadline(time.Now().Add(pc.ReadTimeout))
				return nil
			})

			// Handle server-initiated PING frames (used by Bitget etc.)
			// Extend read deadline so the connection doesn't timeout
			c.SetPingHandler(func(appData string) error {
				c.SetReadDeadline(time.Now().Add(pc.ReadTimeout))
				return nil
			})

			// Client-side ping sender
		pingStop := make(chan struct{})
		if pc.PingInterval > 0 {
			go func() {
				ticker := time.NewTicker(pc.PingInterval)
				defer ticker.Stop()
				for {
					select {
					case <-ticker.C:
						if pc.TextPing {
							// Bitget v2 expects a TEXT message "ping", not WebSocket PingMessage frame
							if err := c.WriteMessage(websocket.TextMessage, []byte("ping")); err != nil {
								return
							}
						} else {
							if err := c.WriteMessage(websocket.PingMessage, nil); err != nil {
								return
							}
						}
					case <-pingStop:
						return
					case <-pc.done:
						return
					}
				}
			}()
		}

		if pc.OnConnect != nil {
			pc.OnConnect()
		}

		// Read loop
	readLoop:
		for {
			_, msg, err := c.ReadMessage()
			if err != nil {
				log.Printf("[%s WS] Read error: %v", pc.Name, err)
				break readLoop
			}
			c.SetReadDeadline(time.Now().Add(pc.ReadTimeout))

			if pc.OnMessage != nil {
				pc.OnMessage(msg)
			}
		}

		close(pingStop)
		c.Close()
	}
}

func (pc *PriceConnector) Stop() {
	close(pc.done)
	if pc.conn != nil {
		pc.conn.Close()
	}
}

// Done returns a channel that's closed when the connector is stopped.
func (pc *PriceConnector) Done() <-chan struct{} {
	return pc.done
}

// SendJSON sends a JSON message over the WebSocket.
func (pc *PriceConnector) SendJSON(v interface{}) error {
	if pc.conn == nil {
		return nil
	}
	return pc.conn.WriteJSON(v)
}

// Helper: parse float from string
// B#8: use strconv.ParseFloat instead of fmt.Sscanf
func parseFloat(s string) float64 {
	f, err := strconv.ParseFloat(s, 64)
	if err != nil {
		return 0
	}
	return f
}

// Helper: convert "BTCUSDT" with suffix "USDT" to "BTC"
func symbolToCoin(symbol, suffix string) string {
	if len(symbol) <= len(suffix) {
		return ""
	}
	if symbol[len(symbol)-len(suffix):] == suffix {
		return symbol[:len(symbol)-len(suffix)]
	}
	return ""
}