jetstream: subscriber pool to support many watched dids

Introduces rebalancing of dids across multiple subscribers if we exceed
a threshold of 100 wantedDids.

Turns out jetstream has a limit that we hit with our 300+ subscriptions
in a single query.

jetstream/jetstream.go

···
 	UpdateLastTimeUs(int64) error
 }
 
+type JetstreamSubscriber struct {
+	client  *client.Client
+	cancel  context.CancelFunc
+	dids    []string
+	ident   string
+	running bool
+}
+
 type JetstreamClient struct {
-	cfg    *client.ClientConfig
-	client *client.Client
-	ident  string
-	l      *slog.Logger
-
-	db         DB
-	waitForDid bool
-	mu         sync.RWMutex
+	cfg                  *client.ClientConfig
+	baseIdent            string
+	l                    *slog.Logger
+	db                   DB
+	waitForDid           bool
+	maxDidsPerSubscriber int
 
-	cancel   context.CancelFunc
-	cancelMu sync.Mutex
+	mu           sync.RWMutex
+	subscribers  []*JetstreamSubscriber
+	processFunc  func(context.Context, *models.Event) error
+	subscriberWg sync.WaitGroup
 }
 
 func (j *JetstreamClient) AddDid(did string) {
···
 		return
 	}
 	j.mu.Lock()
+	defer j.mu.Unlock()
+
+	// Just add to the config for now, actual subscriber management happens in UpdateDids
 	j.cfg.WantedDids = append(j.cfg.WantedDids, did)
-	j.mu.Unlock()
 }
 
 func (j *JetstreamClient) UpdateDids(dids []string) {
···
 			j.cfg.WantedDids = append(j.cfg.WantedDids, did)
 		}
 	}
+
+	needRebalance := j.processFunc != nil
 	j.mu.Unlock()
 
-	j.cancelMu.Lock()
-	if j.cancel != nil {
-		j.cancel()
+	if needRebalance {
+		j.rebalanceSubscribers()
 	}
-	j.cancelMu.Unlock()
 }
 
-func NewJetstreamClient(ident string, collections []string, cfg *client.ClientConfig, logger *slog.Logger, db DB, waitForDid bool) (*JetstreamClient, error) {
+func NewJetstreamClient(endpoint, ident string, collections []string, cfg *client.ClientConfig, logger *slog.Logger, db DB, waitForDid bool) (*JetstreamClient, error) {
 	if cfg == nil {
 		cfg = client.DefaultClientConfig()
-		cfg.WebsocketURL = "wss://jetstream1.us-west.bsky.network/subscribe"
+		cfg.WebsocketURL = endpoint
 		cfg.WantedCollections = collections
 	}
 
 	return &JetstreamClient{
-		cfg:   cfg,
-		ident: ident,
-		db:    db,
-		l:     logger,
-
-		// This will make the goroutine in StartJetstream wait until
-		// cfg.WantedDids has been populated, typically using UpdateDids.
-		waitForDid: waitForDid,
+		cfg:                  cfg,
+		baseIdent:            ident,
+		db:                   db,
+		l:                    logger,
+		waitForDid:           waitForDid,
+		subscribers:          make([]*JetstreamSubscriber, 0),
+		maxDidsPerSubscriber: 100,
 	}, nil
 }
 
 // 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 {
-	logger := j.l
+	j.mu.Lock()
+	j.processFunc = processFunc
+	j.mu.Unlock()
+
+	if j.waitForDid {
+		// Start a goroutine to wait for DIDs and then start subscribers
+		go func() {
+			for {
+				j.mu.RLock()
+				hasDids := len(j.cfg.WantedDids) > 0
+				j.mu.RUnlock()
+
+				if hasDids {
+					j.l.Info("done waiting for did, starting subscribers")
+					j.rebalanceSubscribers()
+					return
+				}
+				time.Sleep(time.Second)
+			}
+		}()
+	} else {
+		// Start subscribers immediately
+		j.rebalanceSubscribers()
+	}
+
+	return nil
+}
+
+// rebalanceSubscribers creates, updates, or removes subscribers based on the current list of DIDs
+func (j *JetstreamClient) rebalanceSubscribers() {
+	j.mu.Lock()
+	defer j.mu.Unlock()
+
+	if j.processFunc == nil {
+		j.l.Warn("cannot rebalance subscribers without a process function")
+		return
+	}
+
+	// stop all subscribers first
+	for _, sub := range j.subscribers {
+		if sub.running && sub.cancel != nil {
+			sub.cancel()
+			sub.running = false
+		}
+	}
+
+	// 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 := *j.cfg
+		subCfg.WantedDids = subscriberDids
+
+		ident := fmt.Sprintf("%s-%d", j.baseIdent, i)
+		subscriber := &JetstreamSubscriber{
+			dids:  subscriberDids,
+			ident: ident,
+		}
+		j.subscribers = append(j.subscribers, subscriber)
+
+		j.subscriberWg.Add(1)
+		go j.startSubscriber(subscriber, &subCfg)
+	}
+}
 
