Something went wrong. Try again.
Monorepo for Tangled forked from tangled.org/core
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298package jetstream
import ( "context" "fmt" "log/slog" "sync" "sync/atomic" "time"
"tangled.org/core/log")
const ( gapRewind = time.Second // shutdown saves the cursor after this even if events are still held heldGrace = 30 * time.Second)
type DB interface { GetLastTimeUs() (int64, error) SaveLastTimeUs(int64) error}
type Set[T comparable] map[T]struct{}
type JetstreamClient struct { cfg StreamConfig l *slog.Logger
logDids bool wantedDids Set[string] unfilteredNsids Set[string] db DB waitForDid bool mu sync.RWMutex
lastSeenUs atomic.Int64 marks watermark holding atomic.Int64
// the first event of each connection clears it, so the gap check runs once anchorUs atomic.Int64 onGap func(context.Context)
done chan struct{}
saveEvery time.Duration}
func (j *JetstreamClient) AddDid(did string) { if did == "" { return }
if j.logDids { j.l.Info("adding did to in-memory filter", "did", did) } j.mu.Lock() j.wantedDids[did] = struct{}{} j.mu.Unlock()}
// call before StartJetstream, fn runs on the reader so don't block in itfunc (j *JetstreamClient) SetGapHandler(fn func(context.Context)) { j.onGap = fn}
func (j *JetstreamClient) ExemptCollection(nsid string) { j.mu.Lock() j.unfilteredNsids[nsid] = struct{}{} j.mu.Unlock()}
func (j *JetstreamClient) RemoveDid(did string) { if did == "" { return }
if j.logDids { j.l.Info("removing did from in-memory filter", "did", did) } j.mu.Lock() delete(j.wantedDids, did) j.mu.Unlock()}
type processor func(context.Context, *Event) error
func (j *JetstreamClient) withDidFilter(processFunc processor) processor { // since this closure references j.WantedDids; it should auto-update // existing instances of the closure when j.WantedDids is mutated return func(ctx context.Context, evt *Event) error { if anchor := j.anchorUs.Swap(0); anchor != 0 && evt.TimeUS > anchor && j.onGap != nil { j.l.Warn("jetstream history no longer reaches the last event seen", "last", anchor, "first", evt.TimeUS) j.onGap(ctx) }
ticket := j.ticket(evt.TimeUS) h := &hold{ticket: ticket} ctx = context.WithValue(ctx, holdKey{}, h)
j.mu.RLock() // empty filter => all dids allowed matches := len(j.wantedDids) == 0 if !matches { if _, ok := j.wantedDids[evt.Did]; ok { matches = true } } if !matches && evt.Commit != nil { if _, ok := j.unfilteredNsids[evt.Commit.Collection]; ok { matches = true } } j.mu.RUnlock()
var err error if matches { err = processFunc(ctx, evt) } if err != nil && ctx.Err() != nil { // shutdown cut this event short, so the cursor stays before it and the next start replays it return err }
if !h.held.Load() { ticket.Done() } return err }}
// StartJetstream replaces cfg.Cursor with the db one, so don't set itfunc NewJetstreamClient(endpoint string, collections []string, cfg *StreamConfig, logger *slog.Logger, db DB, waitForDid, logDids bool) (*JetstreamClient, error) { if cfg == nil { cfg = &StreamConfig{Endpoint: endpoint, Collections: collections, Compress: true} }
return &JetstreamClient{ cfg: *cfg, db: db, l: logger, wantedDids: make(map[string]struct{}), unfilteredNsids: make(map[string]struct{}),
logDids: logDids,
// This will make the goroutine in StartJetstream wait until // j.wantedDids has been populated, typically using addDids. waitForDid: waitForDid,
done: make(chan struct{}),
saveEvery: time.Minute, }, nil}
// StartJetstream starts the jetstream client and processes events using the provided processFunc.// The client persists the last time_us cursor itself via the DB it was constructed with.func (j *JetstreamClient) StartJetstream(ctx context.Context, processFunc func(context.Context, *Event) error) error { cfg := j.cfg cfg.TailWhenNoCursor = true cfg.Cursor = func() (int64, bool, error) { return j.cursor(ctx) } stream, err := NewStream(cfg, j.l) if err != nil { return fmt.Errorf("failed to create jetstream client: %w", err) }
handle := j.withDidFilter(processFunc) go func() { defer close(j.done) if j.waitForDid && !j.waitForWantedDid(ctx) { return } j.l.Info("done waiting for did")
ticked := make(chan struct{}) go func() { defer close(ticked) j.periodicLastTimeSave(ctx) }()
// returning the error would replay it forever, so log and skip stream.Run(ctx, func(ctx context.Context, e *Event) error { if err := handle(ctx, e); err != nil { if ctx.Err() != nil { return err } j.l.Error("error processing event", "error", err, "time_us", e.TimeUS) } return nil }) <-ticked j.waitForHeld() // save last, nothing may move the cursor once it runs j.saveLastSeen(ctx) }()
return nil}
func (j *JetstreamClient) waitForHeld() { for deadline := time.Now().Add(heldGrace); j.holding.Load() > 0 && time.Now().Before(deadline); { time.Sleep(10 * time.Millisecond) }}
// one past the last finished event, not the event itselffunc (j *JetstreamClient) LastSeenUs() int64 { return j.lastSeenUs.Load()}
// wait on it before closing the db, the final cursor save needs it open// it never closes if StartJetstream wasn't calledfunc (j *JetstreamClient) Done() <-chan struct{} { return j.done}
func (j *JetstreamClient) waitForWantedDid(ctx context.Context) bool { for { j.mu.RLock() hasDid := len(j.wantedDids) != 0 j.mu.RUnlock() if hasDid { return true } select { case <-ctx.Done(): return false case <-time.After(time.Second): } }}
// save cursor periodicallyfunc (j *JetstreamClient) periodicLastTimeSave(ctx context.Context) { ticker := time.NewTicker(j.saveEvery) defer ticker.Stop()
for { select { case <-ctx.Done(): return case <-ticker.C: j.saveLastSeen(ctx) } }}
func (j *JetstreamClient) resume(ctx context.Context) *int64 { cursor := j.resumeCursor(ctx) j.anchorUs.Store(0) if cursor == nil || j.onGap == nil || *cursor <= 0 { return cursor } // cursor sits one past its event, hence the -1 j.anchorUs.Store(*cursor - 1) rewound := max(*cursor-gapRewind.Microseconds(), 0) return &rewound}
func (j *JetstreamClient) cursor(ctx context.Context) (int64, bool, error) { cursor := j.resume(ctx) if cursor == nil { return 0, false, nil } return *cursor, true, nil}
func (j *JetstreamClient) resumeCursor(ctx context.Context) *int64 { if seen := j.lastSeenUs.Load(); seen != 0 { return &seen } return j.getLastTimeUs(ctx)}
// the saved cursor is always an event time_us, never a wall clockfunc (j *JetstreamClient) getLastTimeUs(ctx context.Context) *int64 { l := log.FromContext(ctx) lastTimeUs, err := j.db.GetLastTimeUs() if err != nil { l.Warn("couldn't get last time us, starting from now", "error", err) return nil }
l.Info("found last time_us", "time_us", lastTimeUs) return &lastTimeUs}
func (j *JetstreamClient) saveLastSeen(ctx context.Context) { if seen := j.lastSeenUs.Load(); seen != 0 { if err := j.db.SaveLastTimeUs(seen); err != nil { log.FromContext(ctx).Error("failed to save cursor", "error", err) } }}