Cyrene_v2/internal/channel/adapter/wsutil/client.go

package wsutil

import (
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

const (
	// defaultHeartbeatInterval is used when a client enables heartbeat without a
	// configured interval.
	defaultHeartbeatInterval = 30 * time.Second
	// defaultPongTimeout is the read deadline margin added to the heartbeat
	// interval; if the peer sends nothing within interval+pongTimeout the socket
	// is considered dead and the client reconnects.
	defaultPongTimeout = 6 * time.Second
	// writeWait bounds a single control-frame write.
	writeWait = 10 * time.Second
)

// DialFunc returns the gateway URL to connect to. On a reconnect the platform may
// return a resume-augmented URL (splice session/sn via AppendQuery) so the
// gateway can resynchronize the offline message window.
type DialFunc func(ctx context.Context) (string, error)

// Handler carries the platform-specific protocol hooks the Client invokes. A
// nil field disables that hook. The hooks run on the Client's read goroutine
// (OnConnect / OnMessage) or a dedicated heartbeat goroutine (Heartbeat).
type Handler struct {
	// OnConnect is called after the WebSocket connection is established and the
	// application handshake (for example KOOK's hello) succeeds. Returning an
	// error aborts the session and triggers a reconnect with backoff.
	OnConnect func(ctx context.Context, c *Conn) error
	// OnMessage is called for every application text frame received. Returning an
	// error forces a reconnect.
	OnMessage func(ctx context.Context, payload []byte) error
	// Heartbeat is called on the heartbeat ticker (when HeartbeatInterval > 0) to
	// send an application-level keepalive frame. A nil Heartbeat disables it.
	Heartbeat func(ctx context.Context, c *Conn) error
}

// Client is a reusable WebSocket gateway connection manager. It dials a gateway,
// runs the read loop, maintains an optional application-level heartbeat with a
// liveness read deadline, and reconnects with exponential backoff until the
// context is canceled. Only the platform signaling is left to the Handler.
type Client struct {
	dial     DialFunc
	handler  Handler
	dialer   *websocket.Dialer
	header   http.Header
	interval time.Duration
	pongWait time.Duration
	backoff  *Backoff
	log      *slog.Logger

	mu     sync.Mutex
	conn   *Conn
	cancel context.CancelFunc
	closed bool
}

// Option customizes a Client at construction time.
type Option func(*Client)

// WithDialer injects the WebSocket dialer used for the handshake. Tests pass a
// dialer whose net dial is redirected to a fake server.
func WithDialer(d *websocket.Dialer) Option {
	return func(c *Client) { c.dialer = d }
}

// WithHeader injects extra HTTP headers on the handshake (for example an
// Authorization: Bearer token).
func WithHeader(h http.Header) Option {
	return func(c *Client) { c.header = h }
}

// WithHeartbeat enables an application-level heartbeat at the given interval.
// A nil interval disables the heartbeat ticker (default off).
func WithHeartbeat(interval time.Duration) Option {
	return func(c *Client) { c.interval = interval }
}

// WithPongTimeout sets the read deadline margin added to the heartbeat interval.
// A non-positive value uses the default (6s).
func WithPongTimeout(d time.Duration) Option {
	return func(c *Client) { c.pongWait = d }
}

// WithBackoff injects the reconnect backoff scheduler.
func WithBackoff(b *Backoff) Option {
	return func(c *Client) { c.backoff = b }
}

// WithLogger injects the logger used by the client. A nil logger stays silent.
func WithLogger(l *slog.Logger) Option {
	return func(c *Client) {
		if l == nil {
			c.log = logger.New(logger.Options{Out: io.Discard})
		} else {
			c.log = l
		}
	}
}

// New builds a Client. dial must be non-nil.
func New(dial DialFunc, h Handler, opts ...Option) *Client {
	if dial == nil {
		dial = func(context.Context) (string, error) { return "", nil }
	}
	c := &Client{
		dial:     dial,
		handler:  h,
		dialer:   &websocket.Dialer{HandshakeTimeout: 15 * time.Second},
		pongWait: defaultPongTimeout,
		backoff:  NewBackoff(2*time.Second, 60*time.Second),
		log:      logger.New(logger.Options{Out: io.Discard}),
	}
	for _, opt := range opts {
		opt(c)
	}
	return c
}

// Run executes the connect/reconnect loop until ctx is canceled. It returns
// nil when ctx is canceled and a non-nil error on a non-context termination. It
// is intended to be run in its own goroutine by the adapter's Connect.
func (c *Client) Run(ctx context.Context) error {
	for {
		if ctx.Err() != nil {
			return ctx.Err()
		}
		url, err := c.dial(ctx)
		if err != nil {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			c.log.Warn("ws gateway dial failed", "err", err)
			if !c.sleepBackoff(ctx) {
				return ctx.Err()
			}
			continue
		}
		if err := c.connectOnce(ctx, url); err != nil {
			c.log.Warn("ws session ended", "err", err)
		}
		if ctx.Err() != nil {
			return ctx.Err()
		}
	}
}

// connectOnce dials one URL and runs the connect hook plus the read loop until
// the connection drops, the read loop errors, or ctx is canceled.
func (c *Client) connectOnce(ctx context.Context, url string) error {
	if c.header == nil {
		c.header = http.Header{}
	}
	ws, _, err := c.dialer.DialContext(ctx, url, c.header)
	if err != nil {
		if ctx.Err() != nil {
			return ctx.Err()
		}
		c.log.Warn("ws dial failed", "err", err)
		c.sleepBackoff(ctx)
		return err
	}
	conn := &Conn{ws: ws}
	defer conn.Close()

	// A watcher closes the connection when ctx ends so a pending read (including
	// the OnConnect handshake) unblocks and Run can exit promptly.
	stop := make(chan struct{})
	defer close(stop)
	go c.watchCtx(ctx, conn, stop)

	if c.handler.OnConnect != nil {
		if cerr := c.handler.OnConnect(ctx, conn); cerr != nil {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			c.log.Warn("ws connect hook failed", "err", cerr)
			c.sleepBackoff(ctx)
			return cerr
		}
	}

	c.setConn(conn)
	c.backoff.Reset()
	c.runReadLoop(ctx, conn)
	c.setConn(nil)
	return nil
}

// watchCtx closes the connection when ctx is canceled so a pending read
// unblocks, letting Run exit promptly.
func (c *Client) watchCtx(ctx context.Context, conn *Conn, stop <-chan struct{}) {
	select {
	case <-ctx.Done():
		_ = conn.Close()
	case <-stop:
	}
}

// runReadLoop consumes text frames until the read errors, the platform handler
// returns an error, or ctx is canceled. When heartbeat is enabled it also runs a
// heartbeat ticker and refreshes a read deadline so a silent-but-live socket
// stays open while a dead one is detected.
func (c *Client) runReadLoop(ctx context.Context, conn *Conn) {
	hbCtx, hbCancel := context.WithCancel(ctx)
	defer hbCancel()
	if c.handler.Heartbeat != nil && c.interval > 0 {
		go c.heartbeatLoop(hbCtx, conn)
	}

	// A liveness read deadline is only meaningful when an application heartbeat
	// keeps the socket active; without one, an idle-but-open connection must stay
	// open indefinitely.
	deadline := time.Duration(0)
	if c.handler.Heartbeat != nil && c.interval > 0 {
		deadline = c.interval + c.pongWait
	}
	for {
		if ctx.Err() != nil {
			return
		}
		if deadline > 0 {
			_ = conn.SetReadDeadline(time.Now().Add(deadline))
		} else {
			_ = conn.SetReadDeadline(time.Time{})
		}
		payload, err := conn.ReadMessage()
		if err != nil {
			return
		}
		if c.handler.OnMessage != nil {
			if err := c.handler.OnMessage(ctx, payload); err != nil {
				if ctx.Err() == nil {
					c.log.Warn("ws on-message failed", "err", err)
				}
				return
			}
		}
	}
}

// heartbeatLoop sends an application-level keepalive on every interval. A
// failing heartbeat closes the connection to force a reconnect.
func (c *Client) heartbeatLoop(ctx context.Context, conn *Conn) {
	t := time.NewTicker(c.interval)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			if err := c.handler.Heartbeat(ctx, conn); err != nil {
				c.log.Warn("ws heartbeat failed", "err", err)
				_ = conn.Close()
				return
			}
		}
	}
}

