From 65dda68913c3fd73f219606ebd875ca7e1365c48 Mon Sep 17 00:00:00 2001 From: Anirudh Oppiliappan Date: Fri, 21 Mar 2025 16:24:27 +0200 Subject: [PATCH] knotserver: remove wantedDids filter for now --- cmd/knotserver/main.go | 2 +- jetstream/jetstream.go | 303 ++++++---------------------------------- knotserver/handler.go | 2 +- knotserver/jetstream.go | 2 +- knotserver/routes.go | 2 +- 5 files changed, 50 insertions(+), 261 deletions(-) diff --git a/cmd/knotserver/main.go b/cmd/knotserver/main.go index 0a05a564..3a910bd5 100644 --- a/cmd/knotserver/main.go +++ b/cmd/knotserver/main.go @@ -49,7 +49,7 @@ func main() { jc, err := jetstream.NewJetstreamClient(c.Server.JetstreamEndpoint, "knotserver", []string{ tangled.PublicKeyNSID, tangled.KnotMemberNSID, - }, nil, l, db, true) + }, nil, l, db, false) if err != nil { l.Error("failed to setup jetstream", "error", err) } diff --git a/jetstream/jetstream.go b/jetstream/jetstream.go index 5cd9577b..a35bb12c 100644 --- a/jetstream/jetstream.go +++ b/jetstream/jetstream.go @@ -19,26 +19,18 @@ type DB interface { UpdateLastTimeUs(int64) error } -type JetstreamSubscriber struct { - client *client.Client - cancel context.CancelFunc - dids []string - ident string - running bool -} - type JetstreamClient struct { - cfg *client.ClientConfig - baseIdent string - l *slog.Logger - db DB - waitForDid bool - maxDidsPerSubscriber int + cfg *client.ClientConfig + client *client.Client + ident string + l *slog.Logger + + db DB + waitForDid bool + mu sync.RWMutex - mu sync.RWMutex - subscribers []*JetstreamSubscriber - processFunc func(context.Context, *models.Event) error - subscriberWg sync.WaitGroup + cancel context.CancelFunc + cancelMu sync.Mutex } func (j *JetstreamClient) AddDid(did string) { @@ -46,10 +38,8 @@ 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) { @@ -59,13 +49,13 @@ func (j *JetstreamClient) UpdateDids(dids []string) { j.cfg.WantedDids = append(j.cfg.WantedDids, did) } } - - needRebalance := j.processFunc != nil j.mu.Unlock() - if needRebalance { - j.rebalanceSubscribers() + j.cancelMu.Lock() + if j.cancel != nil { + j.cancel() } + j.cancelMu.Unlock() } func NewJetstreamClient(endpoint, ident string, collections []string, cfg *client.ClientConfig, logger *slog.Logger, db DB, waitForDid bool) (*JetstreamClient, error) { @@ -76,229 +66,62 @@ func NewJetstreamClient(endpoint, ident string, collections []string, cfg *clien } return &JetstreamClient{ - cfg: cfg, - baseIdent: ident, - db: db, - l: logger, - waitForDid: waitForDid, - subscribers: make([]*JetstreamSubscriber, 0), - maxDidsPerSubscriber: 100, + 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, }, 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 { - 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 - } - - // calculate how many subscribers we need - totalDids := len(j.cfg.WantedDids) - subscribersNeeded := (totalDids + j.maxDidsPerSubscriber - 1) / j.maxDidsPerSubscriber // ceiling division - - // first case: no subscribers yet; create all needed subscribers - if len(j.subscribers) == 0 { - for i := range subscribersNeeded { - startIdx := i * j.maxDidsPerSubscriber - endIdx := min((i+1)*j.maxDidsPerSubscriber, totalDids) + logger := j.l - subscriberDids := j.cfg.WantedDids[startIdx:endIdx] + sched := sequential.NewScheduler(j.ident, logger, processFunc) - 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) - } - return - } - - // second case: we have more subscribers than needed, stop extra subscribers - if len(j.subscribers) > subscribersNeeded { - for i := subscribersNeeded; i < len(j.subscribers); i++ { - sub := j.subscribers[i] - if sub.running && sub.cancel != nil { - sub.cancel() - sub.running = false - } - } - j.subscribers = j.subscribers[:subscribersNeeded] - } - - // third case: we need more subscribers - if len(j.subscribers) < subscribersNeeded { - existingCount := len(j.subscribers) - // Create additional subscribers - for i := existingCount; i < subscribersNeeded; i++ { - 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) - } + client, err := client.NewClient(j.cfg, log.New("jetstream"), sched) + if err != nil { + return fmt.Errorf("failed to create jetstream client: %w", err) } + j.client = client - // fourth case: update existing subscribers with new wantedDids - for i := 0; i < subscribersNeeded && i < len(j.subscribers); i++ { - startIdx := i * j.maxDidsPerSubscriber - endIdx := min((i+1)*j.maxDidsPerSubscriber, totalDids) - newDids := j.cfg.WantedDids[startIdx:endIdx] - - // if the dids for this subscriber have changed, restart it - sub := j.subscribers[i] - if !didSlicesEqual(sub.dids, newDids) { - j.l.Info("subscriber DIDs changed, updating", - "subscriber", sub.ident, - "old_count", len(sub.dids), - "new_count", len(newDids)) - - if sub.running && sub.cancel != nil { - sub.cancel() - sub.running = false + go func() { + if j.waitForDid { + for len(j.cfg.WantedDids) == 0 { + time.Sleep(time.Second) } - - subCfg := *j.cfg - subCfg.WantedDids = newDids - - sub.dids = newDids - - j.subscriberWg.Add(1) - go j.startSubscriber(sub, &subCfg) - } - } -} - -func didSlicesEqual(a, b []string) bool { - if len(a) != len(b) { - return false - } - - aMap := make(map[string]struct{}, len(a)) - for _, did := range a { - aMap[did] = struct{}{} - } - - for _, did := range b { - if _, exists := aMap[did]; !exists { - return false } - } - - return true -} - -// 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) - if err != nil { - logger.Error("failed to create jetstream client", "error", err) - return - } - - sub.client = client - - j.mu.Lock() - sub.running = true - j.mu.Unlock() + logger.Info("done waiting for did") + j.connectAndRead(ctx) + }() - j.connectAndReadForSubscriber(sub) + return nil } -func (j *JetstreamClient) connectAndReadForSubscriber(sub *JetstreamSubscriber) { - ctx := context.Background() - l := j.l.With("subscriber", sub.ident) - +func (j *JetstreamClient) connectAndRead(ctx context.Context) { + l := log.FromContext(ctx) 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() - - l.Info("connecting subscriber to jetstream") - if err := sub.client.ConnectAndRead(connCtx, cursor); err != nil { + if err := j.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 subscriber") + l.Info("context done, stopping jetstream") return case <-connCtx.Done(): l.Info("connection context done, reconnecting") @@ -307,40 +130,6 @@ func (j *JetstreamClient) connectAndReadForSubscriber(sub *JetstreamSubscriber) } } -// 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() @@ -353,7 +142,7 @@ func (j *JetstreamClient) getLastTimeUs(ctx context.Context) *int64 { } } - // If last time is older than 2 days, start from now + // 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") @@ -363,6 +152,6 @@ func (j *JetstreamClient) getLastTimeUs(ctx context.Context) *int64 { } } - l.Info("found last time_us", "time_us", lastTimeUs, "running_subscribers", j.GetRunningSubscribersCount()) + l.Info("found last time_us", "time_us", lastTimeUs) return &lastTimeUs } diff --git a/knotserver/handler.go b/knotserver/handler.go index a828cbb8..e26d8044 100644 --- a/knotserver/handler.go +++ b/knotserver/handler.go @@ -62,7 +62,7 @@ func Setup(ctx context.Context, c *config.Config, db *db.DB, e *rbac.Enforcer, j if len(dids) > 0 { h.knotInitialized = true close(h.init) - h.jc.UpdateDids(dids) + // h.jc.UpdateDids(dids) } r.Get("/", h.Index) diff --git a/knotserver/jetstream.go b/knotserver/jetstream.go index e7d97d3c..735842b2 100644 --- a/knotserver/jetstream.go +++ b/knotserver/jetstream.go @@ -118,7 +118,7 @@ func (h *Handle) processMessages(ctx context.Context, event *models.Event) error if err := h.db.UpdateLastTimeUs(lastTimeUs); err != nil { err = fmt.Errorf("(deferred) failed to save last time us: %w", err) } - h.jc.UpdateDids([]string{did}) + // h.jc.UpdateDids([]string{did}) }() raw := json.RawMessage(event.Commit.Record) diff --git a/knotserver/routes.go b/knotserver/routes.go index 625c1aa1..1a328712 100644 --- a/knotserver/routes.go +++ b/knotserver/routes.go @@ -769,7 +769,7 @@ func (h *Handle) Init(w http.ResponseWriter, r *http.Request) { return } - h.jc.UpdateDids([]string{data.Did}) + // h.jc.UpdateDids([]string{data.Did}) if err := h.e.AddOwner(ThisServer, data.Did); err != nil { l.Error("adding owner", "error", err.Error()) writeError(w, err.Error(), http.StatusInternalServerError) -- 2.54.0