package lexutil
import (
"cmp"
"context"
"fmt"
"log/slog"
"net/http"
"net/url"
"time"
indigoxrpc "github.com/bluesky-social/indigo/xrpc"
"github.com/carlmjohnson/versioninfo"
"github.com/gorilla/websocket"
cbg "github.com/whyrusleeping/cbor-gen"
)
const minHealthyConn = 30 * time.Second
type Client struct {
indigoxrpc.Client
Dialer websocket.Dialer
Logger *slog.Logger
}
var _ LexClient = (*Client)(nil)
func makeParams(p map[string]any) url.Values {
params := url.Values{}
for k, v := range p {
if s, ok := v.([]string); ok {
for _, v := range s {
params.Add(k, v)
}
} else {
params.Add(k, fmt.Sprint(v))
}
}
return params
}
type processFn func(ctx context.Context, cr *cbg.CborReader) error
func (c *Client) LexDo(ctx context.Context, method string, inputEncoding string, endpoint string, params map[string]any, bodyData any, out any) error {
switch method {
case Subscription:
if process, ok := out.(func(context.Context, *cbg.CborReader) error); ok {
return c.LexSubscribe(ctx, endpoint, params, process)
} else if process, ok := out.(processFn); ok {
return c.LexSubscribe(ctx, endpoint, params, process)
} else if redialer, ok := out.(Redialer); ok {
return c.LexSubscribeWithRedialer(ctx, endpoint, params, redialer)
} else {
return fmt.Errorf("unknown output type: %T", out)
}
default:
return c.Client.LexDo(ctx, method, inputEncoding, endpoint, params, bodyData, out)
}
}
func (c *Client) getHeader() http.Header {
header := http.Header{}
if c.UserAgent != nil {
header.Set("User-Agent", *c.UserAgent)
} else {
header.Set("User-Agent", "extlexutil/"+versioninfo.Short())
}
if c.Headers != nil {
for k, v := range c.Headers {
header.Set(k, v)
}
}
return header
}
func (c *Client) LexSubscribe(ctx context.Context, endpoint string, params map[string]any, process func(ctx context.Context, cr *cbg.CborReader) error) error {
logger := cmp.Or(c.Logger, slog.Default().With("system", "events"))
rurl, err := url.Parse(c.Host)
if err != nil {
return err
}
if rurl.Scheme == "http" {
rurl.Scheme = "ws"
} else {
rurl.Scheme = "wss"
}
surl := rurl.JoinPath("/xrpc", endpoint)
surl.RawQuery = makeParams(params).Encode()
header := c.getHeader()
u := surl.String()
conn, resp, err := c.Dialer.DialContext(ctx, u, header)
if err != nil {
return fmt.Errorf("%w: %w", ErrDialFailure, err)
}
logger.Debug("event subscription response", "code", resp.StatusCode, "url", u)
return c.handleConn(ctx, conn, process)
}
func (c *Client) LexSubscribeWithRedialer(ctx context.Context, endpoint string, params map[string]any, redialer Redialer) error {
logger := cmp.Or(c.Logger, slog.Default().With("system", "events"))
rurl, err := url.Parse(c.Host)
if err != nil {
return err
}
if rurl.Scheme == "http" {
rurl.Scheme = "ws"
} else {
rurl.Scheme = "wss"
}
surl := rurl.JoinPath("/xrpc", endpoint)
header := c.getHeader()
var backoff int
// returns false if the retry budget is exhausted
sleepBackoff := func() bool {
select {
case <-ctx.Done():
case <-time.After(time.Duration(5+backoff) * time.Second):
}
backoff++
return backoff <= 15
}
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
surl.RawQuery = makeParams(params).Encode()
u := surl.String()
conn, resp, err := c.Dialer.DialContext(ctx, u, header)
if err != nil {
logger.Warn("dialing failed", "err", err, "backoff", backoff)
if !sleepBackoff() {
return fmt.Errorf("%w: %w", ErrDialFailure, err)
}
continue
}
logger.Debug("event subscription response", "code", resp.StatusCode, "url", u)
connectedAt := time.Now()
connErr := c.handleConn(ctx, conn, redialer.Process)
if connErr != nil {
logger.Warn("host connection failed", "err", connErr, "backoff", backoff)
}
// updates cursor
updated := redialer.UpdateParams(ctx, params)
// a connection that drops immediately shouldnt reset backoff
// this to avoid reconnect storms
if updated || time.Since(connectedAt) >= minHealthyConn {
backoff = 0
continue
}
if !sleepBackoff() {
return fmt.Errorf("%w: %w", ErrConnFailure, connErr)
}
}
}
func (c *Client) handleConn(ctx context.Context, conn *websocket.Conn, process func(ctx context.Context, cr *cbg.CborReader) error) error {
logger := cmp.Or(c.Logger, slog.Default().With("system", "events"))
ctx, cancel := context.WithCancel(ctx)
defer cancel()
go func() {
t := time.NewTicker(time.Second * 30)
defer t.Stop()
failcount := 0
for {
select {
case <-t.C:
if err := conn.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(time.Second*10)); err != nil {
logger.Warn("failed to ping", "err", err)
failcount++
if failcount >= 4 {
logger.Error("too many ping fails", "count", failcount)
conn.Close()
return
}
} else {
failcount = 0 // ok ping
}
case <-ctx.Done():
conn.Close()
return
}
}
}()
conn.SetPingHandler(func(message string) error {
err := conn.WriteControl(websocket.PongMessage, []byte(message), time.Now().Add(time.Second*60))
if err == websocket.ErrCloseSent {
return nil
}
return err
})
conn.SetPongHandler(func(_ string) error {
if err := conn.SetReadDeadline(time.Now().Add(time.Minute)); err != nil {
logger.Error("failed to set read deadline", "err", err)
}
return nil
})
cr := new(cbg.CborReader)
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
mt, rawReader, err := conn.NextReader()
if err != nil {
return fmt.Errorf("conn err at read: %w", err)
}
if mt != websocket.BinaryMessage {
return fmt.Errorf("expected binary message from subscription endpoint")
}
cr.SetReader(rawReader)
if err := process(ctx, cr); err != nil {
return err
}
}
}