···
UpdateLastTimeUs(int64) error
22
+
type JetstreamSubscriber struct {
23
+
client *client.Client
24
+
cancel context.CancelFunc
type JetstreamClient struct {
23
-
cfg *client.ClientConfig
24
-
client *client.Client
31
+
cfg *client.ClientConfig
36
+
maxDidsPerSubscriber int
32
-
cancel context.CancelFunc
39
+
subscribers []*JetstreamSubscriber
40
+
processFunc func(context.Context, *models.Event) error
41
+
subscriberWg sync.WaitGroup
func (j *JetstreamClient) AddDid(did string) {
···
51
+
// 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)
63
+
needRebalance := j.processFunc != nil
55
-
if j.cancel != nil {
67
+
j.rebalanceSubscribers()
61
-
func NewJetstreamClient(ident string, collections []string, cfg *client.ClientConfig, logger *slog.Logger, db DB, waitForDid bool) (*JetstreamClient, error) {
71
+
func NewJetstreamClient(endpoint, ident string, collections []string, cfg *client.ClientConfig, logger *slog.Logger, db DB, waitForDid bool) (*JetstreamClient, error) {
cfg = client.DefaultClientConfig()
64
-
cfg.WebsocketURL = "wss://jetstream1.us-west.bsky.network/subscribe"
74
+
cfg.WebsocketURL = endpoint
cfg.WantedCollections = collections
74
-
// This will make the goroutine in StartJetstream wait until
75
-
// cfg.WantedDids has been populated, typically using UpdateDids.
76
-
waitForDid: waitForDid,
83
+
waitForDid: waitForDid,
84
+
subscribers: make([]*JetstreamSubscriber, 0),
85
+
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 {
93
+
j.processFunc = processFunc
97
+
// Start a goroutine to wait for DIDs and then start subscribers
101
+
hasDids := len(j.cfg.WantedDids) > 0
105
+
j.l.Info("done waiting for did, starting subscribers")
106
+
j.rebalanceSubscribers()
109
+
time.Sleep(time.Second)
113
+
// Start subscribers immediately
114
+
j.rebalanceSubscribers()
120
+
// rebalanceSubscribers creates, updates, or removes subscribers based on the current list of DIDs
121
+
func (j *JetstreamClient) rebalanceSubscribers() {
123
+
defer j.mu.Unlock()
125
+
if j.processFunc == nil {
126
+
j.l.Warn("cannot rebalance subscribers without a process function")
130
+
// stop all subscribers first
131
+
for _, sub := range j.subscribers {
132
+
if sub.running && sub.cancel != nil {
134
+
sub.running = false
138
+
// calculate how many subscribers we need
139
+
totalDids := len(j.cfg.WantedDids)
140
+
subscribersNeeded := (totalDids + j.maxDidsPerSubscriber - 1) / j.maxDidsPerSubscriber // ceiling division
142
+
// create or reuse subscribers as needed
143
+
j.subscribers = j.subscribers[:0]
145
+
for i := range subscribersNeeded {
146
+
startIdx := i * j.maxDidsPerSubscriber
147
+
endIdx := min((i+1)*j.maxDidsPerSubscriber, totalDids)
149
+
subscriberDids := j.cfg.WantedDids[startIdx:endIdx]
152
+
subCfg.WantedDids = subscriberDids
154
+
ident := fmt.Sprintf("%s-%d", j.baseIdent, i)
155
+
subscriber := &JetstreamSubscriber{
156
+
dids: subscriberDids,
159
+
j.subscribers = append(j.subscribers, subscriber)
161
+
j.subscriberWg.Add(1)
162
+
go j.startSubscriber(subscriber, &subCfg)
85
-
sched := sequential.NewScheduler(j.ident, logger, processFunc)
166
+
// startSubscriber initializes and starts a single subscriber
167
+
func (j *JetstreamClient) startSubscriber(sub *JetstreamSubscriber, cfg *client.ClientConfig) {
168
+
defer j.subscriberWg.Done()
87
-
client, err := client.NewClient(j.cfg, log.New("jetstream"), sched)
170
+
logger := j.l.With("subscriber", sub.ident)
171
+
logger.Info("starting subscriber", "dids_count", len(sub.dids))
173
+
sched := sequential.NewScheduler(sub.ident, logger, j.processFunc)
175
+
client, err := client.NewClient(cfg, log.New("jetstream-"+sub.ident), sched)
89
-
return fmt.Errorf("failed to create jetstream client: %w", err)
177
+
logger.Error("failed to create jetstream client", "error", err)
95
-
for len(j.cfg.WantedDids) == 0 {
96
-
time.Sleep(time.Second)
99
-
logger.Info("done waiting for did")
100
-
j.connectAndRead(ctx)
181
+
sub.client = client
187
+
j.connectAndReadForSubscriber(sub)
106
-
func (j *JetstreamClient) connectAndRead(ctx context.Context) {
107
-
l := log.FromContext(ctx)
190
+
func (j *JetstreamClient) connectAndReadForSubscriber(sub *JetstreamSubscriber) {
191
+
ctx := context.Background()
192
+
l := j.l.With("subscriber", sub.ident)
195
+
// Check if this subscriber should still be running
197
+
running := sub.running
201
+
l.Info("subscriber marked for shutdown")
cursor := j.getLastTimeUs(ctx)
connCtx, cancel := context.WithCancel(ctx)
114
-
j.cancelMu.Unlock()
210
+
sub.cancel = cancel
116
-
if err := j.client.ConnectAndRead(connCtx, cursor); err != nil {
213
+
l.Info("connecting subscriber to jetstream")
214
+
if err := sub.client.ConnectAndRead(connCtx, cursor); err != nil {
l.Error("error reading jetstream", "error", err)
217
+
time.Sleep(time.Second) // Small backoff before retry
124
-
l.Info("context done, stopping jetstream")
223
+
l.Info("context done, stopping subscriber")
l.Info("connection context done, reconnecting")
···
232
+
// GetRunningSubscribersCount returns the total number of currently running subscribers
233
+
func (j *JetstreamClient) GetRunningSubscribersCount() int {
235
+
defer j.mu.RUnlock()
238
+
for _, sub := range j.subscribers {
244
+
return runningCount
247
+
// Shutdown gracefully stops all subscribers
248
+
func (j *JetstreamClient) Shutdown() {
251
+
// Cancel all subscribers
252
+
for _, sub := range j.subscribers {
253
+
if sub.running && sub.cancel != nil {
255
+
sub.running = false
261
+
// Wait for all subscribers to complete
262
+
j.subscriberWg.Wait()
263
+
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()
···
145
-
// If last time is older than a week, start from now
278
+
// 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")
···
155
-
l.Info("found last time_us", "time_us", lastTimeUs)
288
+
l.Info("found last time_us", "time_us", lastTimeUs, "running_subscribers", j.GetRunningSubscribersCount())
kot.pink
anirudh.fi