Something went wrong. Try again.
Monorepo for Tangled
Something went wrong. Try again.
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274package jetstream
import ( "context" "fmt" "log/slog" "os" "os/signal" "sync" "sync/atomic" "syscall" "time"
"github.com/bluesky-social/jetstream/pkg/client" "github.com/bluesky-social/jetstream/pkg/client/schedulers/sequential" "github.com/bluesky-social/jetstream/pkg/models" "tangled.org/core/log")
type DB interface { GetLastTimeUs() (int64, error) SaveLastTimeUs(int64) error}
type Set[T comparable] map[T]struct{}
type JetstreamClient struct { cfg *client.ClientConfig client *client.Client ident string l *slog.Logger
logDids bool wantedDids Set[string] unfilteredNsids Set[string] db DB waitForDid bool mu sync.RWMutex
lastSeenUs atomic.Int64
cancel context.CancelFunc cancelMu sync.Mutex}
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()}
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, *models.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 *models.Event) error { 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) }
j.lastSeenUs.Store(evt.TimeUS + 1) return err }}
func NewJetstreamClient(endpoint, ident string, collections []string, cfg *client.ClientConfig, logger *slog.Logger, db DB, waitForDid, logDids bool) (*JetstreamClient, error) { if cfg == nil { cfg = client.DefaultClientConfig() cfg.WebsocketURL = endpoint cfg.WantedCollections = collections }
return &JetstreamClient{ cfg: cfg, ident: ident, 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, }, 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, *models.Event) error) error { logger := j.l
sched := sequential.NewScheduler(j.ident, logger, j.withDidFilter(processFunc))
client, err := client.NewClient(j.cfg, logger, sched) if err != nil { return fmt.Errorf("failed to create jetstream client: %w", err) } j.client = client
go func() { if j.waitForDid { for { j.mu.RLock() hasDid := len(j.wantedDids) != 0 j.mu.RUnlock() if hasDid { break } time.Sleep(time.Second) } } logger.Info("done waiting for did")
go j.periodicLastTimeSave(ctx) j.saveIfKilled(ctx)
j.connectAndRead(ctx) }()
return nil}
func (j *JetstreamClient) connectAndRead(ctx context.Context) { l := log.FromContext(ctx) for { cursor := j.resumeCursor(ctx)
connCtx, cancel := context.WithCancel(ctx) j.cancelMu.Lock() j.cancel = cancel j.cancelMu.Unlock()
if err := j.client.ConnectAndRead(connCtx, cursor); err != nil { l.Error("error reading jetstream", "error", err) cancel() continue }
select { case <-ctx.Done(): l.Info("context done, stopping jetstream") return case <-connCtx.Done(): l.Info("connection context done, reconnecting") continue } }}
// save cursor periodicallyfunc (j *JetstreamClient) periodicLastTimeSave(ctx context.Context) { ticker := time.NewTicker(time.Minute) defer ticker.Stop()
for { select { case <-ctx.Done(): return case <-ticker.C: 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) } } } }}
func (j *JetstreamClient) resumeCursor(ctx context.Context) *int64 { if seen := j.lastSeenUs.Load(); seen != 0 { return &seen } return j.getLastTimeUs(ctx)}
func (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) lastTimeUs = time.Now().UnixMicro() if err = j.db.SaveLastTimeUs(lastTimeUs); err != nil { l.Error("failed to save last time us", "error", err) } }
l.Info("found last time_us", "time_us", lastTimeUs) return &lastTimeUs}
func (j *JetstreamClient) saveIfKilled(ctx context.Context) context.Context { ctxWithCancel, cancel := context.WithCancel(ctx)
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGHUP, syscall.SIGKILL, syscall.SIGSTOP, )
go func() { sig := <-sigChan j.l.Info("Received signal, initiating graceful shutdown", "signal", sig)
if seen := j.lastSeenUs.Load(); seen != 0 { if err := j.db.SaveLastTimeUs(seen); err != nil { j.l.Error("Failed to save last time during shutdown", "error", err) } j.l.Info("Saved lastTimeUs before shutdown", "lastTimeUs", seen) }
j.cancelMu.Lock() if j.cancel != nil { j.cancel() } j.cancelMu.Unlock()
cancel()
os.Exit(0) }()
return ctxWithCancel}