This repository has no description
9.3 kB
356 lines
1package knotstream
2
3import (
4 "context"
5 "database/sql"
6 "encoding/json"
7 "fmt"
8 "log/slog"
9 "math/rand"
10 "net/http"
11 "sync"
12 "time"
13
14 "github.com/bluesky-social/indigo/atproto/syntax"
15 "github.com/bluesky-social/indigo/util/ssrf"
16 "github.com/carlmjohnson/versioninfo"
17 "github.com/gorilla/websocket"
18 "tangled.org/core/api/tangled"
19 "tangled.org/core/knotmirror/config"
20 "tangled.org/core/knotmirror/db"
21 "tangled.org/core/knotmirror/models"
22 "tangled.org/core/log"
23)
24
25type KnotSlurper struct {
26 logger *slog.Logger
27 db *sql.DB
28 cfg config.SlurperConfig
29
30 subsLk sync.Mutex
31 subs map[string]*subscription
32}
33
34func NewKnotSlurper(l *slog.Logger, db *sql.DB, cfg config.SlurperConfig) *KnotSlurper {
35 return &KnotSlurper{
36 logger: log.SubLogger(l, "slurper"),
37 db: db,
38 cfg: cfg,
39 subs: make(map[string]*subscription),
40 }
41}
42
43func (s *KnotSlurper) Run(ctx context.Context) {
44 for {
45 select {
46 case <-ctx.Done():
47 return
48 case <-time.After(s.cfg.PersistCursorPeriod):
49 if err := s.persistCursors(ctx); err != nil {
50 s.logger.Error("failed to flush cursors", "err", err)
51 }
52 }
53 }
54}
55
56func (s *KnotSlurper) CheckIfSubscribed(hostname string) bool {
57 s.subsLk.Lock()
58 defer s.subsLk.Unlock()
59
60 _, ok := s.subs[hostname]
61 return ok
62}
63
64func (s *KnotSlurper) Shutdown(ctx context.Context) error {
65 s.logger.Info("starting shutdown host cursor flush")
66 err := s.persistCursors(ctx)
67 if err != nil {
68 s.logger.Error("shutdown error", "err", err)
69 }
70 s.logger.Info("slurper shutdown complete")
71 return err
72}
73
74func (s *KnotSlurper) persistCursors(ctx context.Context) error {
75 // // gather cursor list from subscriptions and store them to DB
76 // start := time.Now()
77
78 s.subsLk.Lock()
79 cursors := make([]models.HostCursor, len(s.subs))
80 i := 0
81 for _, sub := range s.subs {
82 cursors[i] = sub.HostCursor()
83 i++
84 }
85 s.subsLk.Unlock()
86
87 err := db.StoreCursors(ctx, s.db, cursors)
88 // s.logger.Info("finished persisting cursors", "count", len(cursors), "duration", time.Since(start).String(), "err", err)
89 return err
90}
91
92func (s *KnotSlurper) Subscribe(host models.Host) error {
93 s.subsLk.Lock()
94 defer s.subsLk.Unlock()
95
96 _, ok := s.subs[host.Hostname]
97 if ok {
98 return fmt.Errorf("already subscribed: %s", host.Hostname)
99 }
100
101 // TODO: include `cancel` function to kill subscription by hostname
102 sub := &subscription{
103 hostname: host.Hostname,
104 scheduler: NewParallelScheduler(
105 s.cfg.ConcurrencyPerHost,
106 host.Hostname,
107 s.ProcessEvent,
108 ),
109 }
110 s.subs[host.Hostname] = sub
111
112 // TODO: use service level context, not the top-most one.
113 // Using top-most context should be avoided to do graceful shutdown.
114 ctx := context.TODO()
115
116 sub.scheduler.Start(ctx)
117 go s.subscribeWithRedialer(ctx, host, sub)
118 return nil
119}
120
121func (s *KnotSlurper) subscribeWithRedialer(ctx context.Context, host models.Host, sub *subscription) {
122 l := s.logger.With("host", host.Hostname)
123 defer func() {
124 s.subsLk.Lock()
125 defer s.subsLk.Unlock()
126
127 l.Info("unsubscribing knot")
128 delete(s.subs, host.Hostname)
129 }()
130
131 dialer := websocket.Dialer{
132 HandshakeTimeout: time.Second * 5,
133 }
134
135 // if this isn't a localhost / private connection, then we should enable SSRF protections
136 if !host.NoSSL {
137 netDialer := ssrf.PublicOnlyDialer()
138 dialer.NetDialContext = netDialer.DialContext
139 }
140
141 cursor := host.LastSeq
142
143 connectedInbound.Inc()
144 defer connectedInbound.Dec()
145
146 var backoff int
147 for {
148 select {
149 case <-ctx.Done():
150 return
151 default:
152 }
153 u := host.LegacyEventsURL(cursor)
154 l.Debug("made url with cursor", "cursor", cursor, "url", u)
155
156 // NOTE: manual backoff retry implementation to explicitly handle fails
157 hdr := make(http.Header)
158 hdr.Add("User-Agent", userAgent())
159 conn, resp, err := dialer.DialContext(ctx, u, hdr)
160 if err != nil {
161 l.Warn("dialing failed", "err", err, "backoff", backoff)
162 time.Sleep(sleepForBackoff(backoff))
163 backoff++
164 if backoff > 30 {
165 l.Warn("host does not appear to be online, disabling for now")
166 host.Status = models.HostStatusOffline
167 if err := db.UpsertHost(ctx, s.db, &host); err != nil {
168 l.Error("failed to update host status", "err", err)
169 }
170 return
171 }
172 continue
173 }
174
175 l.Debug("knot event subscription response", "code", resp.StatusCode, "url", u)
176
177 if err := s.handleConnection(ctx, conn, sub); err != nil {
178 // TODO: measure the last N connection error times and if they're coming too fast reconnect slower or don't reconnect and wait for requestCrawl
179 l.Warn("host connection failed", "err", err, "backoff", backoff)
180 }
181
182 updatedCursor := sub.LastSeq()
183 didProgress := updatedCursor > cursor
184 l.Debug("cursor compare", "cursor", cursor, "updatedCursor", updatedCursor, "didProgress", didProgress)
185 if cursor == 0 || didProgress {
186 cursor = updatedCursor
187 backoff = 0
188
189 batch := []models.HostCursor{sub.HostCursor()}
190 if err := db.StoreCursors(ctx, s.db, batch); err != nil {
191 l.Error("failed to store cursors", "err", err)
192 }
193 }
194 }
195}
196
197// handleConnection handles websocket connection.
198// Schedules task from received event and return when connection is closed
199func (s *KnotSlurper) handleConnection(ctx context.Context, conn *websocket.Conn, sub *subscription) error {
200 // ping on every 30s
201 ctx, cancel := context.WithCancel(ctx)
202 defer cancel() // close the background ping job on connection close
203 go func() {
204 t := time.NewTicker(30 * time.Second)
205 defer t.Stop()
206 failcount := 0
207
208 for {
209 select {
210 case <-t.C:
211 if err := conn.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(time.Second*10)); err != nil {
212 s.logger.Warn("failed to ping", "err", err)
213 failcount++
214 if failcount >= 4 {
215 s.logger.Error("too many ping fails", "count", failcount)
216 _ = conn.Close()
217 return
218 }
219 } else {
220 failcount = 0 // ok ping
221 }
222 case <-ctx.Done():
223 _ = conn.Close()
224 return
225 }
226 }
227 }()
228
229 conn.SetPingHandler(func(message string) error {
230 err := conn.WriteControl(websocket.PongMessage, []byte(message), time.Now().Add(time.Minute))
231 if err == websocket.ErrCloseSent {
232 return nil
233 }
234 return err
235 })
236 conn.SetPongHandler(func(_ string) error {
237 if err := conn.SetReadDeadline(time.Now().Add(time.Minute)); err != nil {
238 s.logger.Error("failed to set read deadline", "err", err)
239 }
240 return nil
241 })
242
243 for {
244 select {
245 case <-ctx.Done():
246 return ctx.Err()
247 default:
248 }
249 msgType, msg, err := conn.ReadMessage()
250 if err != nil {
251 return err
252 }
253
254 if msgType != websocket.TextMessage {
255 continue
256 }
257
258 sub.scheduler.AddTask(ctx, &Task{
259 key: sub.hostname, // TODO: replace to repository AT-URI for better concurrency
260 message: msg,
261 })
262 }
263}
264
265type LegacyGitEvent struct {
266 Rkey string
267 Nsid string
268 Event tangled.GitRefUpdate
269}
270
271func (s *KnotSlurper) ProcessEvent(ctx context.Context, task *Task) error {
272 var legacyMessage LegacyGitEvent
273 if err := json.Unmarshal(task.message, &legacyMessage); err != nil {
274 return fmt.Errorf("unmarshaling message: %w", err)
275 }
276
277 if err := s.ProcessLegacyGitRefUpdate(ctx, task.key, &legacyMessage); err != nil {
278 return fmt.Errorf("processing gitRefUpdate: %w", err)
279 }
280 return nil
281}
282
283func (s *KnotSlurper) ProcessLegacyGitRefUpdate(ctx context.Context, source string, evt *LegacyGitEvent) error {
284 knotstreamEventsReceived.Inc()
285
286 l := s.logger.With("src", source)
287
288 ownerDid := ""
289 if evt.Event.OwnerDid != nil {
290 ownerDid = *evt.Event.OwnerDid
291 } else {
292 // handle legacy event
293 if evt.Event.RepoDid != nil {
294 ownerDid = *evt.Event.RepoDid
295 }
296 }
297 curr, err := db.GetRepoByName(ctx, s.db, syntax.DID(ownerDid), evt.Event.RepoName)
298 if err != nil {
299 return fmt.Errorf("failed to get repo '%s': %w", ownerDid+"/"+evt.Event.RepoName, err)
300 }
301 if curr == nil {
302 // if repo doesn't exist in DB, just ignore the event. That repo is unknown.
303 //
304 // Normally did+name is already enough to perform git-fetch as that's
305 // what needed to fetch the repository.
306 // But we want to store that in did/rkey in knot-mirror.
307 // Therefore, we should ignore when the repository is unknown.
308 // Hopefully crawler will sync it later.
309 l.Warn("skipping event from unknown repo", "did/name", ownerDid+"/"+evt.Event.RepoName)
310 knotstreamEventsSkipped.Inc()
311 return nil
312 }
313 l = l.With("repoAt", curr.AtUri())
314
315 // TODO: should plan resync to resyncBuffer on RepoStateResyncing
316 if curr.State != models.RepoStateActive {
317 l.Debug("skipping non-active repo")
318 knotstreamEventsSkipped.Inc()
319 return nil
320 }
321
322 if curr.GitRev != "" && evt.Rkey <= curr.GitRev.String() {
323 l.Debug("skipping replayed event", "event.Rkey", evt.Rkey, "currentRev", curr.GitRev)
324 knotstreamEventsSkipped.Inc()
325 return nil
326 }
327
328 // if curr.State == models.RepoStateResyncing {
329 // firehoseEventsSkipped.Inc()
330 // return fp.events.addToResyncBuffer(ctx, commit)
331 // }
332
333 // can't skip anything, update repo state
334 if err := db.UpdateRepoState(ctx, s.db, curr.Did, curr.Rkey, models.RepoStateDesynchronized); err != nil {
335 return err
336 }
337
338 l.Info("event processed", "eventRev", evt.Rkey)
339
340 knotstreamEventsProcessed.Inc()
341 return nil
342}
343
344func userAgent() string {
345 return fmt.Sprintf("knotmirror/%s", versioninfo.Short())
346}
347
348func sleepForBackoff(b int) time.Duration {
349 if b == 0 {
350 return 0
351 }
352 if b < 10 {
353 return time.Millisecond * time.Duration((50*b)+rand.Intn(500))
354 }
355 return time.Second * 30
356}