proto-socket/go/ws_client.go

package toki_socket

import (
	"context"
	"crypto/tls"
	"fmt"
	"net/http"
	"sync"
	"time"

	"google.golang.org/protobuf/proto"
	"nhooyr.io/websocket"

	"toki-labs.com/toki_socket/go/packets"
)

const (
	DefaultHeartbeatIntervalSec = 30
	DefaultHeartbeatWaitSec     = 10
)

type WsClient struct {
	Communicator
	conn                *websocket.Conn
	writeMu             sync.Mutex // Defensive if WritePacket is called outside the queued writer.
	readCtx             context.Context
	cancelRead          context.CancelFunc
	heartbeatInterval   time.Duration
	heartbeatWait       time.Duration
	waitingHBResponse   bool
	hbTimer             *HeartbeatTimer
	hbMu                sync.Mutex
	disconnectListeners []func(*WsClient)
	disconnectMu        sync.Mutex
	wsCloseOnce         sync.Once
}

func NewWsClient(conn *websocket.Conn, intervalSec, waitSec int, parserMap ParserMap) *WsClient {
	readCtx, cancelRead := context.WithCancel(context.Background())
	c := &WsClient{
		conn:              conn,
		readCtx:           readCtx,
		cancelRead:        cancelRead,
		heartbeatInterval: time.Duration(intervalSec) * time.Second,
		heartbeatWait:     time.Duration(waitSec) * time.Second,
	}
	c.Communicator.Initialize(c, parserMap)
	c.SetWriteErrorHandler(func(error) {
		c.onDisconnected()
	})
	c.AddListener(TypeNameOf(&packets.HeartBeat{}), func(m proto.Message) {
		c.onHeartBeat()
	})
	go c.readLoop()
	c.sendHeartBeat()
	return c
}

func DialWs(ctx context.Context, host string, port int, path string, parserMap ParserMap) (*WsClient, error) {
	return DialWsWithHeartbeat(ctx, host, port, path, DefaultHeartbeatIntervalSec, DefaultHeartbeatWaitSec, parserMap)
}

func DialWsWithHeartbeat(ctx context.Context, host string, port int, path string, intervalSec, waitSec int, parserMap ParserMap) (*WsClient, error) {
	conn, _, err := websocket.Dial(ctx, fmt.Sprintf("ws://%s:%d%s", host, port, path), nil)
	if err != nil {
		return nil, err
	}
	return NewWsClient(conn, intervalSec, waitSec, parserMap), nil
}

func DialWss(ctx context.Context, host string, port int, path string, tlsCfg *tls.Config, parserMap ParserMap) (*WsClient, error) {
	return DialWssWithHeartbeat(ctx, host, port, path, tlsCfg, DefaultHeartbeatIntervalSec, DefaultHeartbeatWaitSec, parserMap)
}

func DialWssWithHeartbeat(ctx context.Context, host string, port int, path string, tlsCfg *tls.Config, intervalSec, waitSec int, parserMap ParserMap) (*WsClient, error) {
	opts := &websocket.DialOptions{}
	if tlsCfg != nil {
		opts.HTTPClient = &http.Client{
			Transport: &http.Transport{TLSClientConfig: tlsCfg},
		}
	}
	conn, _, err := websocket.Dial(ctx, fmt.Sprintf("wss://%s:%d%s", host, port, path), opts)
	if err != nil {
		return nil, err
	}
	return NewWsClient(conn, intervalSec, waitSec, parserMap), nil
}

func (c *WsClient) readLoop() {
	for c.IsAlive() {
		msgType, b, err := c.conn.Read(c.readCtx)
		if err != nil {
			c.onDisconnected()
			return
		}
		if msgType != websocket.MessageBinary {
			continue
		}
		base := &packets.PacketBase{}
		if err := proto.Unmarshal(b, base); err != nil {
			c.onDisconnected()
			return
		}
		c.OnReceivedData(base.GetTypeName(), base.GetData(), base.GetNonce(), base.GetResponseNonce())
		c.sendHeartBeat()
	}
}

func (c *WsClient) WritePacket(base *packets.PacketBase) error {
	b, err := proto.Marshal(base)
	if err != nil {
		return err
	}
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	return c.conn.Write(context.Background(), websocket.MessageBinary, b)
}

func (c *WsClient) AddDisconnectListener(handler func(*WsClient)) {
	c.disconnectMu.Lock()
	defer c.disconnectMu.Unlock()
	c.disconnectListeners = append(c.disconnectListeners, handler)
}

func (c *WsClient) RemoveDisconnectListeners() {
	c.disconnectMu.Lock()
	defer c.disconnectMu.Unlock()
	c.disconnectListeners = nil
}

func (c *WsClient) Close() error {
	var err error
	c.wsCloseOnce.Do(func() {
		c.Communicator.shutdown()
		c.stopHeartbeat()
		c.cancelRead()
		err = c.conn.Close(websocket.StatusNormalClosure, "")
		c.notifyDisconnected()
	})
	return err
}

func (c *WsClient) sendHeartBeat() {
	if !c.IsAlive() || c.heartbeatInterval <= 0 {
		return
	}
	c.hbMu.Lock()
	if c.hbTimer != nil {
		c.hbTimer.Stop()
	}
	c.hbTimer = NewHeartbeatTimer(c.heartbeatInterval, func() {
		if !c.IsAlive() {
			return
		}
		c.hbMu.Lock()
		c.waitingHBResponse = true
		c.hbMu.Unlock()
		_ = c.Send(&packets.HeartBeat{})
		c.hbMu.Lock()
		if c.hbTimer != nil {
			c.hbTimer.Stop()
		}
		c.hbTimer = NewHeartbeatTimer(c.heartbeatWait, func() {
			if c.IsAlive() {
				c.onDisconnected()
			}
		})
		c.hbMu.Unlock()
	})
	c.hbMu.Unlock()
}

func (c *WsClient) onHeartBeat() {
	c.hbMu.Lock()
	if c.waitingHBResponse {
		c.waitingHBResponse = false
		c.hbMu.Unlock()
		return
	}
	c.hbMu.Unlock()
	_ = c.Send(&packets.HeartBeat{})
}

func (c *WsClient) stopHeartbeat() {
	c.hbMu.Lock()
	defer c.hbMu.Unlock()
	if c.hbTimer != nil {
		c.hbTimer.Stop()
		c.hbTimer = nil
	}
}

func (c *WsClient) onDisconnected() {
	_ = c.Close()
}

func (c *WsClient) notifyDisconnected() {
	c.disconnectMu.Lock()
	listeners := append([]func(*WsClient){}, c.disconnectListeners...)
	c.disconnectListeners = nil
	c.disconnectMu.Unlock()

	for _, listener := range listeners {
		listener(c)
	}
}