+97 -71

main.go

reviewed

··· 6 6 "log" 7 7 "log/slog" 8 8 "net/http" 9 9 + "sync/atomic" 9 10 10 11 "github.com/bluesky-social/indigo/api/bsky" 11 12 "github.com/bluesky-social/indigo/atproto/syntax" ··· 13 14 "github.com/bluesky-social/jetstream/pkg/client" 14 15 "github.com/bluesky-social/jetstream/pkg/models" 15 16 "github.com/cornelk/hashmap" 17 17 + "github.com/google/uuid" 16 18 "github.com/gorilla/mux" 17 19 "github.com/gorilla/websocket" 18 20 ) ··· 23 25 const ListenTypeFollows = "follows" 24 26 25 27 type SubscriberData struct { 26 26 - DID syntax.DID 27 27 - Conn *websocket.Conn 28 28 - ListenType string 29 29 - ListenTo Set[syntax.DID] 30 30 - follows map[syntax.RecordKey]bsky.GraphFollow 28 28 + SubscribedTo syntax.DID 29 29 + Conn *websocket.Conn 30 30 + ListenType string 31 31 + ListenTo Set[syntax.DID] 31 32 } 32 33 33 33 - type ListeneeData struct { 34 34 - targets *hashmap.Map[syntax.DID, *SubscriberData] 35 35 - likes map[syntax.RecordKey]bsky.FeedLike 34 34 + type UserData struct { 35 35 + targets *hashmap.Map[string, *SubscriberData] 36 36 + likes map[syntax.RecordKey]bsky.FeedLike 37 37 + follows *hashmap.Map[syntax.RecordKey, bsky.GraphFollow] 38 38 + followsCursor atomic.Pointer[string] 36 39 } 37 40 38 41 type NotificationMessage struct { ··· 53 56 var ( 54 57 // storing the subscriber data in both Should Be Fine 55 58 // we dont modify subscriber data at the same time in two places 56 56 - subscribers = hashmap.New[syntax.DID, *SubscriberData]() 57 57 - listeningTo = hashmap.New[syntax.DID, *ListeneeData]() 59 59 + subscribers = hashmap.New[string, *SubscriberData]() 60 60 + userData = hashmap.New[syntax.DID, *UserData]() 58 61 59 59 - likeStream *client.Client 60 60 - subscriberStream *client.Client 62 62 + likeStream *client.Client 63 63 + followStream *client.Client 61 64 62 65 upgrader = websocket.Upgrader{ 63 66 CheckOrigin: func(r *http.Request) bool { ··· 70 73 71 74 func getSubscriberDids() []string { 72 75 dids := make([]string, 0, subscribers.Len()) 73 73 - subscribers.Range(func(s syntax.DID, sd *SubscriberData) bool { 74 74 - dids = append(dids, string(s)) 76 76 + subscribers.Range(func(s string, sd *SubscriberData) bool { 77 77 + dids = append(dids, string(sd.SubscribedTo)) 75 78 return true 76 79 }) 77 80 return dids 78 81 } 79 82 80 80 - func startListeningTo(sd *SubscriberData, did syntax.DID) { 81 81 - ld, _ := listeningTo.GetOrInsert(did, &ListeneeData{ 82 82 - targets: hashmap.New[syntax.DID, *SubscriberData](), 83 83 + func getUserData(did syntax.DID) *UserData { 84 84 + ud, _ := userData.GetOrInsert(did, &UserData{ 85 85 + targets: hashmap.New[string, *SubscriberData](), 83 86 likes: make(map[syntax.RecordKey]bsky.FeedLike), 87 87 + follows: hashmap.New[syntax.RecordKey, bsky.GraphFollow](), 84 88 }) 85 85 - ld.targets.Insert(sd.DID, sd) 89 89 + return ud 86 90 } 87 91 88 88 - func stopListeningTo(subscriberDid, did syntax.DID) { 89 89 - if ld, exists := listeningTo.Get(did); exists { 90 90 - ld.targets.Del(subscriberDid) 92 92 + func startListeningTo(sid string, sd *SubscriberData, did syntax.DID) { 93 93 + ud := getUserData(did) 94 94 + ud.targets.Insert(sid, sd) 95 95 + } 96 96 + 97 97 + func stopListeningTo(sid string, did syntax.DID) { 98 98 + if ud, exists := userData.Get(did); exists { 99 99 + ud.targets.Del(sid) 91 100 } 92 101 } 93 102 ··· 95 104 logger = slog.Default() 96 105 97 106 go startJetstreamLoop(logger, &likeStream, "like_tracker", HandleLikeEvent, getLikeStreamOpts) 98 98 - go startJetstreamLoop(logger, &subscriberStream, "subscriber", HandleSubscriberEvent, getSubscriberStreamOpts) 107 107 + go startJetstreamLoop(logger, &followStream, "subscriber", HandleFollowEvent, getFollowStreamOpts) 99 108 100 109 r := mux.NewRouter() 101 110 r.HandleFunc("/subscribe/{did}", handleSubscribe).Methods("GET") ··· 113 122 http.Error(w, "not a valid did", http.StatusBadRequest) 114 123 return 115 124 } 125 125 + sid := uuid.New().String() 116 126 117 127 query := r.URL.Query() 118 128 listenType := query.Get("listenTo") ··· 120 130 listenType = ListenTypeFollows 121 131 } 122 132 123 123 - logger := logger.With("did", did) 133 133 + logger := logger.With("did", did, "subscriberId", sid) 124 134 125 135 conn, err := upgrader.Upgrade(w, r, nil) 126 136 if err != nil { ··· 142 152 Host: pdsURI, 143 153 } 144 154 155 155 + ud := getUserData(did) 145 156 sd := &SubscriberData{ 146 146 - DID: did, 147 147 - Conn: conn, 148 148 - ListenType: listenType, 157 157 + SubscribedTo: did, 158 158 + Conn: conn, 159 159 + ListenType: listenType, 149 160 } 150 161 151 162 switch listenType { 152 163 case ListenTypeFollows: 153 153 - follows, err := fetchFollows(r.Context(), xrpcClient, did) 164 164 + follows, err := fetchFollows(r.Context(), xrpcClient, ud.followsCursor.Load(), did) 154 165 if err != nil { 155 166 logger.Error("error fetching follows", "error", err) 156 167 return 157 168 } 158 158 - logger.Info("fetched follows") 159 159 - sd.follows = follows 160 169 sd.ListenTo = make(Set[syntax.DID]) 161 161 - for _, follow := range follows { 162 162 - sd.ListenTo[syntax.DID(follow.Subject)] = struct{}{} 170 170 + if len(follows) > 0 { 171 171 + // store cursor for later requests so we dont have to fetch the whole thing again 172 172 + ud.followsCursor.Store((*string)(&follows[len(follows)-1].rkey)) 173 173 + for _, f := range follows { 174 174 + ud.follows.Insert(f.rkey, f.follow) 175 175 + sd.ListenTo[syntax.DID(f.follow.Subject)] = struct{}{} 176 176 + } 163 177 } 178 178 + logger.Info("fetched follows") 164 179 case ListenTypeNone: 165 180 sd.ListenTo = make(Set[syntax.DID]) 166 181 default: ··· 168 183 return 169 184 } 170 185 171 171 - subscribers.Set(sd.DID, sd) 186 186 + subscribers.Set(sid, sd) 172 187 for listenDid := range sd.ListenTo { 173 173 - startListeningTo(sd, listenDid) 188 188 + startListeningTo(sid, sd, listenDid) 174 189 } 175 175 - updateSubscriberStreamOpts() 190 190 + updateFollowStreamOpts() 176 191 // delete subscriber after we are done 177 192 defer func() { 178 193 for listenDid := range sd.ListenTo { 179 179 - stopListeningTo(sd.DID, listenDid) 194 194 + stopListeningTo(sid, listenDid) 180 195 } 181 181 - subscribers.Del(sd.DID) 182 182 - updateSubscriberStreamOpts() 196 196 + subscribers.Del(sid) 197 197 + updateFollowStreamOpts() 183 198 }() 184 199 185 200 logger.Info("serving subscriber") ··· 205 220 } 206 221 // remove all current listens and add the ones the user requested 207 222 for listenDid := range sd.ListenTo { 208 208 - stopListeningTo(sd.DID, listenDid) 223 223 + stopListeningTo(sid, listenDid) 209 224 delete(sd.ListenTo, listenDid) 210 225 } 211 226 for _, listenDid := range innerMsg.ListenTo { 212 227 sd.ListenTo[listenDid] = struct{}{} 213 213 - startListeningTo(sd, listenDid) 228 228 + startListeningTo(sid, sd, listenDid) 214 229 } 215 230 } 216 231 } ··· 222 237 } 223 238 } 224 239 225 225 - func getSubscriberStreamOpts() models.SubscriberOptionsUpdatePayload { 240 240 + func getFollowStreamOpts() models.SubscriberOptionsUpdatePayload { 226 241 return models.SubscriberOptionsUpdatePayload{ 227 227 - WantedCollections: []string{"app.bsky.feed.repost", "app.bsky.graph.follow"}, 242 242 + WantedCollections: []string{"app.bsky.graph.follow"}, 228 243 WantedDIDs: getSubscriberDids(), 229 244 } 230 245 } 231 246 232 232 - func updateSubscriberStreamOpts() { 233 233 - opts := getSubscriberStreamOpts() 234 234 - err := subscriberStream.SendOptionsUpdate(opts) 247 247 + func updateFollowStreamOpts() { 248 248 + opts := getFollowStreamOpts() 249 249 + err := followStream.SendOptionsUpdate(opts) 235 250 if err != nil { 236 236 - logger.Error("couldnt update subscriber stream opts", "error", err) 251 251 + logger.Error("couldnt update follow stream opts", "error", err) 237 252 return 238 253 } 239 239 - logger.Info("updated subscriber stream opts", "userCount", len(opts.WantedDIDs)) 254 254 + logger.Info("updated follow stream opts", "userCount", len(opts.WantedDIDs)) 240 255 } 241 256 242 257 func HandleLikeEvent(ctx context.Context, event *models.Event) error { ··· 246 261 247 262 byDid := syntax.DID(event.Did) 248 263 // skip handling event if its not from a source we are listening to 249 249 - ld, exists := listeningTo.Get(byDid) 250 250 - if !exists { 264 264 + ud, exists := userData.Get(byDid) 265 265 + if !exists || ud.targets.Len() == 0 { 251 266 return nil 252 267 } 253 268 ··· 256 271 257 272 var like bsky.FeedLike 258 273 if deleted { 259 259 - if l, exists := ld.likes[rkey]; exists { 274 274 + if l, exists := ud.likes[rkey]; exists { 260 275 like = l 261 261 - defer delete(ld.likes, rkey) 276 276 + defer delete(ud.likes, rkey) 262 277 } else { 263 278 logger.Error("like record not found", "rkey", rkey) 264 279 return nil ··· 277 292 278 293 // store for later when it gets deleted so we can fetch the record 279 294 if !deleted { 280 280 - ld.likes[rkey] = like 295 295 + ud.likes[rkey] = like 281 296 } 282 297 283 298 repostURI := syntax.ATURI(like.Via.Uri) ··· 289 304 if err != nil { 290 305 return err 291 306 } 292 292 - if sd, exists := ld.targets.Get(reposterDID); exists { 307 307 + ud.targets.Range(func(sid string, sd *SubscriberData) bool { 308 308 + if sd.SubscribedTo != reposterDID { 309 309 + return true 310 310 + } 311 311 + 293 312 notification := NotificationMessage{ 294 313 Liked: !deleted, 295 314 ByDid: byDid, ··· 297 316 } 298 317 299 318 if err := sd.Conn.WriteJSON(notification); err != nil { 300 300 - logger.Error("failed to send notification", "subscriber", sd.DID, "error", err) 319 319 + logger.Error("failed to send notification", "subscriber", sd.SubscribedTo, "error", err) 301 320 } 302 302 - } 321 321 + return true 322 322 + }) 303 323 304 324 return nil 305 325 } 306 326 307 307 - func HandleSubscriberEvent(ctx context.Context, event *models.Event) error { 327 327 + func HandleFollowEvent(ctx context.Context, event *models.Event) error { 308 328 if event == nil || event.Commit == nil { 309 329 return nil 310 330 } 311 331 312 332 byDid := syntax.DID(event.Did) 313 313 - sd, exists := subscribers.Get(byDid) 314 314 - if !exists { 333 333 + ud, exists := userData.Get(byDid) 334 334 + if !exists || ud.targets.Len() == 0 { 315 335 return nil 316 336 } 317 337 ··· 320 340 321 341 switch event.Commit.Collection { 322 342 case "app.bsky.graph.follow": 323 323 - // if we arent managing then we dont need to update anything 324 324 - if sd.ListenType != ListenTypeFollows { 325 325 - return nil 326 326 - } 327 343 var r bsky.GraphFollow 328 344 if deleted { 329 329 - if f, exists := sd.follows[rkey]; exists { 345 345 + if f, exists := ud.follows.Get(rkey); exists { 330 346 r = f 331 347 } else { 332 348 logger.Error("follow record not found", "rkey", rkey) 333 349 return nil 334 350 } 335 335 - subjectDid := syntax.DID(r.Subject) 336 336 - stopListeningTo(sd.DID, subjectDid) 337 337 - delete(sd.ListenTo, subjectDid) 338 338 - delete(sd.follows, rkey) 351 351 + ud.follows.Del(rkey) 339 352 } else { 340 353 if err := unmarshalEvent(event, &r); err != nil { 341 341 - return err 354 354 + logger.Error("could not unmarshal follow event", "error", err) 355 355 + return nil 356 356 + } 357 357 + ud.follows.Insert(rkey, r) 358 358 + } 359 359 + ud.targets.Range(func(sid string, sd *SubscriberData) bool { 360 360 + // if we arent managing then we dont need to update anything 361 361 + if sd.ListenType != ListenTypeFollows { 362 362 + return true 342 363 } 343 364 subjectDid := syntax.DID(r.Subject) 344 344 - sd.ListenTo[subjectDid] = struct{}{} 345 345 - sd.follows[rkey] = r 346 346 - startListeningTo(sd, subjectDid) 347 347 - } 365 365 + if deleted { 366 366 + stopListeningTo(sid, subjectDid) 367 367 + delete(sd.ListenTo, subjectDid) 368 368 + } else { 369 369 + sd.ListenTo[subjectDid] = struct{}{} 370 370 + startListeningTo(sid, sd, subjectDid) 371 371 + } 372 372 + return true 373 373 + }) 348 374 } 349 375 350 376 return nil

+19 -11

xrpc.go

reviewed

··· 38 38 return nil 39 39 } 40 40 41 41 - func fetchRecords[v any](ctx context.Context, xrpcClient *xrpc.Client, cb func(syntax.ATURI, v), collection string, did syntax.DID) error { 41 41 + func fetchRecords[v any](ctx context.Context, xrpcClient *xrpc.Client, cb func(syntax.ATURI, v), cursor *string, collection string, did syntax.DID) error { 42 42 if xrpcClient == nil { 43 43 pdsURI, err := findUserPDS(ctx, did) 44 44 if err != nil { ··· 49 49 } 50 50 } 51 51 52 52 - cursor := "" 52 52 + var cur string = "" 53 53 + if cursor != nil { 54 54 + cur = *cursor 55 55 + } 53 56 54 57 for { 55 58 // todo: ratelimits?? idk what this does for those 56 56 - out, err := atproto.RepoListRecords(ctx, xrpcClient, collection, cursor, 100, string(did), false) 59 59 + out, err := atproto.RepoListRecords(ctx, xrpcClient, collection, cur, 100, string(did), true) 57 60 if err != nil { 58 61 return err 59 62 } ··· 70 73 if out.Cursor == nil || *out.Cursor == "" { 71 74 break 72 75 } 73 73 - cursor = *out.Cursor 74 74 - 75 75 - break 76 76 + cur = *out.Cursor 76 77 } 77 78 78 79 return nil 79 80 } 80 81 81 81 - func fetchFollows(ctx context.Context, xrpcClient *xrpc.Client, did syntax.DID) (map[syntax.RecordKey]bsky.GraphFollow, error) { 82 82 - out := make(map[syntax.RecordKey]bsky.GraphFollow) 83 83 - fetchRecords(ctx, xrpcClient, func(uri syntax.ATURI, f bsky.GraphFollow) { out[uri.RecordKey()] = f }, "app.bsky.graph.follow", did) 82 82 + type FetchFollowItem struct { 83 83 + rkey syntax.RecordKey 84 84 + follow bsky.GraphFollow 85 85 + } 86 86 + 87 87 + func fetchFollows(ctx context.Context, xrpcClient *xrpc.Client, cursor *string, did syntax.DID) ([]FetchFollowItem, error) { 88 88 + out := make([]FetchFollowItem, 0) 89 89 + fetchRecords(ctx, xrpcClient, func(uri syntax.ATURI, f bsky.GraphFollow) { 90 90 + out = append(out, FetchFollowItem{rkey: uri.RecordKey(), follow: f}) 91 91 + }, cursor, "app.bsky.graph.follow", did) 84 92 return out, nil 85 93 } 86 94 87 87 - func fetchRepostLikes(ctx context.Context, xrpcClient *xrpc.Client, did syntax.DID) (map[syntax.RecordKey]bsky.FeedLike, error) { 95 95 + func fetchRepostLikes(ctx context.Context, xrpcClient *xrpc.Client, cursor *string, did syntax.DID) (map[syntax.RecordKey]bsky.FeedLike, error) { 88 96 out := make(map[syntax.RecordKey]bsky.FeedLike) 89 97 fetchRecords(ctx, xrpcClient, func(uri syntax.ATURI, f bsky.FeedLike) { 90 98 if f.Via != nil && syntax.ATURI(f.Via.Uri).Collection() == "app.bsky.feed.repost" { 91 99 out[uri.RecordKey()] = f 92 100 } 93 93 - }, "app.bsky.feed.like", did) 101 101 + }, cursor, "app.bsky.feed.like", did) 94 102 return out, nil 95 103 }