Mirror of @tangled.org/core. Running on a Raspberry Pi Zero 2 (Please be gentle).

knotserver: remove wantedDids filter for now

anirudh.fi 93c95212 cfc4b6e6

verified
+51 -262
+1 -1
cmd/knotserver/main.go
··· 49 49 jc, err := jetstream.NewJetstreamClient(c.Server.JetstreamEndpoint, "knotserver", []string{ 50 50 tangled.PublicKeyNSID, 51 51 tangled.KnotMemberNSID, 52 - }, nil, l, db, true) 52 + }, nil, l, db, false) 53 53 if err != nil { 54 54 l.Error("failed to setup jetstream", "error", err) 55 55 }
+47 -258
jetstream/jetstream.go
··· 19 19 UpdateLastTimeUs(int64) error 20 20 } 21 21 22 - type JetstreamSubscriber struct { 23 - client *client.Client 24 - cancel context.CancelFunc 25 - dids []string 26 - ident string 27 - running bool 28 - } 29 - 30 22 type JetstreamClient struct { 31 - cfg *client.ClientConfig 32 - baseIdent string 33 - l *slog.Logger 34 - db DB 35 - waitForDid bool 36 - maxDidsPerSubscriber int 23 + cfg *client.ClientConfig 24 + client *client.Client 25 + ident string 26 + l *slog.Logger 37 27 38 - mu sync.RWMutex 39 - subscribers []*JetstreamSubscriber 40 - processFunc func(context.Context, *models.Event) error 41 - subscriberWg sync.WaitGroup 28 + db DB 29 + waitForDid bool 30 + mu sync.RWMutex 31 + 32 + cancel context.CancelFunc 33 + cancelMu sync.Mutex 42 34 } 43 35 44 36 func (j *JetstreamClient) AddDid(did string) { ··· 38 46 return 39 47 } 40 48 j.mu.Lock() 41 - defer j.mu.Unlock() 42 - 43 - // Just add to the config for now, actual subscriber management happens in UpdateDids 44 49 j.cfg.WantedDids = append(j.cfg.WantedDids, did) 50 + j.mu.Unlock() 45 51 } 46 52 47 53 func (j *JetstreamClient) UpdateDids(dids []string) { ··· 49 59 j.cfg.WantedDids = append(j.cfg.WantedDids, did) 50 60 } 51 61 } 52 - 53 - needRebalance := j.processFunc != nil 54 62 j.mu.Unlock() 55 63 56 - if needRebalance { 57 - j.rebalanceSubscribers() 64 + j.cancelMu.Lock() 65 + if j.cancel != nil { 66 + j.cancel() 58 67 } 68 + j.cancelMu.Unlock() 59 69 } 60 70 61 71 func NewJetstreamClient(endpoint, ident string, collections []string, cfg *client.ClientConfig, logger *slog.Logger, db DB, waitForDid bool) (*JetstreamClient, error) { ··· 66 76 } 67 77 68 78 return &JetstreamClient{ 69 - cfg: cfg, 70 - baseIdent: ident, 71 - db: db, 72 - l: logger, 73 - waitForDid: waitForDid, 74 - subscribers: make([]*JetstreamSubscriber, 0), 75 - maxDidsPerSubscriber: 100, 79 + cfg: cfg, 80 + ident: ident, 81 + db: db, 82 + l: logger, 83 + 84 + // This will make the goroutine in StartJetstream wait until 85 + // cfg.WantedDids has been populated, typically using UpdateDids. 86 + waitForDid: waitForDid, 76 87 }, nil 77 88 } 78 89 79 90 // StartJetstream starts the jetstream client and processes events using the provided processFunc. 80 91 // The caller is responsible for saving the last time_us to the database (just use your db.SaveLastTimeUs). 81 92 func (j *JetstreamClient) StartJetstream(ctx context.Context, processFunc func(context.Context, *models.Event) error) error { 82 - j.mu.Lock() 83 - j.processFunc = processFunc 84 - j.mu.Unlock() 93 + logger := j.l 85 94 86 - if j.waitForDid { 87 - // Start a goroutine to wait for DIDs and then start subscribers 88 - go func() { 89 - for { 90 - j.mu.RLock() 91 - hasDids := len(j.cfg.WantedDids) > 0 92 - j.mu.RUnlock() 95 + sched := sequential.NewScheduler(j.ident, logger, processFunc) 93 96 94 - if hasDids { 95 - j.l.Info("done waiting for did, starting subscribers") 96 - j.rebalanceSubscribers() 97 - return 98 - } 97 + client, err := client.NewClient(j.cfg, log.New("jetstream"), sched) 98 + if err != nil { 99 + return fmt.Errorf("failed to create jetstream client: %w", err) 100 + } 101 + j.client = client 102 + 103 + go func() { 104 + if j.waitForDid { 105 + for len(j.cfg.WantedDids) == 0 { 99 106 time.Sleep(time.Second) 100 107 } 101 - }() 102 - } else { 103 - // Start subscribers immediately 104 - j.rebalanceSubscribers() 105 - } 108 + } 109 + logger.Info("done waiting for did") 110 + j.connectAndRead(ctx) 111 + }() 106 112 107 113 return nil 108 114 } 109 115 110 - // rebalanceSubscribers creates, updates, or removes subscribers based on the current list of DIDs 111 - func (j *JetstreamClient) rebalanceSubscribers() { 112 - j.mu.Lock() 113 - defer j.mu.Unlock() 114 - 115 - if j.processFunc == nil { 116 - j.l.Warn("cannot rebalance subscribers without a process function") 117 - return 118 - } 119 - 120 - // calculate how many subscribers we need 121 - totalDids := len(j.cfg.WantedDids) 122 - subscribersNeeded := (totalDids + j.maxDidsPerSubscriber - 1) / j.maxDidsPerSubscriber // ceiling division 123 - 124 - // first case: no subscribers yet; create all needed subscribers 125 - if len(j.subscribers) == 0 { 126 - for i := range subscribersNeeded { 127 - startIdx := i * j.maxDidsPerSubscriber 128 - endIdx := min((i+1)*j.maxDidsPerSubscriber, totalDids) 129 - 130 - subscriberDids := j.cfg.WantedDids[startIdx:endIdx] 131 - 132 - subCfg := *j.cfg 133 - subCfg.WantedDids = subscriberDids 134 - 135 - ident := fmt.Sprintf("%s-%d", j.baseIdent, i) 136 - subscriber := &JetstreamSubscriber{ 137 - dids: subscriberDids, 138 - ident: ident, 139 - } 140 - j.subscribers = append(j.subscribers, subscriber) 141 - 142 - j.subscriberWg.Add(1) 143 - go j.startSubscriber(subscriber, &subCfg) 144 - } 145 - return 146 - } 147 - 148 - // second case: we have more subscribers than needed, stop extra subscribers 149 - if len(j.subscribers) > subscribersNeeded { 150 - for i := subscribersNeeded; i < len(j.subscribers); i++ { 151 - sub := j.subscribers[i] 152 - if sub.running && sub.cancel != nil { 153 - sub.cancel() 154 - sub.running = false 155 - } 156 - } 157 - j.subscribers = j.subscribers[:subscribersNeeded] 158 - } 159 - 160 - // third case: we need more subscribers 161 - if len(j.subscribers) < subscribersNeeded { 162 - existingCount := len(j.subscribers) 163 - // Create additional subscribers 164 - for i := existingCount; i < subscribersNeeded; i++ { 165 - startIdx := i * j.maxDidsPerSubscriber 166 - endIdx := min((i+1)*j.maxDidsPerSubscriber, totalDids) 167 - 168 - subscriberDids := j.cfg.WantedDids[startIdx:endIdx] 169 - 170 - subCfg := *j.cfg 171 - subCfg.WantedDids = subscriberDids 172 - 173 - ident := fmt.Sprintf("%s-%d", j.baseIdent, i) 174 - subscriber := &JetstreamSubscriber{ 175 - dids: subscriberDids, 176 - ident: ident, 177 - } 178 - j.subscribers = append(j.subscribers, subscriber) 179 - 180 - j.subscriberWg.Add(1) 181 - go j.startSubscriber(subscriber, &subCfg) 182 - } 183 - } 184 - 185 - // fourth case: update existing subscribers with new wantedDids 186 - for i := 0; i < subscribersNeeded && i < len(j.subscribers); i++ { 187 - startIdx := i * j.maxDidsPerSubscriber 188 - endIdx := min((i+1)*j.maxDidsPerSubscriber, totalDids) 189 - newDids := j.cfg.WantedDids[startIdx:endIdx] 190 - 191 - // if the dids for this subscriber have changed, restart it 192 - sub := j.subscribers[i] 193 - if !didSlicesEqual(sub.dids, newDids) { 194 - j.l.Info("subscriber DIDs changed, updating", 195 - "subscriber", sub.ident, 196 - "old_count", len(sub.dids), 197 - "new_count", len(newDids)) 198 - 199 - if sub.running && sub.cancel != nil { 200 - sub.cancel() 201 - sub.running = false 202 - } 203 - 204 - subCfg := *j.cfg 205 - subCfg.WantedDids = newDids 206 - 207 - sub.dids = newDids 208 - 209 - j.subscriberWg.Add(1) 210 - go j.startSubscriber(sub, &subCfg) 211 - } 212 - } 213 - } 214 - 215 - func didSlicesEqual(a, b []string) bool { 216 - if len(a) != len(b) { 217 - return false 218 - } 219 - 220 - aMap := make(map[string]struct{}, len(a)) 221 - for _, did := range a { 222 - aMap[did] = struct{}{} 223 - } 224 - 225 - for _, did := range b { 226 - if _, exists := aMap[did]; !exists { 227 - return false 228 - } 229 - } 230 - 231 - return true 232 - } 233 - 234 - // startSubscriber initializes and starts a single subscriber 235 - func (j *JetstreamClient) startSubscriber(sub *JetstreamSubscriber, cfg *client.ClientConfig) { 236 - defer j.subscriberWg.Done() 237 - 238 - logger := j.l.With("subscriber", sub.ident) 239 - logger.Info("starting subscriber", "dids_count", len(sub.dids)) 240 - 241 - sched := sequential.NewScheduler(sub.ident, logger, j.processFunc) 242 - 243 - client, err := client.NewClient(cfg, log.New("jetstream-"+sub.ident), sched) 244 - if err != nil { 245 - logger.Error("failed to create jetstream client", "error", err) 246 - return 247 - } 248 - 249 - sub.client = client 250 - 251 - j.mu.Lock() 252 - sub.running = true 253 - j.mu.Unlock() 254 - 255 - j.connectAndReadForSubscriber(sub) 256 - } 257 - 258 - func (j *JetstreamClient) connectAndReadForSubscriber(sub *JetstreamSubscriber) { 259 - ctx := context.Background() 260 - l := j.l.With("subscriber", sub.ident) 261 - 116 + func (j *JetstreamClient) connectAndRead(ctx context.Context) { 117 + l := log.FromContext(ctx) 262 118 for { 263 - // Check if this subscriber should still be running 264 - j.mu.RLock() 265 - running := sub.running 266 - j.mu.RUnlock() 267 - 268 - if !running { 269 - l.Info("subscriber marked for shutdown") 270 - return 271 - } 272 - 273 119 cursor := j.getLastTimeUs(ctx) 274 120 275 121 connCtx, cancel := context.WithCancel(ctx) 122 + j.cancelMu.Lock() 123 + j.cancel = cancel 124 + j.cancelMu.Unlock() 276 125 277 - j.mu.Lock() 278 - sub.cancel = cancel 279 - j.mu.Unlock() 280 - 281 - l.Info("connecting subscriber to jetstream") 282 - if err := sub.client.ConnectAndRead(connCtx, cursor); err != nil { 126 + if err := j.client.ConnectAndRead(connCtx, cursor); err != nil { 283 127 l.Error("error reading jetstream", "error", err) 284 128 cancel() 285 - time.Sleep(time.Second) // Small backoff before retry 286 129 continue 287 130 } 288 131 289 132 select { 290 133 case <-ctx.Done(): 291 - l.Info("context done, stopping subscriber") 134 + l.Info("context done, stopping jetstream") 292 135 return 293 136 case <-connCtx.Done(): 294 137 l.Info("connection context done, reconnecting") 295 138 continue 296 139 } 297 140 } 298 - } 299 - 300 - // GetRunningSubscribersCount returns the total number of currently running subscribers 301 - func (j *JetstreamClient) GetRunningSubscribersCount() int { 302 - j.mu.RLock() 303 - defer j.mu.RUnlock() 304 - 305 - runningCount := 0 306 - for _, sub := range j.subscribers { 307 - if sub.running { 308 - runningCount++ 309 - } 310 - } 311 - 312 - return runningCount 313 - } 314 - 315 - // Shutdown gracefully stops all subscribers 316 - func (j *JetstreamClient) Shutdown() { 317 - j.mu.Lock() 318 - 319 - // Cancel all subscribers 320 - for _, sub := range j.subscribers { 321 - if sub.running && sub.cancel != nil { 322 - sub.cancel() 323 - sub.running = false 324 - } 325 - } 326 - 327 - j.mu.Unlock() 328 - 329 - // Wait for all subscribers to complete 330 - j.subscriberWg.Wait() 331 - j.l.Info("all subscribers shut down", "total_subscribers", len(j.subscribers), "running_subscribers", j.GetRunningSubscribersCount()) 332 141 } 333 142 334 143 func (j *JetstreamClient) getLastTimeUs(ctx context.Context) *int64 { ··· 142 353 } 143 354 } 144 355 145 - // If last time is older than 2 days, start from now 356 + // If last time is older than a week, start from now 146 357 if time.Now().UnixMicro()-lastTimeUs > 2*24*60*60*1000*1000 { 147 358 lastTimeUs = time.Now().UnixMicro() 148 359 l.Warn("last time us is older than 2 days; discarding that and starting from now") ··· 152 363 } 153 364 } 154 365 155 - l.Info("found last time_us", "time_us", lastTimeUs, "running_subscribers", j.GetRunningSubscribersCount()) 366 + l.Info("found last time_us", "time_us", lastTimeUs) 156 367 return &lastTimeUs 157 368 }
+1 -1
knotserver/handler.go
··· 62 62 if len(dids) > 0 { 63 63 h.knotInitialized = true 64 64 close(h.init) 65 - h.jc.UpdateDids(dids) 65 + // h.jc.UpdateDids(dids) 66 66 } 67 67 68 68 r.Get("/", h.Index)
+1 -1
knotserver/jetstream.go
··· 118 118 if err := h.db.UpdateLastTimeUs(lastTimeUs); err != nil { 119 119 err = fmt.Errorf("(deferred) failed to save last time us: %w", err) 120 120 } 121 - h.jc.UpdateDids([]string{did}) 121 + // h.jc.UpdateDids([]string{did}) 122 122 }() 123 123 124 124 raw := json.RawMessage(event.Commit.Record)
+1 -1
knotserver/routes.go
··· 769 769 return 770 770 } 771 771 772 - h.jc.UpdateDids([]string{data.Did}) 772 + // h.jc.UpdateDids([]string{data.Did}) 773 773 if err := h.e.AddOwner(ThisServer, data.Did); err != nil { 774 774 l.Error("adding owner", "error", err.Error()) 775 775 writeError(w, err.Error(), http.StatusInternalServerError)