this repo has no description
edavis.dev
1package main
2
3import (
4 "context"
5 "database/sql"
6 _ "embed"
7 "log"
8 "os/signal"
9 "sync"
10 "syscall"
11 "time"
12
13 jetstream "github.com/bluesky-social/jetstream/pkg/models"
14 "github.com/gorilla/websocket"
15 _ "github.com/mattn/go-sqlite3"
16)
17
18type CheckpointResults struct {
19 Blocked int
20 Pages int
21 Transferred int
22}
23
24type Queue struct {
25 lk sync.Mutex
26 events []jetstream.Event
27}
28
29func NewQueue(capacity int) *Queue {
30 return &Queue{
31 events: make([]jetstream.Event, 0, capacity),
32 }
33}
34
35func (q *Queue) Enqueue(event jetstream.Event) {
36 q.lk.Lock()
37 defer q.lk.Unlock()
38
39 q.events = append(q.events, event)
40}
41
42func (q *Queue) Dequeue() (jetstream.Event, bool) {
43 q.lk.Lock()
44 defer q.lk.Unlock()
45
46 if len(q.events) == 0 {
47 var e jetstream.Event
48 return e, false
49 }
50
51 event := q.events[0]
52 q.events = q.events[1:]
53 return event, true
54}
55
56func (q *Queue) Size() int {
57 q.lk.Lock()
58 defer q.lk.Unlock()
59
60 return len(q.events)
61}
62
63var AppBskyAllowlist = map[string]bool{
64 "app.bsky.actor.profile": true,
65 "app.bsky.feed.generator": true,
66 "app.bsky.feed.like": true,
67 "app.bsky.feed.post": true,
68 "app.bsky.feed.postgate": true,
69 "app.bsky.feed.repost": true,
70 "app.bsky.feed.threadgate": true,
71 "app.bsky.graph.block": true,
72 "app.bsky.graph.follow": true,
73 "app.bsky.graph.list": true,
74 "app.bsky.graph.listblock": true,
75 "app.bsky.graph.listitem": true,
76 "app.bsky.graph.starterpack": true,
77 "app.bsky.labeler.service": true,
78 "chat.bsky.actor.declaration": true,
79}
80
81const JetstreamUrl = `wss://jetstream1.us-west.bsky.network/subscribe` // TODO(ejd): attach a reconnect cursor
82
83const userTimestampUpdate = `insert into users (did, ts) values (?, ?) on conflict (did) do update set ts = ?`
84
85//go:embed schema.sql
86var ddl string
87
88func handler(ctx context.Context, queue *Queue, dbCnx *sql.DB) {
89 if _, err := dbCnx.ExecContext(ctx, ddl); err != nil {
90 log.Printf("could not create tables: %v\n", err)
91 }
92 if _, err := dbCnx.ExecContext(ctx, "PRAGMA wal_autocheckpoint = 0"); err != nil {
93 log.Printf("could not set PRAGMA wal_autocheckpoint: %v\n", err)
94 }
95
96 var (
97 dbTx *sql.Tx
98 err error
99 eventCount int
100 )
101
102 for {
103 event, ok := queue.Dequeue()
104 if !ok {
105 time.Sleep(100 * time.Millisecond)
106 continue
107 }
108
109 if dbTx == nil {
110 dbTx, err = dbCnx.BeginTx(ctx, nil)
111 if err != nil {
112 log.Printf("failed to begin transaction: %v\n", err)
113 }
114 }
115
116 if event.Kind != jetstream.EventKindCommit {
117 continue
118 }
119 if event.Commit.Operation != jetstream.CommitOperationCreate {
120 // we're missing deletes and updates but this matches how bsky-activity
121 // does it so we stay consistent
122 continue
123 }
124
125 did := event.Did
126 commit := *event.Commit
127 ts := time.Now().UTC().Unix()
128
129 if _, ok := AppBskyAllowlist[commit.Collection]; !ok {
130 continue
131 }
132
133 dbTx.ExecContext(ctx, userTimestampUpdate, did, ts, ts)
134
135 eventCount += 1
136 if eventCount%2500 == 0 {
137 if err = dbTx.Commit(); err != nil {
138 log.Printf("commit failed: %v\n", err)
139 } else {
140 log.Printf("commit successful\n")
141 }
142
143 if eventCount%25_000 == 0 {
144 var results CheckpointResults
145 err = dbCnx.QueryRowContext(ctx, "PRAGMA wal_checkpoint(RESTART)").Scan(&results.Blocked, &results.Pages, &results.Transferred)
146 switch {
147 case err != nil:
148 log.Printf("failed checkpoint: %v\n", err)
149 case results.Blocked == 1:
150 log.Printf("checkpoint: blocked\n")
151 case results.Pages == results.Transferred:
152 log.Printf("checkpoint: %d pages transferred\n", results.Transferred)
153 case results.Pages != results.Transferred:
154 log.Printf("checkpoint: %d pages, %d transferred\n", results.Pages, results.Transferred)
155 }
156 }
157
158 dbTx, err = dbCnx.BeginTx(ctx, nil)
159 if err != nil {
160 log.Printf("failed to begin transaction: %v\n", err)
161 }
162
163 log.Printf("queue size: %d\n", queue.Size())
164 }
165 }
166}
167
168func main() {
169 ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
170 defer stop()
171
172 conn, _, err := websocket.DefaultDialer.DialContext(ctx, JetstreamUrl, nil)
173 if err != nil {
174 log.Fatalf("failed to open websocket: %v\n", err)
175 }
176 defer func() {
177 if err := conn.Close(); err != nil {
178 log.Printf("failed to close websocket: %v\n", err)
179 }
180 log.Printf("websocket closed\n")
181 }()
182
183 // TODO(ejd): use more readable URL params for this
184 dbCnx, err := sql.Open("sqlite3", "data/bsky-users.db?_journal=WAL&_fk=on&_timeout=5000&_sync=1&_txlock=immediate")
185 if err != nil {
186 log.Fatalf("failed to open database: %v\n", err)
187 }
188 defer func() {
189 if _, err := dbCnx.Exec("PRAGMA wal_checkpoint(TRUNCATE)"); err != nil {
190 log.Printf("error doing final WAL checkpoint: %v\n", err)
191 }
192 if err := dbCnx.Close(); err != nil {
193 log.Printf("failed to close db: %v\n", err)
194 }
195 log.Printf("db closed\n")
196 }()
197
198 queue := NewQueue(100_000)
199 go handler(ctx, queue, dbCnx)
200
201 log.Printf("starting up\n")
202 go func() {
203 for {
204 var event jetstream.Event
205 err := conn.ReadJSON(&event)
206 if err != nil {
207 log.Printf("ReadJSON error: %v\n", err)
208 stop()
209 break
210 }
211 queue.Enqueue(event)
212 }
213 }()
214
215 <-ctx.Done()
216 log.Printf("shutting down\n")
217}