toki
4dbeecf2fb
- Add gateway real path implementation (06_gateway_real_path) - Add gateway cleanup tasks (07+06_gateway_cleanup) - Update Go WebSocket client and benchmark - Update Kotlin Communicator, TcpClient, WsClient and stress tests - Update TypeScript base_client, communicator, inbound_gateway, node variants, tcp/ws clients and tests - Update Kotlin and TypeScript stress benchmarks - Update roadmap milestone documentation
121 lines
3.6 KiB
Go
121 lines
3.6 KiB
Go
package proto_socket
|
|
|
|
import (
|
|
"context"
|
|
"crypto/tls"
|
|
"fmt"
|
|
"net/http"
|
|
"sync"
|
|
|
|
"google.golang.org/protobuf/proto"
|
|
"nhooyr.io/websocket"
|
|
|
|
"git.toki-labs.com/toki/proto-socket/go/packets"
|
|
)
|
|
|
|
const (
|
|
DefaultHeartbeatIntervalSec = 30
|
|
DefaultHeartbeatWaitSec = 10
|
|
)
|
|
|
|
type WsClient struct {
|
|
baseClient[*WsClient]
|
|
conn *websocket.Conn
|
|
writeMu sync.Mutex // Defensive if WritePacket is called outside the queued writer.
|
|
readCtx context.Context
|
|
cancelRead context.CancelFunc
|
|
writeCtx context.Context
|
|
cancelWrite context.CancelFunc
|
|
}
|
|
|
|
func NewWsClient(conn *websocket.Conn, intervalSec, waitSec int, parserMap ParserMap) *WsClient {
|
|
conn.SetReadLimit(MaxPacketSize)
|
|
readCtx, cancelRead := context.WithCancel(context.Background())
|
|
writeCtx, cancelWrite := context.WithCancel(context.Background())
|
|
c := &WsClient{
|
|
conn: conn,
|
|
readCtx: readCtx,
|
|
cancelRead: cancelRead,
|
|
writeCtx: writeCtx,
|
|
cancelWrite: cancelWrite,
|
|
}
|
|
c.baseClient = newBaseClient(c, intervalSec, waitSec, func() error {
|
|
c.cancelRead()
|
|
c.cancelWrite()
|
|
return c.conn.Close(websocket.StatusNormalClosure, "")
|
|
})
|
|
c.Communicator.Initialize(c, parserMap)
|
|
c.SetWriteErrorHandler(func(err error) {
|
|
c.onDisconnected(DisconnectReasonWriteError, err)
|
|
})
|
|
c.SetFrameErrorHandler(func(err error) {
|
|
c.onDisconnected(DisconnectReasonWSPacketParse, err)
|
|
})
|
|
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(DisconnectReasonWSRead, err)
|
|
return
|
|
}
|
|
if msgType != websocket.MessageBinary {
|
|
continue
|
|
}
|
|
// OnReceivedFrame routes the raw frame through the inbound gateway when
|
|
// one is attached, or decodes it inline otherwise. An inline decode
|
|
// error preserves the existing parse-error disconnect semantics.
|
|
if err := c.OnReceivedFrame(b); err != nil {
|
|
c.onDisconnected(DisconnectReasonWSPacketParse, err)
|
|
return
|
|
}
|
|
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(c.writeCtx, websocket.MessageBinary, b)
|
|
}
|