···
UpdateLastTimeUs(int64) error
22
-
type JetstreamSubscriber struct {
23
-
client *client.Client
24
-
cancel context.CancelFunc
type JetstreamClient struct {
31
-
cfg *client.ClientConfig
36
-
maxDidsPerSubscriber int
23
+
cfg *client.ClientConfig
24
+
client *client.Client
39
-
subscribers []*JetstreamSubscriber
40
-
processFunc func(context.Context, *models.Event) error
41
-
subscriberWg sync.WaitGroup
32
+
cancel context.CancelFunc
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
67
-
j.rebalanceSubscribers()
55
+
if j.cancel != nil {
func NewJetstreamClient(endpoint, ident string, collections []string, cfg *client.ClientConfig, logger *slog.Logger, db DB, waitForDid bool) (*JetstreamClient, error) {
···
83
-
waitForDid: waitForDid,
84
-
subscribers: make([]*JetstreamSubscriber, 0),
85
-
maxDidsPerSubscriber: 100,
74
+
// This will make the goroutine in StartJetstream wait until
75
+
// cfg.WantedDids has been populated, typically using UpdateDids.
76
+
waitForDid: waitForDid,
// 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
-
// calculate how many subscribers we need
131
-
totalDids := len(j.cfg.WantedDids)
132
-
subscribersNeeded := (totalDids + j.maxDidsPerSubscriber - 1) / j.maxDidsPerSubscriber // ceiling division
134
-
// first case: no subscribers yet; create all needed subscribers
135
-
if len(j.subscribers) == 0 {
136
-
for i := range subscribersNeeded {
137
-
startIdx := i * j.maxDidsPerSubscriber
138
-
endIdx := min((i+1)*j.maxDidsPerSubscriber, totalDids)
140
-
subscriberDids := j.cfg.WantedDids[startIdx:endIdx]
143
-
subCfg.WantedDids = subscriberDids
145
-
ident := fmt.Sprintf("%s-%d", j.baseIdent, i)
146
-
subscriber := &JetstreamSubscriber{
147
-
dids: subscriberDids,
150
-
j.subscribers = append(j.subscribers, subscriber)
85
+
sched := sequential.NewScheduler(j.ident, logger, processFunc)
152
-
j.subscriberWg.Add(1)
153
-
go j.startSubscriber(subscriber, &subCfg)
87
+
client, err := client.NewClient(j.cfg, log.New("jetstream"), sched)
89
+
return fmt.Errorf("failed to create jetstream client: %w", err)
158
-
// second case: we have more subscribers than needed, stop extra subscribers
159
-
if len(j.subscribers) > subscribersNeeded {
160
-
for i := subscribersNeeded; i < len(j.subscribers); i++ {
161
-
sub := j.subscribers[i]
162
-
if sub.running && sub.cancel != nil {
164
-
sub.running = false
95
+
for len(j.cfg.WantedDids) == 0 {
96
+
time.Sleep(time.Second)
167
-
j.subscribers = j.subscribers[:subscribersNeeded]
99
+
logger.Info("done waiting for did")
100
+
j.connectAndRead(ctx)
170
-
// third case: we need more subscribers
171
-
if len(j.subscribers) < subscribersNeeded {
172
-
existingCount := len(j.subscribers)
173
-
// Create additional subscribers
174
-
for i := existingCount; i < subscribersNeeded; i++ {
175
-
startIdx := i * j.maxDidsPerSubscriber
176
-
endIdx := min((i+1)*j.maxDidsPerSubscriber, totalDids)
178
-
subscriberDids := j.cfg.WantedDids[startIdx:endIdx]
181
-
subCfg.WantedDids = subscriberDids
183
-
ident := fmt.Sprintf("%s-%d", j.baseIdent, i)
184
-
subscriber := &JetstreamSubscriber{
185
-
dids: subscriberDids,
188
-
j.subscribers = append(j.subscribers, subscriber)
190
-
j.subscriberWg.Add(1)
191
-
go j.startSubscriber(subscriber, &subCfg)
195
-
// fourth case: update existing subscribers with new wantedDids
196
-
for i := 0; i < subscribersNeeded && i < len(j.subscribers); i++ {
197
-
startIdx := i * j.maxDidsPerSubscriber
198
-
endIdx := min((i+1)*j.maxDidsPerSubscriber, totalDids)
199
-
newDids := j.cfg.WantedDids[startIdx:endIdx]
201
-
// if the dids for this subscriber have changed, restart it
202
-
sub := j.subscribers[i]
203
-
if !didSlicesEqual(sub.dids, newDids) {
204
-
j.l.Info("subscriber DIDs changed, updating",
205
-
"subscriber", sub.ident,
206
-
"old_count", len(sub.dids),
207
-
"new_count", len(newDids))
209
-
if sub.running && sub.cancel != nil {
211
-
sub.running = false
215
-
subCfg.WantedDids = newDids
219
-
j.subscriberWg.Add(1)
220
-
go j.startSubscriber(sub, &subCfg)
225
-
func didSlicesEqual(a, b []string) bool {
226
-
if len(a) != len(b) {
230
-
aMap := make(map[string]struct{}, len(a))
231
-
for _, did := range a {
232
-
aMap[did] = struct{}{}
235
-
for _, did := range b {
236
-
if _, exists := aMap[did]; !exists {
244
-
// startSubscriber initializes and starts a single subscriber
245
-
func (j *JetstreamClient) startSubscriber(sub *JetstreamSubscriber, cfg *client.ClientConfig) {
246
-
defer j.subscriberWg.Done()
248
-
logger := j.l.With("subscriber", sub.ident)
249
-
logger.Info("starting subscriber", "dids_count", len(sub.dids))
251
-
sched := sequential.NewScheduler(sub.ident, logger, j.processFunc)
253
-
client, err := client.NewClient(cfg, log.New("jetstream-"+sub.ident), sched)
255
-
logger.Error("failed to create jetstream client", "error", err)
259
-
sub.client = client
265
-
j.connectAndReadForSubscriber(sub)
268
-
func (j *JetstreamClient) connectAndReadForSubscriber(sub *JetstreamSubscriber) {
269
-
ctx := context.Background()
270
-
l := j.l.With("subscriber", sub.ident)
106
+
func (j *JetstreamClient) connectAndRead(ctx context.Context) {
107
+
l := log.FromContext(ctx)
273
-
// Check if this subscriber should still be running
275
-
running := sub.running
279
-
l.Info("subscriber marked for shutdown")
cursor := j.getLastTimeUs(ctx)
connCtx, cancel := context.WithCancel(ctx)
114
+
j.cancelMu.Unlock()
288
-
sub.cancel = cancel
291
-
l.Info("connecting subscriber to jetstream")
292
-
if err := sub.client.ConnectAndRead(connCtx, cursor); err != nil {
116
+
if err := j.client.ConnectAndRead(connCtx, cursor); err != nil {
l.Error("error reading jetstream", "error", err)
295
-
time.Sleep(time.Second) // Small backoff before retry
301
-
l.Info("context done, stopping subscriber")
124
+
l.Info("context done, stopping jetstream")
l.Info("connection context done, reconnecting")
···
310
-
// GetRunningSubscribersCount returns the total number of currently running subscribers
311
-
func (j *JetstreamClient) GetRunningSubscribersCount() int {
313
-
defer j.mu.RUnlock()
316
-
for _, sub := range j.subscribers {
322
-
return runningCount
325
-
// Shutdown gracefully stops all subscribers
326
-
func (j *JetstreamClient) Shutdown() {
329
-
// Cancel all subscribers
330
-
for _, sub := range j.subscribers {
331
-
if sub.running && sub.cancel != nil {
333
-
sub.running = false
339
-
// Wait for all subscribers to complete
340
-
j.subscriberWg.Wait()
341
-
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()
···
356
-
// If last time is older than 2 days, start from now
145
+
// If last time is older than a week, 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")
···
366
-
l.Info("found last time_us", "time_us", lastTimeUs, "running_subscribers", j.GetRunningSubscribersCount())
155
+
l.Info("found last time_us", "time_us", lastTimeUs)
veryroundbird.house
anirudh.fi