package jetstream
import (
"context"
"errors"
"fmt"
"log/slog"
"net/http"
"net/http/httptest"
"slices"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/gorilla/websocket"
"github.com/klauspost/compress/zstd"
)
type fakeCursorDB struct {
saved int64
savedErr error
mu sync.Mutex
got int64
}
func (f *fakeCursorDB) GetLastTimeUs() (int64, error) { return f.saved, f.savedErr }
func (f *fakeCursorDB) SaveLastTimeUs(us int64) error {
f.mu.Lock()
defer f.mu.Unlock()
f.got = us
return nil
}
func cursorFor(db DB) int64 {
j := &JetstreamClient{db: db, wantedDids: make(Set[string])}
return *j.getLastTimeUs(context.Background())
}
func newGapClient(db DB, gaps *int) *JetstreamClient {
j := &JetstreamClient{db: db, l: slog.Default(), wantedDids: make(Set[string])}
j.SetGapHandler(func(context.Context) { *gaps++ })
return j
}
const twoDaysUs = int64(2 * 24 * 60 * 60 * 1000 * 1000)
func TestStaleCursorIsHonored(t *testing.T) {
old := time.Now().UnixMicro() - 5*twoDaysUs
if got := cursorFor(&fakeCursorDB{saved: old}); got != old {
t.Fatalf("stale cursor must be honored, not snapped forward: got %d, want %d", got, old)
}
}
func TestZeroCursorIsHonored(t *testing.T) {
if got := cursorFor(&fakeCursorDB{saved: 0}); got != 0 {
t.Fatalf("cursor 0 must replay from the start: got %d, want 0", got)
}
}
func TestMissingCursorConnectsLive(t *testing.T) {
var gaps int
j := newGapClient(&fakeCursorDB{savedErr: errors.New("no row")}, &gaps)
if got := j.resume(context.Background()); got != nil {
t.Fatalf("missing cursor must connect live, got %d", *got)
}
_ = j.withDidFilter(func(context.Context, *Event) error { return nil })(context.Background(), &Event{Did: "did:plc:a", TimeUS: 5})
if gaps != 0 {
t.Fatal("a first boot has nothing to catch up on")
}
}
func TestResumeRewindsASecondFromTheLastEvent(t *testing.T) {
var gaps int
j := newGapClient(&fakeCursorDB{saved: 10_000_001}, &gaps)
if got := j.resume(context.Background()); got == nil || *got != 9_000_001 {
t.Fatalf("cursor = %v, want 9000001", got)
}
}
func TestResumeKeepsTheStoredCursorWithoutAGapHandler(t *testing.T) {
j := &JetstreamClient{db: &fakeCursorDB{saved: 10_000_001}, l: slog.Default(), wantedDids: make(Set[string])}
if got := j.resume(context.Background()); got == nil || *got != 10_000_001 {
t.Fatalf("cursor = %v, want the stored one", got)
}
}
func firstEventAfterResume(t *testing.T, saved, first int64) int {
t.Helper()
var gaps int
j := newGapClient(&fakeCursorDB{saved: saved}, &gaps)
j.resume(context.Background())
wrapped := j.withDidFilter(func(context.Context, *Event) error { return nil })
_ = wrapped(context.Background(), &Event{Did: "did:plc:a", TimeUS: first})
_ = wrapped(context.Background(), &Event{Did: "did:plc:a", TimeUS: first + 1_000_000})
return gaps
}
func TestAGapIsReportedWhenTheStreamStartsAfterTheLastEvent(t *testing.T) {
// last event was 10_000_000, so the stored cursor is 10_000_001
if got := firstEventAfterResume(t, 10_000_001, 10_000_001); got != 1 {
t.Fatalf("gap handler ran %d times, want once", got)
}
}
func TestNoGapWhenTheStreamReplaysTheLastEvent(t *testing.T) {
for name, first := range map[string]int64{
"the last event itself": 10_000_000,
"an earlier event": 9_500_000,
} {
if got := firstEventAfterResume(t, 10_000_001, first); got != 0 {
t.Fatalf("%s: gap handler ran %d times, want none", name, got)
}
}
}
func TestLastSeenTracksEveryEventPreFilter(t *testing.T) {
j := &JetstreamClient{wantedDids: Set[string]{"did:plc:boltless": {}}}
wrapped := j.withDidFilter(func(context.Context, *Event) error { return nil })
_ = wrapped(context.Background(), &Event{Did: "did:plc:akshay", TimeUS: 100})
if got := j.lastSeenUs.Load(); got != 101 {
t.Fatalf("filtered-out event must still advance last-seen: got %d, want 101", got)
}
_ = wrapped(context.Background(), &Event{Did: "did:plc:boltless", TimeUS: 200})
if got := j.lastSeenUs.Load(); got != 201 {
t.Fatalf("matching event must advance last-seen: got %d, want 201", got)
}
}
func TestExemptCollectionShouldBypassDidFilter(t *testing.T) {
j := &JetstreamClient{
wantedDids: Set[string]{"did:plc:boltless": {}},
unfilteredNsids: make(Set[string]),
}
j.ExemptCollection("sh.tangled.repo")
var processed int
wrapped := j.withDidFilter(func(context.Context, *Event) error {
processed++
return nil
})
_ = wrapped(context.Background(), &Event{
Did: "did:plc:akshay", TimeUS: 100,
Commit: &Commit{Collection: "sh.tangled.repo"},
})
if processed != 1 {
t.Fatal("An exempt collection must reach this processor from an unwanted DID; enrollment is decided later, not by this filter.")
}
_ = wrapped(context.Background(), &Event{
Did: "did:plc:akshay", TimeUS: 200,
Commit: &Commit{Collection: "sh.tangled.repo.pull"},
})
if processed != 1 {
t.Fatal("A non-exempt collection must still be filtered by DID.")
}
}
func TestLastSeenAdvancesAfterProcessing(t *testing.T) {
j := &JetstreamClient{wantedDids: make(Set[string])}
var seenDuringProcess int64
wrapped := j.withDidFilter(func(context.Context, *Event) error {
seenDuringProcess = j.lastSeenUs.Load()
return nil
})
_ = wrapped(context.Background(), &Event{Did: "did:plc:boltless", TimeUS: 500})
if seenDuringProcess == 501 {
t.Fatal("cursor advanced before the event finished processing; a crash mid-process would skip it")
}
if got := j.lastSeenUs.Load(); got != 501 {
t.Fatalf("cursor must advance once processing returns: got %d, want 501", got)
}
}
func TestShutdownLeavesTheCursorBeforeTheCutEvent(t *testing.T) {
j := &JetstreamClient{wantedDids: make(Set[string])}
ctx, cancel := context.WithCancel(context.Background())
wrapped := j.withDidFilter(func(ctx context.Context, _ *Event) error {
cancel()
return ctx.Err()
})
_ = wrapped(ctx, &Event{Did: "did:plc:boltless", TimeUS: 700})
if got := j.lastSeenUs.Load(); got != 0 {
t.Fatalf("cursor moved past an event shutdown cut short: got %d, want it untouched", got)
}
}
func TestAFailingEventStillAdvancesTheCursor(t *testing.T) {
j := &JetstreamClient{wantedDids: make(Set[string])}
wrapped := j.withDidFilter(func(context.Context, *Event) error {
return errors.New("bad record")
})
_ = wrapped(context.Background(), &Event{Did: "did:plc:boltless", TimeUS: 800})
if got := j.lastSeenUs.Load(); got != 801 {
t.Fatalf("a record that fails on its own must not stall the cursor: got %d, want 801", got)
}
}
// nothing listens on port 1, so every connect fails at once
const refusedEndpoint = "ws://127.0.0.1:1/subscribe"
func startClient(t *testing.T, endpoint string, db DB, waitForDid bool) (*JetstreamClient, context.CancelFunc) {
t.Helper()
j, err := NewJetstreamClient(endpoint, nil, nil, slog.Default(), db, waitForDid, false)
if err != nil {
t.Fatal(err)
}
j.lastSeenUs.Store(42)
ctx, cancel := context.WithCancel(context.Background())
if err := j.StartJetstream(ctx, func(context.Context, *Event) error { return nil }); err != nil {
t.Fatal(err)
}
return j, cancel
}
func requireStopsWithSavedCursor(t *testing.T, j *JetstreamClient, db *fakeCursorDB, msg string) {
t.Helper()
select {
case <-j.Done():
case <-time.After(5 * time.Second):
t.Fatal(msg)
}
db.mu.Lock()
defer db.mu.Unlock()
if db.got != 42 {
t.Fatalf("cursor saved at stop = %d, want 42", db.got)
}
}
func TestDoneClosesOnlyAfterTheFinalCursorSave(t *testing.T) {
db := &fakeCursorDB{savedErr: errors.New("no row")}
j, cancel := startClient(t, refusedEndpoint, db, false)
// long enough to be waiting out a failed connect when the cancel lands
time.Sleep(50 * time.Millisecond)
cancel()
requireStopsWithSavedCursor(t, j, db, "the reconnect loop kept going after its context was cancelled")
}
func TestDoneClosesWhenCancelledBeforeAnyDidArrives(t *testing.T) {
j, cancel := startClient(t, refusedEndpoint, &fakeCursorDB{}, true)
cancel()
select {
case <-j.Done():
case <-time.After(5 * time.Second):
t.Fatal("a client still waiting for its first did ignored shutdown")
}
}
func TestDoneClosesEvenWhenTheStreamIsQuiet(t *testing.T) {
connected := make(chan struct{})
release := make(chan struct{})
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
if err != nil {
return
}
defer conn.Close()
close(connected)
<-release
}))
defer srv.Close()
defer close(release)
db := &fakeCursorDB{savedErr: errors.New("no row")}
j, cancel := startClient(t, "ws"+strings.TrimPrefix(srv.URL, "http")+"/subscribe", db, false)
<-connected
cancel()
requireStopsWithSavedCursor(t, j, db, "a stream with no traffic held shutdown hostage, so the final cursor save never ran")
}
func frameServer(t *testing.T, frames ...string) *streamServer {
t.Helper()
enc, err := zstd.NewWriter(nil, zstd.WithEncoderDict(zstdDictionary))
if err != nil {
t.Fatal(err)
}
var srv *streamServer
srv = newStreamServer(t, func(n int, con *websocket.Conn) {
srv.mu.Lock()
compressed := srv.headers[n].Get("Socket-Encoding") == "zstd"
srv.mu.Unlock()
for _, f := range frames {
if compressed {
con.WriteMessage(websocket.BinaryMessage, enc.EncodeAll([]byte(f), nil))
} else {
con.WriteMessage(websocket.TextMessage, []byte(f))
}
}
drain(con)
})
return srv
}
func commitFrame(did string, us int, collection string, live bool) string {
return fmt.Sprintf(`{"did":%q,"time_us":%d,"kind":"commit","commit":{"operation":"create","collection":%q,"rkey":"k","live":%t}}`, did, us, collection, live)
}
func runClient(t *testing.T, j *JetstreamClient, stop func(*Event) bool) {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := j.StartJetstream(ctx, func(_ context.Context, e *Event) error {
if stop(e) {
cancel()
}
return nil
}); err != nil {
t.Fatal(err)
}
<-j.Done()
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
t.Fatal("the client never delivered what the test waited for")
}
}
func newTestClient(t *testing.T, srv *streamServer, db DB, collections ...string) *JetstreamClient {
t.Helper()
j, err := NewJetstreamClient(srv.endpoint(), collections, nil, slog.Default(), db, false, false)
if err != nil {
t.Fatal(err)
}
return j
}
func TestClientDeliversInOrderAndOneAtATime(t *testing.T) {
var frames []string
for i := 1; i <= 50; i++ {
frames = append(frames, commitFrame("did:plc:a", i, "c", true))
}
j := newTestClient(t, frameServer(t, frames...), &fakeCursorDB{savedErr: errors.New("no row")})
var inflight atomic.Int32
var got []int64
runClient(t, j, func(e *Event) bool {
if inflight.Add(1) != 1 {
t.Error("two events were being handled at once")
}
time.Sleep(time.Millisecond)
inflight.Add(-1)
got = append(got, e.TimeUS)
return e.TimeUS == 50
})
for i, us := range got {
if us != int64(i+1) {
t.Fatalf("event %d arrived as time_us %d, want %d", i, us, i+1)
}
}
}
func TestClientDidFilterChangesAtRuntime(t *testing.T) {
srv := frameServer(t,
commitFrame("did:plc:a", 1, "c", true),
commitFrame("did:plc:b", 2, "c", true),
commitFrame("did:plc:a", 3, "c", true),
commitFrame("did:plc:b", 4, "c", true),
commitFrame("did:plc:a", 5, "c", true),
)
j := newTestClient(t, srv, &fakeCursorDB{savedErr: errors.New("no row")})
j.AddDid("did:plc:b")
var got []int64
runClient(t, j, func(e *Event) bool {
got = append(got, e.TimeUS)
if e.TimeUS == 2 {
j.AddDid("did:plc:a")
j.RemoveDid("did:plc:b")
}
return e.TimeUS == 5
})
if want := []int64{2, 3, 5}; !slices.Equal(got, want) {
t.Fatalf("delivered %v, want %v", got, want)
}
}
func TestClientUnfilteredCollectionIgnoresTheDidFilter(t *testing.T) {
srv := frameServer(t,
commitFrame("did:plc:a", 1, "other", true),
commitFrame("did:plc:a", 2, "open", true),
commitFrame("did:plc:b", 3, "other", true),
)
j := newTestClient(t, srv, &fakeCursorDB{savedErr: errors.New("no row")})
j.AddDid("did:plc:b")
j.ExemptCollection("open")
var got []int64
runClient(t, j, func(e *Event) bool {
got = append(got, e.TimeUS)
return e.TimeUS == 3
})
if len(got) != 2 || got[0] != 2 || got[1] != 3 {
t.Fatalf("delivered %v, want the exempt event 2 and the wanted did's event 3", got)
}
}
// only deliberi sends wantedEventTypes, a plain jetstream client must not
func TestClientAsksForNoEventTypes(t *testing.T) {
srv := frameServer(t, commitFrame("did:plc:a", 1, "c", true))
j := newTestClient(t, srv, &fakeCursorDB{savedErr: errors.New("no row")}, "c")
runClient(t, j, func(e *Event) bool { return true })
if q := srv.query(0); q.Has("wantedEventTypes") {
t.Fatalf("query = %v, want no wantedEventTypes", q)
}
}
func TestClientSavesTheCursorOnTheTimer(t *testing.T) {
srv := frameServer(t, commitFrame("did:plc:a", 10, "c", true))
db := &fakeCursorDB{savedErr: errors.New("no row")}
j := newTestClient(t, srv, db)
j.saveEvery = 10 * time.Millisecond
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
if err := j.StartJetstream(ctx, func(context.Context, *Event) error { return nil }); err != nil {
t.Fatal(err)
}
deadline := time.Now().Add(5 * time.Second)
for {
db.mu.Lock()
got := db.got
db.mu.Unlock()
if got == 11 {
return
}
if time.Now().After(deadline) {
t.Fatalf("saved cursor = %d, want 11 from the timer before any shutdown", got)
}
time.Sleep(5 * time.Millisecond)
}
}
func TestClientReportsAGapAndResumesARewoundCursor(t *testing.T) {
// cursor is one past 10_000_000, but the server's history starts after it
srv := frameServer(t, commitFrame("did:plc:a", 20_000_000, "c", true))
j := newTestClient(t, srv, &fakeCursorDB{saved: 10_000_001})
gaps := make(chan struct{}, 2)
j.SetGapHandler(func(context.Context) { gaps <- struct{}{} })
runClient(t, j, func(e *Event) bool { return true })
if len(gaps) != 1 {
t.Fatalf("gap handler ran %d times, want once", len(gaps))
}
if got := srv.query(0).Get("cursor"); got != "9000001" {
t.Fatalf("connected with cursor %q, want the stored one rewound a second", got)
}
}
func TestClientFailingEventsDontDropTheConnection(t *testing.T) {
srv := frameServer(t,
commitFrame("did:plc:a", 1, "c", true),
commitFrame("did:plc:a", 2, "c", true),
)
j := newTestClient(t, srv, &fakeCursorDB{savedErr: errors.New("no row")})
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
var seen atomic.Int32
_ = j.StartJetstream(ctx, func(_ context.Context, e *Event) error {
if seen.Add(1) == 2 {
cancel()
}
return errors.New("bad record")
})
<-j.Done()
if seen.Load() != 2 || srv.connects() != 1 {
t.Fatalf("saw %d events over %d connections, want 2 over 1", seen.Load(), srv.connects())
}
}