Something went wrong. Try again.
Monorepo for Tangled tangled.org
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479package 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 onceconst 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 notfunc 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()) }}