Something went wrong. Try again.
Go boilerplate library for building atproto apps
atproto go
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565// Package jetstream consumes real-time AT Protocol events from a Jetstream relay.//// Jetstream is a WebSocket-based relay that delivers a filtered stream of AT// Protocol repository events (commits, identity changes, account changes).// This package handles connection management, reconnection with backoff,// endpoint rotation, cursor tracking, optional zstd decompression, and// runtime-mutable subscription filters via the jetstream options_update// control frame.//// Basic usage://// consumer := jetstream.New(&jetstream.Config{// WantedCollections: []string{"app.bsky.feed.post"},// }, func(ctx context.Context, evt *jetstream.Event) error {// fmt.Printf("new post from %s\n", evt.DID)// return nil// })// consumer.Start(ctx)// defer consumer.Stop()package jetstream
import ( "context" _ "embed" "encoding/json" "fmt" "net/url" "slices" "sync" "sync/atomic" "time"
"github.com/gorilla/websocket" "github.com/klauspost/compress/zstd")
// zstdDictionary is the shared dictionary the Jetstream relay compresses// each frame against. Sourced from://// https://github.com/bluesky-social/jetstream/raw/main/pkg/models/zstd_dictionary////go:embed zstd_dictionaryvar zstdDictionary []byte
// DefaultEndpoints are the public Jetstream relay endpoints.var DefaultEndpoints = []string{ "wss://jetstream1.us-east.bsky.network/subscribe", "wss://jetstream2.us-east.bsky.network/subscribe", "wss://jetstream1.us-west.bsky.network/subscribe", "wss://jetstream2.us-west.bsky.network/subscribe",}
// Default liveness settings. Jetstream's server-side filters mean a connection// can sit silent indefinitely if no matching events are flowing. Application// pings keep the socket warm and let callers distinguish "upstream has nothing// for us" from "upstream is dead". The read deadline must be comfortably// greater than 2× the ping interval so a single dropped pong doesn't trip it.const ( defaultPingInterval = 30 * time.Second defaultReadDeadline = 90 * time.Second defaultPingWriteTimeout = 10 * time.Second)
// Event is a single event from the Jetstream relay.type Event struct { DID string `json:"did"` TimeUS int64 `json:"time_us"` Kind string `json:"kind"` // "commit", "identity", "account" Commit *Commit `json:"commit,omitempty"` Identity *Identity `json:"identity,omitempty"` Account *Account `json:"account,omitempty"`}
// Commit is the commit payload within an Event.type Commit struct { Rev string `json:"rev"` Operation string `json:"operation"` // "create", "update", "delete" Collection string `json:"collection"` RKey string `json:"rkey"` Record json.RawMessage `json:"record,omitempty"` CID string `json:"cid"`}
// Identity is the identity payload of an Event. Emitted when a DID's handle// or PDS endpoint changes.type Identity struct { DID string `json:"did"` Handle string `json:"handle,omitempty"` Seq int64 `json:"seq"` Time string `json:"time"`}
// Account is the account payload of an Event. Status is one of "active",// "deleted", "deactivated", "suspended", "takendown".type Account struct { Active bool `json:"active"` DID string `json:"did"` Seq int64 `json:"seq"` Status string `json:"status,omitempty"` Time string `json:"time"`}
// Handler is called for each event received from Jetstream.// Returning an error logs a warning but does not stop the consumer.type Handler func(ctx context.Context, event *Event) error
// CursorStore persists the Jetstream cursor across restarts.// If nil, the cursor is tracked in memory only (replay from live on restart).type CursorStore interface { GetCursor(ctx context.Context) (int64, error) SetCursor(ctx context.Context, cursor int64) error}
// Config configures a Jetstream consumer.type Config struct { // Endpoints is the list of Jetstream WebSocket URLs. Defaults to DefaultEndpoints. Endpoints []string
// WantedCollections filters events to specific NSIDs. // Empty means all collections (high volume). WantedCollections []string
// WantedDIDs filters events to specific DID subjects. // Empty means all DIDs. When combined with WantedCollections, the // jetstream relay applies AND semantics — events must match both filters. WantedDIDs []string
// Compress enables zstd compression on the wire. The Jetstream zstd // dictionary is embedded in the package and applied automatically. Compress bool
// CursorStore persists the cursor for resume after restart. // If nil, the consumer starts from live on each restart. CursorStore CursorStore
// CursorPersistEvery controls how often the cursor is flushed to CursorStore. // Defaults to every 1000 events. CursorPersistEvery int64
// PingInterval is how often application-level ping frames are sent to keep // the socket warm. Defaults to 30s. Set to a negative value to disable // pings entirely (rely on the read deadline alone). PingInterval time.Duration
// ReadDeadline is the per-read websocket deadline, refreshed on every // successful read and every pong. Should be comfortably greater than 2× // PingInterval. Defaults to 90s. ReadDeadline time.Duration
// OnConnect is called each time a WebSocket connection is established. // Call Consumer.CurrentEndpoint to learn which endpoint we connected to. OnConnect func()
// OnDisconnect is called each time a connection is lost. OnDisconnect func()
// OnError is called when the handler returns an error. // If nil, errors are silently dropped (caller should log in the handler). OnError func(err error, event *Event)}
func (c *Config) endpoints() []string { if len(c.Endpoints) > 0 { return c.Endpoints } return DefaultEndpoints}
func (c *Config) cursorPersistEvery() int64 { if c.CursorPersistEvery > 0 { return c.CursorPersistEvery } return 1000}
// pingInterval returns the configured ping interval, the default if zero, or// zero if pings are explicitly disabled (negative config value).func (c *Config) pingInterval() time.Duration { if c.PingInterval < 0 { return 0 } if c.PingInterval == 0 { return defaultPingInterval } return c.PingInterval}
func (c *Config) readDeadline() time.Duration { if c.ReadDeadline > 0 { return c.ReadDeadline } return defaultReadDeadline}
// optionsUpdateMessage is the jetstream control frame for runtime filter// updates. See https://github.com/bluesky-social/jetstream for the schema.type optionsUpdateMessage struct { Type string `json:"type"` Payload optionsUpdatePayload `json:"payload"`}
type optionsUpdatePayload struct { WantedCollections []string `json:"wantedCollections"` WantedDIDs []string `json:"wantedDids"`}
// Consumer consumes events from a Jetstream relay.type Consumer struct { cfg *Config handler Handler
conn *websocket.Conn connMu sync.Mutex currentEndpointIdx int currentEndpoint atomic.Value // string
// writeMu serializes control-frame writes (pings and options_update). // gorilla/websocket is not safe for concurrent writes. writeMu sync.Mutex
// Mutable subscription state. Mirrors the config at construction time but // can be replaced at runtime via UpdateOptions. optMu sync.RWMutex wantedCollections []string wantedDIDs []string
zstdDecoder *zstd.Decoder
cursor atomic.Int64 eventsReceived atomic.Int64 bytesReceived atomic.Int64 connected atomic.Bool
stopCh chan struct{} wg sync.WaitGroup}
// New creates a new Consumer. Call Start to begin consuming events.func New(cfg *Config, handler Handler) *Consumer { opts := []zstd.DOption{zstd.WithDecoderConcurrency(1)} if cfg.Compress { opts = append(opts, zstd.WithDecoderDicts(zstdDictionary)) } decoder, err := zstd.NewReader(nil, opts...) if err != nil { // zstd.NewReader with nil src only fails on bad options panic(fmt.Sprintf("jetstream: create zstd decoder: %v", err)) }
c := &Consumer{ cfg: cfg, handler: handler, stopCh: make(chan struct{}), zstdDecoder: decoder, wantedCollections: slices.Clone(cfg.WantedCollections), wantedDIDs: slices.Clone(cfg.WantedDIDs), } c.currentEndpoint.Store("")
if cfg.CursorStore != nil { if cursor, err := cfg.CursorStore.GetCursor(context.Background()); err == nil && cursor > 0 { c.cursor.Store(cursor) } }
return c}
// Start begins consuming events in a background goroutine.func (c *Consumer) Start(ctx context.Context) { c.wg.Go(func() { c.run(ctx) })}
// Stop gracefully shuts down the consumer and waits for it to finish.func (c *Consumer) Stop() { close(c.stopCh) c.connMu.Lock() if c.conn != nil { c.conn.Close() } c.connMu.Unlock() c.wg.Wait() c.zstdDecoder.Close()}
// IsConnected reports whether the consumer is currently connected.func (c *Consumer) IsConnected() bool { return c.connected.Load()}
// CurrentEndpoint returns the endpoint URL of the active connection, or the// empty string if not connected. Safe for concurrent use.func (c *Consumer) CurrentEndpoint() string { v, _ := c.currentEndpoint.Load().(string) return v}
// Stats returns cumulative event and byte counts since Start was called.func (c *Consumer) Stats() (eventsReceived, bytesReceived int64) { return c.eventsReceived.Load(), c.bytesReceived.Load()}
// UpdateOptions atomically replaces the subscription filters. If currently// connected, an options_update control frame is sent to the relay so the new// filters take effect immediately. Otherwise the new filters apply on the// next reconnect via the WebSocket URL.//// Passing nil slices is equivalent to passing empty slices: the filter is// cleared and the relay delivers all events for that dimension.func (c *Consumer) UpdateOptions(wantedCollections, wantedDIDs []string) error { collsCopy := slices.Clone(wantedCollections) didsCopy := slices.Clone(wantedDIDs)
c.optMu.Lock() c.wantedCollections = collsCopy c.wantedDIDs = didsCopy c.optMu.Unlock()
c.connMu.Lock() conn := c.conn c.connMu.Unlock() if conn == nil { return nil // applied via URL on next reconnect }
msg := optionsUpdateMessage{ Type: "options_update", Payload: optionsUpdatePayload{ WantedCollections: collsCopy, WantedDIDs: didsCopy, }, } data, err := json.Marshal(msg) if err != nil { return fmt.Errorf("marshal options_update: %w", err) }
c.writeMu.Lock() defer c.writeMu.Unlock() if err := conn.WriteMessage(websocket.TextMessage, data); err != nil { return fmt.Errorf("send options_update: %w", err) } return nil}
func (c *Consumer) run(ctx context.Context) { backoff := time.Second const maxBackoff = 30 * time.Second
for { select { case <-ctx.Done(): return case <-c.stopCh: return default: }
endpoints := c.cfg.endpoints() endpoint := endpoints[c.currentEndpointIdx]
if err := c.connectAndConsume(ctx, endpoint); err != nil { c.connected.Store(false) if c.cfg.OnDisconnect != nil { c.cfg.OnDisconnect() }
// Rotate to next endpoint c.currentEndpointIdx = (c.currentEndpointIdx + 1) % len(endpoints)
select { case <-ctx.Done(): return case <-c.stopCh: return case <-time.After(backoff): }
backoff *= 2 if backoff > maxBackoff { backoff = maxBackoff } } else { backoff = time.Second } }}
func (c *Consumer) connectAndConsume(ctx context.Context, endpoint string) error { wsURL, err := c.buildURL(endpoint) if err != nil { return fmt.Errorf("build URL: %w", err) }
dialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second} conn, _, err := dialer.DialContext(ctx, wsURL, nil) if err != nil { return fmt.Errorf("dial: %w", err) }
c.connMu.Lock() c.conn = conn c.connMu.Unlock()
c.currentEndpoint.Store(endpoint) c.connected.Store(true) if c.cfg.OnConnect != nil { c.cfg.OnConnect() }
defer func() { c.connMu.Lock() if c.conn != nil { c.conn.Close() c.conn = nil } c.connMu.Unlock() c.currentEndpoint.Store("") c.connected.Store(false) }()
deadline := c.cfg.readDeadline() conn.SetReadDeadline(time.Now().Add(deadline)) conn.SetPongHandler(func(string) error { return conn.SetReadDeadline(time.Now().Add(deadline)) })
// Pings run in their own goroutine. Cancel them before the connection // teardown above so we don't write to a closed conn. pingCtx, cancelPing := context.WithCancel(ctx) defer cancelPing() c.startPingLoop(pingCtx, conn)
for { select { case <-ctx.Done(): return ctx.Err() case <-c.stopCh: return nil default: }
_, msg, err := conn.ReadMessage() if err != nil { return fmt.Errorf("read: %w", err) }
// Refresh deadline on every successful read. The PongHandler covers // idle periods; this covers active ones. conn.SetReadDeadline(time.Now().Add(deadline))
c.bytesReceived.Add(int64(len(msg)))
if err := c.process(ctx, msg); err != nil { if c.cfg.OnError != nil { // We don't have the event here since parsing may have failed, // pass nil to signal a parse/process error c.cfg.OnError(err, nil) } } }}
func (c *Consumer) startPingLoop(ctx context.Context, conn *websocket.Conn) { interval := c.cfg.pingInterval() if interval <= 0 { return } go func() { t := time.NewTicker(interval) defer t.Stop() for { select { case <-ctx.Done(): return case <-c.stopCh: return case <-t.C: c.writeMu.Lock() err := conn.WriteControl( websocket.PingMessage, nil, time.Now().Add(defaultPingWriteTimeout), ) c.writeMu.Unlock() if err != nil { // The read loop will surface the underlying connection // failure; we just stop pinging. return } } } }()}
func (c *Consumer) buildURL(endpoint string) (string, error) { u, err := url.Parse(endpoint) if err != nil { return "", err }
q := u.Query()
c.optMu.RLock() for _, coll := range c.wantedCollections { q.Add("wantedCollections", coll) } for _, did := range c.wantedDIDs { q.Add("wantedDids", did) } c.optMu.RUnlock()
if c.cfg.Compress { q.Set("compress", "true") } if cursor := c.cursor.Load(); cursor > 0 { // Rewind 5 seconds to cover any gaps at reconnect rewind := cursor - (5 * time.Second.Microseconds()) q.Set("cursor", fmt.Sprintf("%d", rewind)) }
u.RawQuery = q.Encode() return u.String(), nil}
func (c *Consumer) process(ctx context.Context, data []byte) error { if c.cfg.Compress { decompressed, err := c.zstdDecoder.DecodeAll(data, nil) if err != nil { return fmt.Errorf("decompress: %w", err) } data = decompressed }
var event Event if err := json.Unmarshal(data, &event); err != nil { return fmt.Errorf("unmarshal event: %w", err) }
c.eventsReceived.Add(1)
if event.TimeUS > 0 { c.cursor.Store(event.TimeUS)
if c.cfg.CursorStore != nil && c.eventsReceived.Load()%c.cfg.cursorPersistEvery() == 0 { if err := c.cfg.CursorStore.SetCursor(ctx, event.TimeUS); err != nil { // Non-fatal: log via OnError if configured if c.cfg.OnError != nil { c.cfg.OnError(fmt.Errorf("persist cursor: %w", err), nil) } } } }
if err := c.handler(ctx, &event); err != nil { if c.cfg.OnError != nil { c.cfg.OnError(err, &event) } }
return nil}