its for when you want to get like notifications for your reposts
feat: rest of the rest of the owl
- use like records Via field (dusk is stupid)
- keep records so we can actually detect them being deleted (like records that were there before the server started arent counted)
- overall make the code betterer
+176
-138
+2
-5
jetstream.go
+2
-5
jetstream.go
···
3
3
import (
4
4
"context"
5
5
"log/slog"
6
-
"time"
7
6
8
7
"github.com/bluesky-social/jetstream/pkg/client"
9
8
"github.com/bluesky-social/jetstream/pkg/client/schedulers/sequential"
···
33
32
config.Compress = true
34
33
config.WantedCollections = opts.WantedCollections
35
34
config.WantedDids = opts.WantedDIDs
36
-
config.RequireHello = len(config.WantedDids) == 0
35
+
config.RequireHello = false
37
36
38
37
scheduler := sequential.NewScheduler(name, logger, handleEvent)
39
38
···
44
43
}
45
44
46
45
startFn := func() error {
47
-
cursor := time.Now().UnixMicro()
48
-
49
46
logger.Info("starting jetstream client", "name", name, "collections", opts.WantedCollections, "wanted_dids", len(opts.WantedDIDs))
50
-
if err := c.ConnectAndRead(ctx, &cursor); err != nil {
47
+
if err := c.ConnectAndRead(ctx, nil); err != nil {
51
48
logger.Error("jetstream client failed", "name", name, "error", err)
52
49
return err
53
50
}
+129
-117
main.go
+129
-117
main.go
···
8
8
"net/http"
9
9
10
10
"github.com/bluesky-social/indigo/api/bsky"
11
+
"github.com/bluesky-social/indigo/atproto/syntax"
11
12
"github.com/bluesky-social/indigo/xrpc"
12
13
"github.com/bluesky-social/jetstream/pkg/client"
13
14
"github.com/bluesky-social/jetstream/pkg/models"
···
22
23
const ListenTypeFollows = "follows"
23
24
24
25
type SubscriberData struct {
25
-
DID string
26
+
DID syntax.DID
26
27
Conn *websocket.Conn
27
28
ListenType string
28
-
ListenTo Set[string]
29
-
Reposts Set[string]
29
+
ListenTo Set[syntax.DID]
30
+
follows map[syntax.RecordKey]bsky.GraphFollow
31
+
}
32
+
33
+
type ListeneeData struct {
34
+
targets *hashmap.Map[syntax.DID, *SubscriberData]
35
+
likes map[syntax.RecordKey]bsky.FeedLike
30
36
}
31
37
32
38
type NotificationMessage struct {
33
-
Liked bool `json:"liked"`
34
-
ByDid string `json:"did"`
35
-
RepostURI string `json:"repost_uri"`
39
+
Liked bool `json:"liked"`
40
+
ByDid syntax.DID `json:"did"`
41
+
RepostURI syntax.ATURI `json:"repost_uri"`
36
42
}
37
43
38
44
type SubscriberMessage struct {
···
41
47
}
42
48
43
49
type SubscriberUpdateListenTo struct {
44
-
ListenTo []string `json:"listen_to"`
50
+
ListenTo []syntax.DID `json:"listen_to"`
45
51
}
46
52
47
53
var (
48
54
// storing the subscriber data in both Should Be Fine
49
55
// we dont modify subscriber data at the same time in two places
50
-
subscribers = hashmap.New[string, *SubscriberData]()
51
-
listeningTo = hashmap.New[string, *hashmap.Map[string, *SubscriberData]]()
56
+
subscribers = hashmap.New[syntax.DID, *SubscriberData]()
57
+
listeningTo = hashmap.New[syntax.DID, *ListeneeData]()
52
58
53
59
likeStream *client.Client
54
60
subscriberStream *client.Client
···
64
70
65
71
func getSubscriberDids() []string {
66
72
dids := make([]string, 0, subscribers.Len())
67
-
subscribers.Range(func(s string, sd *SubscriberData) bool {
68
-
dids = append(dids, s)
73
+
subscribers.Range(func(s syntax.DID, sd *SubscriberData) bool {
74
+
dids = append(dids, string(s))
69
75
return true
70
76
})
71
77
return dids
72
78
}
73
79
74
-
func listenTo(sd *SubscriberData, did string) {
75
-
targetDids, _ := listeningTo.GetOrInsert(did, hashmap.New[string, *SubscriberData]())
76
-
targetDids.Insert(sd.DID, sd)
80
+
func startListeningTo(sd *SubscriberData, did syntax.DID) {
81
+
ld, _ := listeningTo.GetOrInsert(did, &ListeneeData{
82
+
targets: hashmap.New[syntax.DID, *SubscriberData](),
83
+
likes: make(map[syntax.RecordKey]bsky.FeedLike),
84
+
})
85
+
ld.targets.Insert(sd.DID, sd)
77
86
}
78
87
79
-
func stopListeningTo(subscriberDid, did string) {
80
-
if targetDids, exists := listeningTo.Get(did); exists {
81
-
targetDids.Del(subscriberDid)
88
+
func stopListeningTo(subscriberDid, did syntax.DID) {
89
+
if ld, exists := listeningTo.Get(did); exists {
90
+
ld.targets.Del(subscriberDid)
82
91
}
83
92
}
84
93
···
99
108
100
109
func handleSubscribe(w http.ResponseWriter, r *http.Request) {
101
110
vars := mux.Vars(r)
102
-
did := vars["did"]
111
+
did, err := syntax.ParseDID(vars["did"])
112
+
if err != nil {
113
+
http.Error(w, "not a valid did", http.StatusBadRequest)
114
+
return
115
+
}
103
116
104
117
query := r.URL.Query()
105
-
// "follows", everything else is considered as "none"
106
118
listenType := query.Get("listenTo")
107
119
if len(listenType) == 0 {
108
120
listenType = ListenTypeFollows
109
121
}
110
122
111
-
logger = logger.With("did", did)
123
+
logger := logger.With("did", did)
112
124
113
125
conn, err := upgrader.Upgrade(w, r, nil)
114
126
if err != nil {
···
130
142
Host: pdsURI,
131
143
}
132
144
133
-
var subbedTo Set[string]
145
+
sd := &SubscriberData{
146
+
DID: did,
147
+
Conn: conn,
148
+
ListenType: listenType,
149
+
}
134
150
135
151
switch listenType {
136
152
case ListenTypeFollows:
···
140
156
return
141
157
}
142
158
logger.Info("fetched follows")
143
-
subbedTo = follows
159
+
sd.follows = follows
160
+
sd.ListenTo = make(Set[syntax.DID])
161
+
for _, follow := range follows {
162
+
sd.ListenTo[syntax.DID(follow.Subject)] = struct{}{}
163
+
}
144
164
case ListenTypeNone:
145
-
subbedTo = make(Set[string])
165
+
sd.ListenTo = make(Set[syntax.DID])
146
166
default:
147
-
logger.Error("invalid listen type", "requestedType", listenType)
148
-
return
149
-
}
150
-
151
-
reposts, err := fetchReposts(r.Context(), xrpcClient, did)
152
-
if err != nil {
153
-
logger.Error("error fetching reposts", "error", err)
167
+
http.Error(w, "invalid listen type", http.StatusBadRequest)
154
168
return
155
169
}
156
-
logger.Info("fetched reposts")
157
-
158
-
sd := &SubscriberData{
159
-
DID: did,
160
-
Conn: conn,
161
-
ListenType: listenType,
162
-
ListenTo: subbedTo,
163
-
Reposts: reposts,
164
-
}
165
170
166
171
subscribers.Set(sd.DID, sd)
167
172
for listenDid := range sd.ListenTo {
168
-
listenTo(sd, listenDid)
173
+
startListeningTo(sd, listenDid)
169
174
}
170
-
171
175
updateSubscriberStreamOpts()
172
-
updateLikeStreamOpts()
173
176
// delete subscriber after we are done
174
177
defer func() {
175
178
for listenDid := range sd.ListenTo {
176
179
stopListeningTo(sd.DID, listenDid)
177
180
}
178
181
subscribers.Del(sd.DID)
179
-
180
182
updateSubscriberStreamOpts()
181
-
updateLikeStreamOpts()
182
183
}()
183
184
184
185
logger.Info("serving subscriber")
···
209
210
}
210
211
for _, listenDid := range innerMsg.ListenTo {
211
212
sd.ListenTo[listenDid] = struct{}{}
212
-
listenTo(sd, listenDid)
213
+
startListeningTo(sd, listenDid)
213
214
}
214
215
}
215
216
}
···
218
219
func getLikeStreamOpts() models.SubscriberOptionsUpdatePayload {
219
220
return models.SubscriberOptionsUpdatePayload{
220
221
WantedCollections: []string{"app.bsky.feed.like"},
221
-
// WantedDIDs: getFollowsDids(),
222
222
}
223
223
}
224
224
···
229
229
}
230
230
}
231
231
232
-
func updateLikeStreamOpts() {
233
-
opts := getLikeStreamOpts()
234
-
err := likeStream.SendOptionsUpdate(opts)
235
-
if err != nil {
236
-
logger.Error("couldnt update like stream opts", "error", err)
237
-
return
238
-
}
239
-
logger.Info("updated like stream opts", "requestedDids", len(opts.WantedDIDs))
240
-
}
241
-
242
232
func updateSubscriberStreamOpts() {
243
233
opts := getSubscriberStreamOpts()
244
234
err := subscriberStream.SendOptionsUpdate(opts)
···
250
240
}
251
241
252
242
func HandleLikeEvent(ctx context.Context, event *models.Event) error {
253
-
if event == nil || event.Commit == nil || len(event.Commit.Record) == 0 {
243
+
if event == nil || event.Commit == nil {
254
244
return nil
255
245
}
256
246
247
+
byDid := syntax.DID(event.Did)
257
248
// skip handling event if its not from a source we are listening to
258
-
targets, exists := listeningTo.Get(event.Did)
249
+
ld, exists := listeningTo.Get(byDid)
259
250
if !exists {
260
251
return nil
261
252
}
262
253
254
+
deleted := event.Commit.Operation == models.CommitOperationDelete
255
+
rkey := syntax.RecordKey(event.Commit.RKey)
256
+
263
257
var like bsky.FeedLike
264
-
if err := json.Unmarshal(event.Commit.Record, &like); err != nil {
265
-
logger.Error("failed to unmarshal like", "error", err)
258
+
if deleted {
259
+
if l, exists := ld.likes[rkey]; exists {
260
+
like = l
261
+
defer delete(ld.likes, rkey)
262
+
} else {
263
+
logger.Error("like record not found", "rkey", rkey)
264
+
return nil
265
+
}
266
+
} else {
267
+
if err := json.Unmarshal(event.Commit.Record, &like); err != nil {
268
+
logger.Error("failed to unmarshal like", "error", err)
269
+
return nil
270
+
}
271
+
}
272
+
273
+
// if there is no via it means its not a repost anyway
274
+
if like.Via == nil {
266
275
return nil
267
276
}
268
277
269
-
targets.Range(func(s string, sd *SubscriberData) bool {
270
-
for repostURI, _ := range sd.Reposts {
271
-
// (un)liked a post the subscriber reposted
272
-
if like.Subject.Uri == repostURI {
273
-
notification := NotificationMessage{
274
-
Liked: event.Commit.Operation != models.CommitOperationDelete,
275
-
ByDid: event.Did,
276
-
RepostURI: repostURI,
277
-
}
278
+
// store for later when it gets deleted so we can fetch the record
279
+
if !deleted {
280
+
ld.likes[rkey] = like
281
+
}
278
282
279
-
if err := sd.Conn.WriteJSON(notification); err != nil {
280
-
logger.Error("failed to send notification", "subscriber", sd.DID, "error", err)
281
-
}
282
-
}
283
+
repostURI := syntax.ATURI(like.Via.Uri)
284
+
// if not a repost we dont care
285
+
if repostURI.Collection() != "app.bsky.feed.repost" {
286
+
return nil
287
+
}
288
+
reposterDID, err := repostURI.Authority().AsDID()
289
+
if err != nil {
290
+
return err
291
+
}
292
+
if sd, exists := ld.targets.Get(reposterDID); exists {
293
+
notification := NotificationMessage{
294
+
Liked: !deleted,
295
+
ByDid: byDid,
296
+
RepostURI: repostURI,
283
297
}
284
-
return true
285
-
})
298
+
299
+
if err := sd.Conn.WriteJSON(notification); err != nil {
300
+
logger.Error("failed to send notification", "subscriber", sd.DID, "error", err)
301
+
}
302
+
}
286
303
287
304
return nil
288
305
}
···
292
309
return nil
293
310
}
294
311
312
+
byDid := syntax.DID(event.Did)
313
+
sd, exists := subscribers.Get(byDid)
314
+
if !exists {
315
+
return nil
316
+
}
317
+
318
+
deleted := event.Commit.Operation == models.CommitOperationDelete
319
+
rkey := syntax.RecordKey(event.Commit.RKey)
320
+
295
321
switch event.Commit.Collection {
296
-
case "app.bsky.feed.repost":
297
-
modifySubscribersWithEvent(
298
-
event,
299
-
func(s *SubscriberData, r bsky.FeedRepost, deleted bool) {
300
-
if deleted {
301
-
delete(s.Reposts, r.Subject.Uri)
302
-
} else {
303
-
s.Reposts[r.Subject.Uri] = struct{}{}
304
-
}
305
-
},
306
-
)
307
322
case "app.bsky.graph.follow":
308
-
modifySubscribersWithEvent(
309
-
event,
310
-
func(s *SubscriberData, r bsky.GraphFollow, deleted bool) {
311
-
// if we arent managing then we dont need to update anything
312
-
if s.ListenType != ListenTypeFollows {
313
-
return
314
-
}
315
-
if deleted {
316
-
stopListeningTo(s.DID, r.Subject)
317
-
delete(s.ListenTo, r.Subject)
318
-
} else {
319
-
s.ListenTo[r.Subject] = struct{}{}
320
-
listenTo(s, r.Subject)
321
-
}
322
-
},
323
-
)
323
+
// if we arent managing then we dont need to update anything
324
+
if sd.ListenType != ListenTypeFollows {
325
+
return nil
326
+
}
327
+
var r bsky.GraphFollow
328
+
if deleted {
329
+
if f, exists := sd.follows[rkey]; exists {
330
+
r = f
331
+
} else {
332
+
logger.Error("follow record not found", "rkey", rkey)
333
+
return nil
334
+
}
335
+
subjectDid := syntax.DID(r.Subject)
336
+
stopListeningTo(sd.DID, subjectDid)
337
+
delete(sd.ListenTo, subjectDid)
338
+
delete(sd.follows, rkey)
339
+
} else {
340
+
if err := unmarshalEvent(event, &r); err != nil {
341
+
return err
342
+
}
343
+
subjectDid := syntax.DID(r.Subject)
344
+
sd.ListenTo[subjectDid] = struct{}{}
345
+
sd.follows[rkey] = r
346
+
startListeningTo(sd, subjectDid)
347
+
}
324
348
}
325
349
326
350
return nil
327
351
}
328
352
329
-
type ModifyFunc[v any] func(*SubscriberData, v, bool)
330
-
331
-
func modifySubscribersWithEvent[v any](event *models.Event, handle ModifyFunc[v]) error {
332
-
if len(event.Commit.Record) == 0 {
353
+
func unmarshalEvent[v any](event *models.Event, val *v) error {
354
+
if err := json.Unmarshal(event.Commit.Record, val); err != nil {
355
+
logger.Error("failed to unmarshal", "error", err, "raw", event.Commit.Record)
333
356
return nil
334
357
}
335
-
336
-
var data v
337
-
if err := json.Unmarshal(event.Commit.Record, &data); err != nil {
338
-
logger.Error("failed to unmarshal repost", "error", err, "raw", event.Commit.Record)
339
-
return nil
340
-
}
341
-
342
-
if subscriber, exists := subscribers.Get(event.Did); exists {
343
-
handle(subscriber, data, event.Commit.Operation == models.CommitOperationDelete)
344
-
}
345
-
346
358
return nil
347
359
}
+45
-16
xrpc.go
+45
-16
xrpc.go
···
10
10
"github.com/bluesky-social/indigo/atproto/identity"
11
11
"github.com/bluesky-social/indigo/atproto/syntax"
12
12
"github.com/bluesky-social/indigo/xrpc"
13
+
"github.com/bluesky-social/jetstream/pkg/models"
13
14
)
14
15
15
-
func findUserPDS(ctx context.Context, did string) (string, error) {
16
-
id, err := identity.DefaultDirectory().LookupDID(ctx, syntax.DID(did))
16
+
func findUserPDS(ctx context.Context, did syntax.DID) (string, error) {
17
+
id, err := identity.DefaultDirectory().LookupDID(ctx, did)
17
18
if err != nil {
18
19
return "", err
19
20
}
···
25
26
return pdsURI, nil
26
27
}
27
28
28
-
func fetchRecords[v any](ctx context.Context, xrpcClient *xrpc.Client, collection, did string, extractFn func(v) string) (Set[string], error) {
29
-
all := make(Set[string])
29
+
func fetchRecord[v any](ctx context.Context, xrpcClient *xrpc.Client, val *v, event *models.Event) error {
30
+
out, err := atproto.RepoGetRecord(ctx, xrpcClient, "", event.Commit.Collection, event.Did, event.Commit.RKey)
31
+
if err != nil {
32
+
return err
33
+
}
34
+
raw, _ := out.Value.MarshalJSON()
35
+
if err := json.Unmarshal(raw, val); err != nil {
36
+
return err
37
+
}
38
+
return nil
39
+
}
40
+
41
+
func fetchRecords[v any](ctx context.Context, xrpcClient *xrpc.Client, cb func(syntax.ATURI, v), collection string, did syntax.DID) error {
42
+
if xrpcClient == nil {
43
+
pdsURI, err := findUserPDS(ctx, did)
44
+
if err != nil {
45
+
return err
46
+
}
47
+
xrpcClient = &xrpc.Client{
48
+
Host: pdsURI,
49
+
}
50
+
}
51
+
30
52
cursor := ""
31
53
32
54
for {
33
55
// todo: ratelimits?? idk what this does for those
34
-
out, err := atproto.RepoListRecords(ctx, xrpcClient, collection, cursor, 100, did, false)
56
+
out, err := atproto.RepoListRecords(ctx, xrpcClient, collection, cursor, 100, string(did), false)
35
57
if err != nil {
36
-
return nil, err
58
+
return err
37
59
}
38
60
39
61
for _, record := range out.Records {
40
62
raw, _ := record.Value.MarshalJSON()
41
63
var val v
42
64
if err := json.Unmarshal(raw, &val); err != nil {
43
-
return nil, err
65
+
return err
44
66
}
45
-
s := extractFn(val)
46
-
if len(s) > 0 {
47
-
all[s] = struct{}{}
48
-
}
67
+
cb(syntax.ATURI(record.Uri), val)
49
68
}
50
69
51
70
if out.Cursor == nil || *out.Cursor == "" {
52
71
break
53
72
}
54
73
cursor = *out.Cursor
74
+
75
+
break
55
76
}
56
77
57
-
return all, nil
78
+
return nil
58
79
}
59
80
60
-
func fetchReposts(ctx context.Context, xrpcClient *xrpc.Client, did string) (Set[string], error) {
61
-
return fetchRecords(ctx, xrpcClient, "app.bsky.feed.repost", did, func(v bsky.FeedRepost) string { return v.Subject.Uri })
81
+
func fetchFollows(ctx context.Context, xrpcClient *xrpc.Client, did syntax.DID) (map[syntax.RecordKey]bsky.GraphFollow, error) {
82
+
out := make(map[syntax.RecordKey]bsky.GraphFollow)
83
+
fetchRecords(ctx, xrpcClient, func(uri syntax.ATURI, f bsky.GraphFollow) { out[uri.RecordKey()] = f }, "app.bsky.graph.follow", did)
84
+
return out, nil
62
85
}
63
86
64
-
func fetchFollows(ctx context.Context, xrpcClient *xrpc.Client, did string) (Set[string], error) {
65
-
return fetchRecords(ctx, xrpcClient, "app.bsky.graph.follow", did, func(v bsky.GraphFollow) string { return v.Subject })
87
+
func fetchRepostLikes(ctx context.Context, xrpcClient *xrpc.Client, did syntax.DID) (map[syntax.RecordKey]bsky.FeedLike, error) {
88
+
out := make(map[syntax.RecordKey]bsky.FeedLike)
89
+
fetchRecords(ctx, xrpcClient, func(uri syntax.ATURI, f bsky.FeedLike) {
90
+
if f.Via != nil && syntax.ATURI(f.Via.Uri).Collection() == "app.bsky.feed.repost" {
91
+
out[uri.RecordKey()] = f
92
+
}
93
+
}, "app.bsky.feed.like", did)
94
+
return out, nil
66
95
}