Something went wrong. Try again.
Monorepo for Tangled forked from tangled.org/core
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221package jetstream
import ( "context" _ "embed" "encoding/json" "fmt" "log/slog" "net/http" "net/url" "strconv" "time"
comatproto "github.com/bluesky-social/indigo/api/atproto" "github.com/gorilla/websocket" "github.com/klauspost/compress/zstd")
// bsky jetstream compresses with this dictionary, and hydrant's frames decode with it too////go:embed zstd_dictionaryvar zstdDictionary []byte
const ( EventKindCommit = "commit" EventKindAccount = "account" EventKindIdentity = "identity"
CommitOperationCreate = "create" CommitOperationUpdate = "update" CommitOperationDelete = "delete")
// hydrant pings every 30s, so 90s of silence means the connection is deadconst defaultReadTimeout = 90 * time.Second
// refused dials fail instantly, without the wait this loop spins hotconst reconnectDelay = time.Second
// indigo's jetstream event plus a live flag on commits, which is what hydrant's /subscribe sendstype Event struct { Did string `json:"did"` TimeUS int64 `json:"time_us"` Kind string `json:"kind,omitempty"` Commit *Commit `json:"commit,omitempty"` Account *comatproto.SyncSubscribeRepos_Account `json:"account,omitempty"` Identity *comatproto.SyncSubscribeRepos_Identity `json:"identity,omitempty"`}
type Commit struct { Rev string `json:"rev,omitempty"` Operation string `json:"operation,omitempty"` Collection string `json:"collection,omitempty"` RKey string `json:"rkey,omitempty"` Record json.RawMessage `json:"record,omitempty"` CID string `json:"cid,omitempty"` // plain jetstream never sends this, so nil counts as live Live *bool `json:"live,omitempty"`}
func (c *Commit) IsLive() bool { return c.Live == nil || *c.Live}
type StreamConfig struct { Endpoint string Collections []string // hydrant only backfills when asked here, plain jetstream doesn't know the param EventTypes []string Compress bool
// read on every connect, ok false replays from 0 unless TailWhenNoCursor is set Cursor func() (timeUs int64, ok bool, err error) TailWhenNoCursor bool // each hydrant stamps time_us its own way, so resume a little before the cursor Rewind time.Duration
ReadTimeout time.Duration ReconnectDelay time.Duration UserAgent string}
// the handler saves the cursor in the same transaction as the event, so a crash can't keep one without the othertype Stream struct { cfg StreamConfig l *slog.Logger dec *zstd.Decoder}
func NewStream(cfg StreamConfig, logger *slog.Logger) (*Stream, error) { if cfg.Cursor == nil { return nil, fmt.Errorf("stream needs a cursor source") } if cfg.ReadTimeout == 0 { cfg.ReadTimeout = defaultReadTimeout } if cfg.ReconnectDelay == 0 { cfg.ReconnectDelay = reconnectDelay } s := &Stream{cfg: cfg, l: logger} if cfg.Compress { dec, err := zstd.NewReader(nil, zstd.WithDecoderDicts(zstdDictionary)) if err != nil { return nil, fmt.Errorf("creating zstd decoder: %w", err) } s.dec = dec } return s, nil}
// a handler error drops the connection, so the event comes back from the cursorfunc (s *Stream) Run(ctx context.Context, handle func(context.Context, *Event) error) error { for { if err := s.readOnce(ctx, handle); err != nil && ctx.Err() == nil { s.l.Error("reading stream", "err", err) } select { case <-ctx.Done(): return nil case <-time.After(s.cfg.ReconnectDelay): } }}
// nil leaves the cursor param off, which starts at the live tailfunc (s *Stream) resumeCursor() (*int64, error) { us, ok, err := s.cfg.Cursor() if err != nil { return nil, err } if !ok { if s.cfg.TailWhenNoCursor { return nil, nil } us = 0 } cursor := max(us-s.cfg.Rewind.Microseconds(), 0) return &cursor, nil}
func (s *Stream) url(cursor *int64) (string, error) { u, err := url.Parse(s.cfg.Endpoint) if err != nil { return "", fmt.Errorf("parsing endpoint %q: %w", s.cfg.Endpoint, err) } q := u.Query() q.Del("cursor") if cursor != nil { q.Set("cursor", strconv.FormatInt(*cursor, 10)) } q.Del("wantedCollections") for _, nsid := range s.cfg.Collections { q.Add("wantedCollections", nsid) } q.Del("wantedEventTypes") for _, t := range s.cfg.EventTypes { q.Add("wantedEventTypes", t) } u.RawQuery = q.Encode() return u.String(), nil}
func (s *Stream) readOnce(ctx context.Context, handle func(context.Context, *Event) error) error { cursor, err := s.resumeCursor() if err != nil { return fmt.Errorf("reading cursor: %w", err) } target, err := s.url(cursor) if err != nil { return err }
header := http.Header{} if s.cfg.UserAgent != "" { header.Set("User-Agent", s.cfg.UserAgent) } if s.dec != nil { header.Set("Socket-Encoding", "zstd") }
s.l.Info("connecting to stream", "url", target) con, _, err := websocket.DefaultDialer.DialContext(ctx, target, header) if err != nil { return err } defer con.Close() // a quiet stream blocks the read and never checks ctx, so cancel closes the conn defer context.AfterFunc(ctx, func() { con.Close() })()
alive := func() error { return con.SetReadDeadline(time.Now().Add(s.cfg.ReadTimeout)) } con.SetPingHandler(func(data string) error { if err := alive(); err != nil { return err } return con.WriteControl(websocket.PongMessage, []byte(data), time.Now().Add(time.Second)) })
for { if err := alive(); err != nil { return err } _, msg, err := con.ReadMessage() if err != nil { return fmt.Errorf("reading message: %w", err) } if s.dec != nil { if msg, err = s.dec.DecodeAll(msg, nil); err != nil { return fmt.Errorf("decompressing message: %w", err) } } var e Event if err := json.Unmarshal(msg, &e); err != nil { s.l.Warn("skipping undecodable frame", "err", err) continue } if err := handle(ctx, &e); err != nil { return fmt.Errorf("handling event at %d: %w", e.TimeUS, err) } }}