jetstream client
atproto jetstream client
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292package main
import ( "bytes" "context" "encoding/json" "fmt" "net/http" "net/http/httptest" "os" "path/filepath" "strconv" "strings" "sync/atomic" "time"
js "github.com/bluesky-social/jetstream" "github.com/bluesky-social/jetstream/segment" "github.com/coder/websocket" "github.com/jcalabro/atmos/cbor")
const digestSeed uint64 = 14695981039346656037
func recordDigest(h uint64, b []byte) uint64 { for _, v := range b { h = (h ^ uint64(v)) * 1099511628211 } return h}func catchupRecord(seq int) map[string]any { return map[string]any{"$type": "app.bsky.feed.post", "text": fmt.Sprintf("post %d %s", seq, strings.Repeat("synthetic archive content. ", 8))}}
type catchupFixture struct { count int frames [][]byte checksum uint64 delay time.Duration requests, wireBytes, liveConnections atomic.Int64 active, peak atomic.Int64}
func newCatchup(count int) (*catchupFixture, error) { f := &catchupFixture{count: count, checksum: digestSeed} dir, err := os.MkdirTemp("", "jetstream-catchup-") if err != nil { return nil, err } defer os.RemoveAll(dir) path := filepath.Join(dir, "backlog.jss") writer, err := segment.New(segment.Config{Path: path, MaxEventsPerBlock: 1000}) if err != nil { return nil, err } defer writer.Close() for seq := 1; seq <= count+1; seq++ { payload, err := cbor.Marshal(catchupRecord(seq)) if err != nil { return nil, err } f.checksum = recordDigest(f.checksum, payload) if seq > count { break } _, err = writer.Append(segment.Event{Seq: uint64(seq), WitnessedAt: 1783900801000000, Kind: segment.KindCreate, DID: "did:plc:fixture", Collection: "app.bsky.feed.post", Rkey: strconv.Itoa(seq), Rev: "rev1", Payload: payload}) if err != nil { return nil, err } if seq%1000 == 0 && seq < count { if err = writer.Flush(); err != nil { return nil, err } } } if _, err = writer.Seal(); err != nil { return nil, err } raw, err := os.ReadFile(path) if err != nil { return nil, err } header, err := segment.ReadSealedHeader(bytes.NewReader(raw)) if err != nil { return nil, err } for i := 0; i < int(header.BlockCount); i++ { frame, err := segment.ReadBlockFrame(bytes.NewReader(raw), header, i) if err != nil { return nil, err } f.frames = append(f.frames, frame) } return f, nil}func (f *catchupFixture) ServeHTTP(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/xrpc/network.bsky.jetstream.planSnapshot": var in struct { After int `json:"afterSeq"` } if json.NewDecoder(r.Body).Decode(&in) != nil || in.After != 0 { http.Error(w, "invalid catch-up start", 400) return } _ = json.NewEncoder(w).Encode(map[string]any{"plannedThroughSeq": f.count, "sealedTipSeq": f.count, "segments": []any{map[string]any{"name": "backlog", "index": 0, "minSeq": 1, "maxSeq": f.count, "mode": "blocks", "blocks": []any{map[string]any{"first": 0, "last": len(f.frames) - 1}}}}}) case "/xrpc/network.bsky.jetstream.getBlock": active := f.active.Add(1) defer f.active.Add(-1) for old := f.peak.Load(); active > old; old = f.peak.Load() { if f.peak.CompareAndSwap(old, active) { break } } i, err := strconv.Atoi(r.URL.Query().Get("blockIndex")) if err != nil || i < 0 || i >= len(f.frames) { http.Error(w, "invalid block", 400) return } if f.delay > 0 { timer := time.NewTimer(f.delay) defer timer.Stop() select { case <-timer.C: case <-r.Context().Done(): return } } f.requests.Add(1) n, _ := w.Write(f.frames[i]) f.wireBytes.Add(int64(n)) case "/xrpc/network.bsky.jetstream.subscribeEvents": if r.URL.Query().Get("cursor") != strconv.Itoa(f.count) { http.Error(w, "wrong handoff cursor", 400) return } f.liveConnections.Add(1) conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{Subprotocols: []string{"xrpc.v1.json"}}) if err != nil { return } defer conn.CloseNow() ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second) defer cancel() for _, seq := range []int{f.count, f.count + 1} { b, err := json.Marshal(map[string]any{"$type": "message", "payload": map[string]any{"$type": "network.bsky.jetstream.subscribeEvents#commit", "collection": "app.bsky.feed.post", "did": "did:plc:fixture", "operation": "create", "record": catchupRecord(seq), "rev": "rev1", "rkey": strconv.Itoa(seq), "seq": seq, "time": "2026-07-13T00:00:01.000000Z"}}) if err != nil { return } if conn.Write(ctx, websocket.MessageText, b) != nil { return } } _, _, _ = conn.Read(ctx) default: http.NotFound(w, r) }}func catchupClient(args []string) error { if len(args) != 2 { return fmt.Errorf("catchup-client HOST ARCHIVE_COUNT") } count, err := strconv.Atoi(args[1]) if err != nil || count < 1 || count > 1_000_000 { return fmt.Errorf("invalid archive count") } fmt.Println("READY") start := time.Now() client, err := js.Subscribe(args[0], js.WithAfterSeq(0), js.WithBatchSize(1), js.WithDownloadConcurrency(4), js.WithZstdCompression(false)) if err != nil { return err } defer client.Close() ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) defer cancel() seen := 0 digest := digestSeed for batch, err := range client.Events(ctx) { if err != nil { return err } for _, event := range batch.Events() { seen++ if event.Seq != uint64(seen) || event.Commit == nil { return fmt.Errorf("out-of-order catch-up") } digest = recordDigest(digest, event.Commit.RecordCBOR) } if seen == count+1 { break } } if seen != count+1 { return fmt.Errorf("incomplete catch-up") } fmt.Printf("CATCHUP %d %d %d\n", seen, digest, time.Since(start).Nanoseconds()) return nil}
type catchupResult struct { Profile string `json:"profile"` PeakRequests int64 `json:"peak_block_requests"` Client string `json:"client"` Repetition int `json:"repetition"` ArchiveCount int `json:"archive_count"` RequestDelayMS int64 `json:"request_delay_ms"` CatchupSeconds float64 `json:"catchup_seconds"` CPUSeconds float64 `json:"cpu_seconds"` PeakRSSBytes int64 `json:"peak_rss_bytes"` Checksum uint64 `json:"record_checksum"` Blocks int64 `json:"block_requests"` PayloadBytes int64 `json:"archive_payload_bytes"`}
func catchupMain(args []string) error { counts := []int{10000, 100000, 1000000} profile, repetitions := "full", 4 if len(args) > 1 { return fmt.Errorf("catchup [ci|ARCHIVE_COUNT]") } if len(args) == 1 && args[0] == "ci" { profile, repetitions = "ci", 2 counts = []int{10000} } else if len(args) == 1 { n, err := strconv.Atoi(args[0]) if err != nil || n < 1000 || n > 1_000_000 || n%1000 != 0 { return fmt.Errorf("count must be a multiple of 1000 between 1000 and 1000000") } counts = []int{n} profile = "count" } for _, count := range counts { fixture, err := newCatchup(count) if err != nil { return err } for rep := 0; rep < repetitions; rep++ { clients := []string{"go", "zig"} if rep%2 == 1 { clients = []string{"zig", "go"} } for _, delay := range []time.Duration{0, 20 * time.Millisecond} { for _, client := range clients { f := &catchupFixture{count: count, frames: fixture.frames, checksum: fixture.checksum, delay: delay} server := httptest.NewServer(f) exe, err := os.Executable() if err != nil { server.Close() return err } args := []string{"catchup-client", server.URL, strconv.Itoa(count)} if client == "zig" { exe = "zig-out/bin/zig-client" args = []string{server.URL, "catchup", strconv.Itoa(count)} } ctx, cancel := context.WithTimeout(context.Background(), 65*time.Second) measured, err := measureClient(ctx, exe, args) cancel() server.Close() if err != nil { return err } var seen int var digest uint64 var ns int64 if _, err = fmt.Sscanf(catchupSummary(measured.Output), "CATCHUP %d %d %d", &seen, &digest, &ns); err != nil { return err } if seen != count+1 || digest != f.checksum || f.requests.Load() != int64(len(f.frames)) || f.liveConnections.Load() != 1 { return fmt.Errorf("catch-up delivery/work mismatch: count=%d digest=%d/%d blocks=%d live=%d", seen, digest, f.checksum, f.requests.Load(), f.liveConnections.Load()) } result := catchupResult{profile, f.peak.Load(), client, rep, count, delay.Milliseconds(), float64(ns) / 1e9, measured.CPUSeconds, measured.PeakRSSBytes, digest, f.requests.Load(), f.wireBytes.Load()} if err = json.NewEncoder(os.Stdout).Encode(result); err != nil { return err } fmt.Fprintf(os.Stderr, "PASS catchup %s count=%d delay=%s repeat=%d\n", client, count, delay, rep) } } } } return nil}
func catchupSummary(output string) string { for _, line := range strings.Split(output, "\n") { if strings.HasPrefix(line, "CATCHUP ") { return line } } return ""}