// sleepBackoff sleeps the current backoff delay or until ctx is canceled,
// returning false (and leaving the connection for Run to exit) when ctx fires.
func (c *Client) sleepBackoff(ctx context.Context) bool {
	d := c.backoff.Next()
	t := time.NewTimer(d)
	defer t.Stop()
	select {
	case <-ctx.Done():
		return false
	case <-t.C:
		return true
	}
}

// setConn installs the live connection so Close can tear it down.
func (c *Client) setConn(conn *Conn) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.conn = conn
}

// Close tears down the client, closing the live connection and its context.
func (c *Client) Close() error {
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return nil
	}
	c.closed = true
	conn := c.conn
	cancel := c.cancel
	c.mu.Unlock()
	if conn != nil {
		_ = conn.Close()
	}
	if cancel != nil {
		cancel()
	}
	return nil
}

// BindCancel associates a context cancel func so Close can stop a running Run.
func (c *Client) BindCancel(cancel context.CancelFunc) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.cancel = cancel
}

// Conn is a thin, concurrency-safe wrapper over one *websocket.Conn. It
// serializes writes (gorilla allows one reader plus one writer) and exposes only
// the JSON text-frame surface the gateway clients need.
type Conn struct {
	ws      *websocket.Conn
	writeMu sync.Mutex
}

// ReadMessage reads the next complete message, returning its text payload. It
// honors any read deadline set via SetReadDeadline.
func (c *Conn) ReadMessage() ([]byte, error) {
	_, p, err := c.ws.ReadMessage()
	if err != nil {
		return nil, err
	}
	return p, nil
}

// WriteMessage writes a text frame with the given payload.
func (c *Conn) WriteMessage(payload []byte) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	return c.ws.WriteMessage(websocket.TextMessage, payload)
}

// WriteJSON marshals v and writes it as a text frame.
func (c *Conn) WriteJSON(v any) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return c.WriteMessage(data)
}

// SetReadDeadline bounds the next ReadMessage by the given absolute time.
func (c *Conn) SetReadDeadline(t time.Time) error { return c.ws.SetReadDeadline(t) }

// Close closes the underlying connection exactly once.
func (c *Conn) Close() error { return c.ws.Close() }