package knotfeed
import (
"cmp"
"context"
"errors"
"fmt"
"log/slog"
"net/url"
"slices"
"strconv"
"strings"
"time"
"github.com/gorilla/websocket"
)
const (
livenessTimeout = 90 * time.Second
pongWriteTimeout = 5 * time.Second
initialBackoff = time.Second
maxBackoff = 60 * time.Second
maxFrameBytes = 8 << 20
subscribeReposNS = "com.atproto.sync.subscribeRepos"
eventsPath = "/events"
defaultMaxHandlerAttempts = 3
defaultPoisonGrace = 10 * time.Minute
)
type Feed struct {
kind uint8
}
const (
atprotoKind uint8 = iota
legacyKind
)
var (
FeedAtproto = Feed{kind: atprotoKind}
FeedLegacy = Feed{kind: legacyKind}
)
type feedSpec struct {
token string
name string
path string
msgKind int
other Feed
decode func([]byte, *slog.Logger) (Message, error)
}
var feedSpecs = [...]feedSpec{
atprotoKind: {
token: "atproto", name: "atproto firehose", path: "/xrpc/" + subscribeReposNS,
msgKind: websocket.BinaryMessage, other: FeedLegacy, decode: Decode,
},
legacyKind: {
token: "legacy", name: "knot event stream", path: eventsPath,
msgKind: websocket.TextMessage, other: FeedAtproto, decode: decodeLegacy,
},
}
func (f Feed) spec() feedSpec { return feedSpecs[f.kind] }
func (f Feed) Token() string { return f.spec().token }
func (f Feed) String() string { return f.spec().name }
func (f Feed) other() Feed { return f.spec().other }
func (f Feed) Live(now time.Time) Cursor {
if f == FeedLegacy {
return NewCursor(f, now.UnixNano())
}
return NewCursor(f, 0)
}
var ErrUnrecognizedFeed = errors.New("unrecognized feed token")
func ParseFeed(token string) (Feed, error) {
kind := slices.IndexFunc(feedSpecs[:], func(spec feedSpec) bool { return spec.token == token })
if kind < 0 {
return FeedAtproto, fmt.Errorf("%w: %q", ErrUnrecognizedFeed, token)
}
return Feed{kind: uint8(kind)}, nil
}
type Cursor struct {
feed Feed
seq int64
}
func NewCursor(feed Feed, seq int64) Cursor {
return Cursor{feed: feed, seq: seq}
}
func (c Cursor) Feed() Feed { return c.feed }
func (c Cursor) Seq() int64 { return c.seq }
func (c Cursor) Encode() string {
return c.feed.Token() + ":" + strconv.FormatInt(c.seq, 10)
}
func ParseCursor(raw string) (Cursor, error) {
token, rest, named := strings.Cut(raw, ":")
if !named {
token, rest = FeedAtproto.Token(), raw
}
feed, err := ParseFeed(token)
if err != nil {
return Cursor{}, err
}
seq, err := strconv.ParseInt(rest, 10, 64)
if err != nil {
return Cursor{}, fmt.Errorf("cursor %q: %w", raw, err)
}
return NewCursor(feed, seq), nil
}
type Consumer struct {
Host string
NoTLS bool
Dialer *websocket.Dialer
Logger *slog.Logger
LoadCursor func(context.Context) (Cursor, error)
StoreCursor func(context.Context, Cursor) error
Handle func(context.Context, Message) error
OutdatedReplay func(context.Context, Feed) Cursor
OnConnectError func(error)
ReplayFromStart bool
MaxHandlerAttempts int
PoisonGrace time.Duration
now func() time.Time
current Feed
adopted bool
proven bool
poisonSeq int64
poisonAttempts int
poisonSince time.Time
}
func (c *Consumer) Run(ctx context.Context) error {
c.fill()
backoff := initialBackoff
for {
if ctx.Err() != nil {
return nil
}
advanced, resync, err := c.session(ctx)
switch {
case err != nil:
c.reactTo(err)
case resync:
advanced = true
}
if advanced {
backoff = initialBackoff
} else {
backoff = min(backoff*2, maxBackoff)
}
if resync {
continue
}
timer := time.NewTimer(backoff)
select {
case <-ctx.Done():
timer.Stop()
return nil
case <-timer.C:
}
}
}
func (c *Consumer) reactTo(err error) {
var mismatch feedMismatch
switch {
case errors.As(err, &mismatch):
case isConnectError(err):
if c.OnConnectError != nil {
c.OnConnectError(err)
}
if c.proven {
c.Logger.Error("knot turned us away on a feed it had already served",
"host", c.Host, "feed", c.current, "err", err)
return
}
default:
c.Logger.Error("session died for its own reasons",
"host", c.Host, "feed", c.current, "err", err)
return
}
c.Logger.Warn("trying the other feed, since this knot won't serve the one we asked for",
"host", c.Host, "tried", c.current, "next", c.current.other(), "err", err)
c.current, c.proven = c.current.other(), false
}
type connectError struct{ err error }
func (e connectError) Error() string { return e.err.Error() }
func (e connectError) Unwrap() error { return e.err }
func isConnectError(err error) bool {
var ce connectError
return errors.As(err, &ce)
}
type feedMismatch struct{ feed Feed }
func (e feedMismatch) Error() string {
return fmt.Sprintf("%s sent the %s's frames", e.feed, e.feed.other())
}
func (c *Consumer) session(ctx context.Context) (advanced bool, resync bool, err error) {
stored, err := c.LoadCursor(ctx)
if err != nil {
return false, false, fmt.Errorf("loading cursor: %w", err)
}
if !c.adopted {
c.adopted = true
if stored.Seq() > 0 {
c.current = stored.Feed()
}
}
cursor := c.sessionCursor(stored)
conn, _, err := c.Dialer.DialContext(ctx, c.url(cursor), nil)
if err != nil {
return false, false, connectError{err}
}
c.proven = true
defer conn.Close()
watcher := context.AfterFunc(ctx, func() { conn.Close() })
defer watcher()
conn.SetReadLimit(maxFrameBytes)
c.Logger.Info("subscribed to knot feed", "host", c.Host, "feed", c.current, "cursor", cursor)
conn.SetPingHandler(func(payload string) error {
conn.SetReadDeadline(time.Now().Add(livenessTimeout))
return conn.WriteControl(websocket.PongMessage, []byte(payload), time.Now().Add(pongWriteTimeout))
})
conn.SetReadDeadline(time.Now().Add(livenessTimeout))
want := c.current.spec().msgKind
for {
kind, data, err := conn.ReadMessage()
if err != nil {
return advanced, false, err
}
conn.SetReadDeadline(time.Now().Add(livenessTimeout))
if kind != want {
return advanced, false, feedMismatch{feed: c.current}
}
message, err := c.decode(data)
if err != nil {
c.Logger.Error("undecodable knot feed frame", "host", c.Host, "feed", c.current, "err", err)
continue
}
resync, err := c.dispatch(ctx, message, &cursor)
if resync {
return advanced, true, nil
}
if err != nil {
return false, false, err
}
if message.Type == TypeCommit {
advanced = true
}
}
}
func (c *Consumer) decode(data []byte) (Message, error) {
message, err := c.current.spec().decode(data, c.Logger)
message.Feed = c.current
return message, err
}
func (c *Consumer) dispatch(ctx context.Context, message Message, cursor *Cursor) (bool, error) {
if message.Type == TypeCommit {
return c.commit(ctx, message, cursor)
}
if c.unreachable(message) {
c.Logger.Warn("firehose can't serve our cursor, resuming live", "host", c.Host, "error", c.resumeReason(message))
*cursor = c.replayCursor(ctx)
if err := c.store(ctx, *cursor); err != nil {
return false, err
}
return true, nil
}
if c.Handle != nil {
if err := c.Handle(ctx, message); err != nil {
c.Logger.Error("firehose handler failed", "host", c.Host, "frame", message.Type, "err", err)
}
}
return false, nil
}
func (c *Consumer) replayCursor(ctx context.Context) Cursor {
live := c.current.Live(c.now())
if c.OutdatedReplay == nil {
return live
}
replay := c.OutdatedReplay(ctx, c.current)
if replay.Feed() != c.current {
c.Logger.Warn("replay hook named the other feed, so resuming live instead",
"host", c.Host, "feed", c.current, "replay", replay.Feed())
return live
}
return replay
}
func (c *Consumer) commit(ctx context.Context, message Message, cursor *Cursor) (bool, error) {
if c.Handle != nil {
if err := c.Handle(ctx, message); err != nil {
if !c.poison(message.Commit.Seq) {
return false, fmt.Errorf("frame seq %d: %w", message.Commit.Seq, err)
}
c.Logger.Error("handler keeps failing at one frame, skipping it",
"host", c.Host, "seq", message.Commit.Seq, "attempts", c.poisonAttempts, "err", err)
}
}
*cursor = NewCursor(c.current, message.Commit.Seq)
if err := c.store(ctx, *cursor); err != nil {
return false, err
}
return false, nil
}
func (c *Consumer) store(ctx context.Context, cursor Cursor) error {
if err := c.StoreCursor(ctx, cursor); err != nil {
return fmt.Errorf("storing cursor: %w", err)
}
return nil
}
func (c *Consumer) poison(seq int64) bool {
if seq != c.poisonSeq {
c.poisonSeq = seq
c.poisonAttempts = 0
c.poisonSince = c.now()
}
c.poisonAttempts++
return c.poisonAttempts >= c.MaxHandlerAttempts && c.now().Sub(c.poisonSince) >= c.PoisonGrace
}
func (c *Consumer) unreachable(message Message) bool {
if message.Type == TypeError {
return message.Error == "FutureCursor"
}
return message.Type == TypeInfo && message.InfoName == "OutdatedCursor"
}
func (c *Consumer) resumeReason(message Message) string {
if message.Type == TypeError {
return message.Error
}
return message.InfoName
}
func (c *Consumer) url(cursor Cursor) string {
scheme := "wss"
if c.NoTLS {
scheme = "ws"
}
endpoint := url.URL{Scheme: scheme, Host: c.Host, Path: c.current.spec().path}
if seq := cursor.Seq(); seq > 0 || (seq == 0 && c.ReplayFromStart) {
endpoint.RawQuery = "cursor=" + strconv.FormatInt(seq, 10)
}
return endpoint.String()
}
func (c *Consumer) sessionCursor(stored Cursor) Cursor {
switch {
case stored.Feed() == c.current && stored.Seq() > 0:
return stored
case c.ReplayFromStart:
return NewCursor(c.current, 0)
default:
return c.current.Live(c.now())
}
}
func (c *Consumer) fill() {
c.MaxHandlerAttempts = cmp.Or(c.MaxHandlerAttempts, defaultMaxHandlerAttempts)
c.PoisonGrace = cmp.Or(c.PoisonGrace, defaultPoisonGrace)
c.Dialer = cmp.Or(c.Dialer, websocket.DefaultDialer)
c.Logger = cmp.Or(c.Logger, slog.Default())
if c.now == nil {
c.now = time.Now
}
}