Something went wrong. Try again.
Monorepo for Tangled tangled.org
Something went wrong. Try again.
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398package 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 }}