···
"github.com/bluesky-social/indigo/api/bsky"
"github.com/bluesky-social/indigo/atproto/syntax"
···
"github.com/bluesky-social/jetstream/pkg/client"
"github.com/bluesky-social/jetstream/pkg/models"
"github.com/cornelk/hashmap"
"github.com/gorilla/websocket"
···
const ListenTypeFollows = "follows"
type SubscriberData struct {
-
ListenTo Set[syntax.DID]
-
follows map[syntax.RecordKey]bsky.GraphFollow
-
type ListeneeData struct {
-
targets *hashmap.Map[syntax.DID, *SubscriberData]
-
likes map[syntax.RecordKey]bsky.FeedLike
type NotificationMessage struct {
···
// storing the subscriber data in both Should Be Fine
// we dont modify subscriber data at the same time in two places
-
subscribers = hashmap.New[syntax.DID, *SubscriberData]()
-
listeningTo = hashmap.New[syntax.DID, *ListeneeData]()
-
likeStream *client.Client
-
subscriberStream *client.Client
upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
···
func getSubscriberDids() []string {
dids := make([]string, 0, subscribers.Len())
-
subscribers.Range(func(s syntax.DID, sd *SubscriberData) bool {
-
dids = append(dids, string(s))
-
func startListeningTo(sd *SubscriberData, did syntax.DID) {
-
ld, _ := listeningTo.GetOrInsert(did, &ListeneeData{
-
targets: hashmap.New[syntax.DID, *SubscriberData](),
likes: make(map[syntax.RecordKey]bsky.FeedLike),
-
ld.targets.Insert(sd.DID, sd)
-
func stopListeningTo(subscriberDid, did syntax.DID) {
-
if ld, exists := listeningTo.Get(did); exists {
-
ld.targets.Del(subscriberDid)
···
go startJetstreamLoop(logger, &likeStream, "like_tracker", HandleLikeEvent, getLikeStreamOpts)
-
go startJetstreamLoop(logger, &subscriberStream, "subscriber", HandleSubscriberEvent, getSubscriberStreamOpts)
r.HandleFunc("/subscribe/{did}", handleSubscribe).Methods("GET")
···
http.Error(w, "not a valid did", http.StatusBadRequest)
listenType := query.Get("listenTo")
···
listenType = ListenTypeFollows
-
logger := logger.With("did", did)
conn, err := upgrader.Upgrade(w, r, nil)
···
-
ListenType: listenType,
-
follows, err := fetchFollows(r.Context(), xrpcClient, did)
logger.Error("error fetching follows", "error", err)
-
logger.Info("fetched follows")
sd.ListenTo = make(Set[syntax.DID])
-
for _, follow := range follows {
-
sd.ListenTo[syntax.DID(follow.Subject)] = struct{}{}
sd.ListenTo = make(Set[syntax.DID])
···
-
subscribers.Set(sd.DID, sd)
for listenDid := range sd.ListenTo {
-
startListeningTo(sd, listenDid)
-
updateSubscriberStreamOpts()
// delete subscriber after we are done
for listenDid := range sd.ListenTo {
-
stopListeningTo(sd.DID, listenDid)
-
subscribers.Del(sd.DID)
-
updateSubscriberStreamOpts()
logger.Info("serving subscriber")
···
// remove all current listens and add the ones the user requested
for listenDid := range sd.ListenTo {
-
stopListeningTo(sd.DID, listenDid)
delete(sd.ListenTo, listenDid)
for _, listenDid := range innerMsg.ListenTo {
sd.ListenTo[listenDid] = struct{}{}
-
startListeningTo(sd, listenDid)
···
-
func getSubscriberStreamOpts() models.SubscriberOptionsUpdatePayload {
return models.SubscriberOptionsUpdatePayload{
-
WantedCollections: []string{"app.bsky.feed.repost", "app.bsky.graph.follow"},
WantedDIDs: getSubscriberDids(),
-
func updateSubscriberStreamOpts() {
-
opts := getSubscriberStreamOpts()
-
err := subscriberStream.SendOptionsUpdate(opts)
-
logger.Error("couldnt update subscriber stream opts", "error", err)
-
logger.Info("updated subscriber stream opts", "userCount", len(opts.WantedDIDs))
func HandleLikeEvent(ctx context.Context, event *models.Event) error {
···
byDid := syntax.DID(event.Did)
// skip handling event if its not from a source we are listening to
-
ld, exists := listeningTo.Get(byDid)
···
-
if l, exists := ld.likes[rkey]; exists {
-
defer delete(ld.likes, rkey)
logger.Error("like record not found", "rkey", rkey)
···
// store for later when it gets deleted so we can fetch the record
repostURI := syntax.ATURI(like.Via.Uri)
···
-
if sd, exists := ld.targets.Get(reposterDID); exists {
notification := NotificationMessage{
···
if err := sd.Conn.WriteJSON(notification); err != nil {
-
logger.Error("failed to send notification", "subscriber", sd.DID, "error", err)
-
func HandleSubscriberEvent(ctx context.Context, event *models.Event) error {
if event == nil || event.Commit == nil {
byDid := syntax.DID(event.Did)
-
sd, exists := subscribers.Get(byDid)
···
switch event.Commit.Collection {
case "app.bsky.graph.follow":
-
// if we arent managing then we dont need to update anything
-
if sd.ListenType != ListenTypeFollows {
-
if f, exists := sd.follows[rkey]; exists {
logger.Error("follow record not found", "rkey", rkey)
-
subjectDid := syntax.DID(r.Subject)
-
stopListeningTo(sd.DID, subjectDid)
-
delete(sd.ListenTo, subjectDid)
-
delete(sd.follows, rkey)
if err := unmarshalEvent(event, &r); err != nil {
subjectDid := syntax.DID(r.Subject)
-
sd.ListenTo[subjectDid] = struct{}{}
-
startListeningTo(sd, subjectDid)
···
"github.com/bluesky-social/indigo/api/bsky"
"github.com/bluesky-social/indigo/atproto/syntax"
···
"github.com/bluesky-social/jetstream/pkg/client"
"github.com/bluesky-social/jetstream/pkg/models"
"github.com/cornelk/hashmap"
+
"github.com/google/uuid"
"github.com/gorilla/websocket"
···
const ListenTypeFollows = "follows"
type SubscriberData struct {
+
SubscribedTo syntax.DID
+
ListenTo Set[syntax.DID]
+
targets *hashmap.Map[string, *SubscriberData]
+
likes map[syntax.RecordKey]bsky.FeedLike
+
follows *hashmap.Map[syntax.RecordKey, bsky.GraphFollow]
+
followsCursor atomic.Pointer[string]
type NotificationMessage struct {
···
// storing the subscriber data in both Should Be Fine
// we dont modify subscriber data at the same time in two places
+
subscribers = hashmap.New[string, *SubscriberData]()
+
userData = hashmap.New[syntax.DID, *UserData]()
+
likeStream *client.Client
+
followStream *client.Client
upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
···
func getSubscriberDids() []string {
dids := make([]string, 0, subscribers.Len())
+
subscribers.Range(func(s string, sd *SubscriberData) bool {
+
dids = append(dids, string(sd.SubscribedTo))
+
func getUserData(did syntax.DID) *UserData {
+
ud, _ := userData.GetOrInsert(did, &UserData{
+
targets: hashmap.New[string, *SubscriberData](),
likes: make(map[syntax.RecordKey]bsky.FeedLike),
+
follows: hashmap.New[syntax.RecordKey, bsky.GraphFollow](),
+
func startListeningTo(sid string, sd *SubscriberData, did syntax.DID) {
+
ud.targets.Insert(sid, sd)
+
func stopListeningTo(sid string, did syntax.DID) {
+
if ud, exists := userData.Get(did); exists {
···
go startJetstreamLoop(logger, &likeStream, "like_tracker", HandleLikeEvent, getLikeStreamOpts)
+
go startJetstreamLoop(logger, &followStream, "subscriber", HandleFollowEvent, getFollowStreamOpts)
r.HandleFunc("/subscribe/{did}", handleSubscribe).Methods("GET")
···
http.Error(w, "not a valid did", http.StatusBadRequest)
+
sid := uuid.New().String()
listenType := query.Get("listenTo")
···
listenType = ListenTypeFollows
+
logger := logger.With("did", did, "subscriberId", sid)
conn, err := upgrader.Upgrade(w, r, nil)
···
+
ListenType: listenType,
+
follows, err := fetchFollows(r.Context(), xrpcClient, ud.followsCursor.Load(), did)
logger.Error("error fetching follows", "error", err)
sd.ListenTo = make(Set[syntax.DID])
+
// store cursor for later requests so we dont have to fetch the whole thing again
+
ud.followsCursor.Store((*string)(&follows[len(follows)-1].rkey))
+
for _, f := range follows {
+
ud.follows.Insert(f.rkey, f.follow)
+
sd.ListenTo[syntax.DID(f.follow.Subject)] = struct{}{}
+
logger.Info("fetched follows")
sd.ListenTo = make(Set[syntax.DID])
···
+
subscribers.Set(sid, sd)
for listenDid := range sd.ListenTo {
+
startListeningTo(sid, sd, listenDid)
+
updateFollowStreamOpts()
// delete subscriber after we are done
for listenDid := range sd.ListenTo {
+
stopListeningTo(sid, listenDid)
+
updateFollowStreamOpts()
logger.Info("serving subscriber")
···
// remove all current listens and add the ones the user requested
for listenDid := range sd.ListenTo {
+
stopListeningTo(sid, listenDid)
delete(sd.ListenTo, listenDid)
for _, listenDid := range innerMsg.ListenTo {
sd.ListenTo[listenDid] = struct{}{}
+
startListeningTo(sid, sd, listenDid)
···
+
func getFollowStreamOpts() models.SubscriberOptionsUpdatePayload {
return models.SubscriberOptionsUpdatePayload{
+
WantedCollections: []string{"app.bsky.graph.follow"},
WantedDIDs: getSubscriberDids(),
+
func updateFollowStreamOpts() {
+
opts := getFollowStreamOpts()
+
err := followStream.SendOptionsUpdate(opts)
+
logger.Error("couldnt update follow stream opts", "error", err)
+
logger.Info("updated follow stream opts", "userCount", len(opts.WantedDIDs))
func HandleLikeEvent(ctx context.Context, event *models.Event) error {
···
byDid := syntax.DID(event.Did)
// skip handling event if its not from a source we are listening to
+
ud, exists := userData.Get(byDid)
+
if !exists || ud.targets.Len() == 0 {
···
+
if l, exists := ud.likes[rkey]; exists {
+
defer delete(ud.likes, rkey)
logger.Error("like record not found", "rkey", rkey)
···
// store for later when it gets deleted so we can fetch the record
repostURI := syntax.ATURI(like.Via.Uri)
···
+
ud.targets.Range(func(sid string, sd *SubscriberData) bool {
+
if sd.SubscribedTo != reposterDID {
notification := NotificationMessage{
···
if err := sd.Conn.WriteJSON(notification); err != nil {
+
logger.Error("failed to send notification", "subscriber", sd.SubscribedTo, "error", err)
+
func HandleFollowEvent(ctx context.Context, event *models.Event) error {
if event == nil || event.Commit == nil {
byDid := syntax.DID(event.Did)
+
ud, exists := userData.Get(byDid)
+
if !exists || ud.targets.Len() == 0 {
···
switch event.Commit.Collection {
case "app.bsky.graph.follow":
+
if f, exists := ud.follows.Get(rkey); exists {
logger.Error("follow record not found", "rkey", rkey)
if err := unmarshalEvent(event, &r); err != nil {
+
logger.Error("could not unmarshal follow event", "error", err)
+
ud.follows.Insert(rkey, r)
+
ud.targets.Range(func(sid string, sd *SubscriberData) bool {
+
// if we arent managing then we dont need to update anything
+
if sd.ListenType != ListenTypeFollows {
subjectDid := syntax.DID(r.Subject)
+
stopListeningTo(sid, subjectDid)
+
delete(sd.ListenTo, subjectDid)
+
sd.ListenTo[subjectDid] = struct{}{}
+
startListeningTo(sid, sd, subjectDid)
ptr.pet