···
UpdateLastTimeUs(int64) error
+
type JetstreamSubscriber struct {
+
cancel context.CancelFunc
type JetstreamClient struct {
+
cfg *client.ClientConfig
+
maxDidsPerSubscriber int
+
subscribers []*JetstreamSubscriber
+
processFunc func(context.Context, *models.Event) error
+
subscriberWg sync.WaitGroup
func (j *JetstreamClient) AddDid(did string) {
···
+
// Just add to the config for now, actual subscriber management happens in UpdateDids
j.cfg.WantedDids = append(j.cfg.WantedDids, did)
func (j *JetstreamClient) UpdateDids(dids []string) {
···
j.cfg.WantedDids = append(j.cfg.WantedDids, did)
+
needRebalance := j.processFunc != nil
+
j.rebalanceSubscribers()
+
func NewJetstreamClient(endpoint, ident string, collections []string, cfg *client.ClientConfig, logger *slog.Logger, db DB, waitForDid bool) (*JetstreamClient, error) {
cfg = client.DefaultClientConfig()
+
cfg.WebsocketURL = endpoint
cfg.WantedCollections = collections
+
waitForDid: waitForDid,
+
subscribers: make([]*JetstreamSubscriber, 0),
+
maxDidsPerSubscriber: 100,
// StartJetstream starts the jetstream client and processes events using the provided processFunc.
// The caller is responsible for saving the last time_us to the database (just use your db.SaveLastTimeUs).
func (j *JetstreamClient) StartJetstream(ctx context.Context, processFunc func(context.Context, *models.Event) error) error {
+
j.processFunc = processFunc
+
// Start a goroutine to wait for DIDs and then start subscribers
+
hasDids := len(j.cfg.WantedDids) > 0
+
j.l.Info("done waiting for did, starting subscribers")
+
j.rebalanceSubscribers()
+
time.Sleep(time.Second)
+
// Start subscribers immediately
+
j.rebalanceSubscribers()
+
// rebalanceSubscribers creates, updates, or removes subscribers based on the current list of DIDs
+
func (j *JetstreamClient) rebalanceSubscribers() {
+
if j.processFunc == nil {
+
j.l.Warn("cannot rebalance subscribers without a process function")
+
// stop all subscribers first
+
for _, sub := range j.subscribers {
+
if sub.running && sub.cancel != nil {
+
// calculate how many subscribers we need
+
totalDids := len(j.cfg.WantedDids)
+
subscribersNeeded := (totalDids + j.maxDidsPerSubscriber - 1) / j.maxDidsPerSubscriber // ceiling division
+
// create or reuse subscribers as needed
+
j.subscribers = j.subscribers[:0]
+
for i := range subscribersNeeded {
+
startIdx := i * j.maxDidsPerSubscriber
+
endIdx := min((i+1)*j.maxDidsPerSubscriber, totalDids)
+
subscriberDids := j.cfg.WantedDids[startIdx:endIdx]
+
subCfg.WantedDids = subscriberDids
+
ident := fmt.Sprintf("%s-%d", j.baseIdent, i)
+
subscriber := &JetstreamSubscriber{
+
j.subscribers = append(j.subscribers, subscriber)
+
go j.startSubscriber(subscriber, &subCfg)
+
// startSubscriber initializes and starts a single subscriber
+
func (j *JetstreamClient) startSubscriber(sub *JetstreamSubscriber, cfg *client.ClientConfig) {
+
defer j.subscriberWg.Done()
+
logger := j.l.With("subscriber", sub.ident)
+
logger.Info("starting subscriber", "dids_count", len(sub.dids))
+
sched := sequential.NewScheduler(sub.ident, logger, j.processFunc)
+
client, err := client.NewClient(cfg, log.New("jetstream-"+sub.ident), sched)
+
logger.Error("failed to create jetstream client", "error", err)
+
j.connectAndReadForSubscriber(sub)
+
func (j *JetstreamClient) connectAndReadForSubscriber(sub *JetstreamSubscriber) {
+
ctx := context.Background()
+
l := j.l.With("subscriber", sub.ident)
+
// Check if this subscriber should still be running
+
l.Info("subscriber marked for shutdown")
cursor := j.getLastTimeUs(ctx)
connCtx, cancel := context.WithCancel(ctx)
+
l.Info("connecting subscriber to jetstream")
+
if err := sub.client.ConnectAndRead(connCtx, cursor); err != nil {
l.Error("error reading jetstream", "error", err)
+
time.Sleep(time.Second) // Small backoff before retry
+
l.Info("context done, stopping subscriber")
l.Info("connection context done, reconnecting")
···
+
// GetRunningSubscribersCount returns the total number of currently running subscribers
+
func (j *JetstreamClient) GetRunningSubscribersCount() int {
+
for _, sub := range j.subscribers {
+
// Shutdown gracefully stops all subscribers
+
func (j *JetstreamClient) Shutdown() {
+
// Cancel all subscribers
+
for _, sub := range j.subscribers {
+
if sub.running && sub.cancel != nil {
+
// Wait for all subscribers to complete
+
j.l.Info("all subscribers shut down", "total_subscribers", len(j.subscribers), "running_subscribers", j.GetRunningSubscribersCount())
func (j *JetstreamClient) getLastTimeUs(ctx context.Context) *int64 {
l := log.FromContext(ctx)
lastTimeUs, err := j.db.GetLastTimeUs()
···
+
// If last time is older than 2 days, start from now
if time.Now().UnixMicro()-lastTimeUs > 2*24*60*60*1000*1000 {
lastTimeUs = time.Now().UnixMicro()
l.Warn("last time us is older than 2 days; discarding that and starting from now")
···
+
l.Info("found last time_us", "time_us", lastTimeUs, "running_subscribers", j.GetRunningSubscribersCount())
julien.rbrt.fr
anirudh.fi