-	sched := sequential.NewScheduler(j.ident, logger, processFunc)
+// startSubscriber initializes and starts a single subscriber
+func (j *JetstreamClient) startSubscriber(sub *JetstreamSubscriber, cfg *client.ClientConfig) {
+	defer j.subscriberWg.Done()
 
-	client, err := client.NewClient(j.cfg, log.New("jetstream"), sched)
+	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)
 	if err != nil {
-		return fmt.Errorf("failed to create jetstream client: %w", err)
+		logger.Error("failed to create jetstream client", "error", err)
+		return
 	}
-	j.client = client
 
-	go func() {
-		if j.waitForDid {
-			for len(j.cfg.WantedDids) == 0 {
-				time.Sleep(time.Second)
-			}
-		}
-		logger.Info("done waiting for did")
-		j.connectAndRead(ctx)
-	}()
+	sub.client = client
+
+	j.mu.Lock()
+	sub.running = true
+	j.mu.Unlock()
 
-	return nil
+	j.connectAndReadForSubscriber(sub)
 }
 
-func (j *JetstreamClient) connectAndRead(ctx context.Context) {
-	l := log.FromContext(ctx)
+func (j *JetstreamClient) connectAndReadForSubscriber(sub *JetstreamSubscriber) {
+	ctx := context.Background()
+	l := j.l.With("subscriber", sub.ident)
+
 	for {
+		// Check if this subscriber should still be running
+		j.mu.RLock()
+		running := sub.running
+		j.mu.RUnlock()
+
+		if !running {
+			l.Info("subscriber marked for shutdown")
+			return
+		}
+
 		cursor := j.getLastTimeUs(ctx)
 
 		connCtx, cancel := context.WithCancel(ctx)
-		j.cancelMu.Lock()
-		j.cancel = cancel
-		j.cancelMu.Unlock()
+
+		j.mu.Lock()
+		sub.cancel = cancel
+		j.mu.Unlock()
 
-		if err := j.client.ConnectAndRead(connCtx, cursor); err != nil {
+		l.Info("connecting subscriber to jetstream")
+		if err := sub.client.ConnectAndRead(connCtx, cursor); err != nil {
 			l.Error("error reading jetstream", "error", err)
 			cancel()
+			time.Sleep(time.Second) // Small backoff before retry
 			continue
 		}
 
 		select {
 		case <-ctx.Done():
-			l.Info("context done, stopping jetstream")
+			l.Info("context done, stopping subscriber")
 			return
 		case <-connCtx.Done():
 			l.Info("connection context done, reconnecting")
···
 	}
 }
 
+// GetRunningSubscribersCount returns the total number of currently running subscribers
+func (j *JetstreamClient) GetRunningSubscribersCount() int {
+	j.mu.RLock()
+	defer j.mu.RUnlock()
+
+	runningCount := 0
+	for _, sub := range j.subscribers {
+		if sub.running {
+			runningCount++
+		}
+	}
+
+	return runningCount
+}
+
+// Shutdown gracefully stops all subscribers
+func (j *JetstreamClient) Shutdown() {
+	j.mu.Lock()
+
+	// Cancel all subscribers
+	for _, sub := range j.subscribers {
+		if sub.running && sub.cancel != nil {
+			sub.cancel()
+			sub.running = false
+		}
+	}
+
+	j.mu.Unlock()
+
+	// Wait for all subscribers to complete
+	j.subscriberWg.Wait()
+	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 a week, start from now
+	// 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)
+	l.Info("found last time_us", "time_us", lastTimeUs, "running_subscribers", j.GetRunningSubscribersCount())
 	return &lastTimeUs
 }