willdot.net
forked from
hailey.at/cocoon
An atproto PDS written in Go
request crawl when websocket dies (#31)
hailey.at
+31
-11
server/handle_sync_subscribe_repos.go
+31
-11
server/handle_sync_subscribe_repos.go
···
1
1
package server
2
2
3
3
import (
4
-
"fmt"
4
+
"context"
5
+
"time"
5
6
6
7
"github.com/bluesky-social/indigo/events"
7
8
"github.com/bluesky-social/indigo/lex/util"
···
10
11
)
11
12
12
13
func (s *Server) handleSyncSubscribeRepos(e echo.Context) error {
14
+
ctx := e.Request().Context()
15
+
logger := s.logger.With("component", "subscribe-repos-websocket")
16
+
13
17
conn, err := websocket.Upgrade(e.Response().Writer, e.Request(), e.Response().Header(), 1<<10, 1<<10)
14
18
if err != nil {
19
+
logger.Error("unable to establish websocket with relay", "err", err)
15
20
return err
16
21
}
17
-
18
-
s.logger.Info("new connection", "ua", e.Request().UserAgent())
19
-
20
-
ctx := e.Request().Context()
21
22
22
23
ident := e.RealIP() + "-" + e.Request().UserAgent()
24
+
logger = logger.With("ident", ident)
25
+
logger.Info("new connection established")
23
26
24
27
evts, cancel, err := s.evtman.Subscribe(ctx, ident, func(evt *events.XRPCStreamEvent) bool {
25
28
return true
···
33
36
for evt := range evts {
34
37
wc, err := conn.NextWriter(websocket.BinaryMessage)
35
38
if err != nil {
36
-
return err
39
+
logger.Error("error writing message to relay", "err", err)
40
+
break
37
41
}
38
42
39
-
var obj util.CBOR
43
+
if ctx.Err() != nil {
44
+
logger.Error("context error", "err", err)
45
+
break
46
+
}
40
47
48
+
var obj util.CBOR
41
49
switch {
42
50
case evt.Error != nil:
43
51
header.Op = events.EvtKindErrorFrame
···
55
63
header.MsgType = "#info"
56
64
obj = evt.RepoInfo
57
65
default:
58
-
return fmt.Errorf("unrecognized event kind")
66
+
logger.Warn("unrecognized event kind")
67
+
return nil
59
68
}
60
69
61
70
if err := header.MarshalCBOR(wc); err != nil {
62
-
return fmt.Errorf("failed to write header: %w", err)
71
+
logger.Error("failed to write header to relay", "err", err)
72
+
break
63
73
}
64
74
65
75
if err := obj.MarshalCBOR(wc); err != nil {
66
-
return fmt.Errorf("failed to write event: %w", err)
76
+
logger.Error("failed to write event to relay", "err", err)
77
+
break
67
78
}
68
79
69
80
if err := wc.Close(); err != nil {
70
-
return fmt.Errorf("failed to flush-close our event write: %w", err)
81
+
logger.Error("failed to flush-close our event write", "err", err)
82
+
break
71
83
}
84
+
}
85
+
86
+
// we should tell the relay to request a new crawl at this point if we got disconnected
87
+
// use a new context since the old one might be cancelled at this point
88
+
ctx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
89
+
defer cancel()
90
+
if err := s.requestCrawl(ctx); err != nil {
91
+
logger.Error("error requesting crawls", "err", err)
72
92
}
73
93
74
94
return nil
+36
-5
server/server.go
+36
-5
server/server.go
···
77
77
passport *identity.Passport
78
78
fallbackProxy string
79
79
80
+
lastRequestCrawl time.Time
81
+
requestCrawlMu sync.Mutex
82
+
80
83
dbName string
81
84
s3Config *S3Config
82
85
}
···
518
521
519
522
go s.backupRoutine()
520
523
524
+
go func() {
525
+
if err := s.requestCrawl(ctx); err != nil {
526
+
s.logger.Error("error requesting crawls", "err", err)
527
+
}
528
+
}()
529
+
530
+
<-ctx.Done()
531
+
532
+
fmt.Println("shut down")
533
+
534
+
return nil
535
+
}
536
+
537
+
func (s *Server) requestCrawl(ctx context.Context) error {
538
+
logger := s.logger.With("component", "request-crawl")
539
+
s.requestCrawlMu.Lock()
540
+
defer s.requestCrawlMu.Unlock()
541
+
542
+
logger.Info("requesting crawl with configured relays")
543
+
544
+
if time.Now().Sub(s.lastRequestCrawl) <= 1*time.Minute {
545
+
return fmt.Errorf("a crawl request has already been made within the last minute")
546
+
}
547
+
521
548
for _, relay := range s.config.Relays {
549
+
logger := logger.With("relay", relay)
550
+
logger.Info("requesting crawl from relay")
522
551
cli := xrpc.Client{Host: relay}
523
-
atproto.SyncRequestCrawl(ctx, &cli, &atproto.SyncRequestCrawl_Input{
552
+
if err := atproto.SyncRequestCrawl(ctx, &cli, &atproto.SyncRequestCrawl_Input{
524
553
Hostname: s.config.Hostname,
525
-
})
554
+
}); err != nil {
555
+
logger.Error("error requesting crawl", "err", err)
556
+
} else {
557
+
logger.Info("crawl requested successfully")
558
+
}
526
559
}
527
560
528
-
<-ctx.Done()
529
-
530
-
fmt.Println("shut down")
561
+
s.lastRequestCrawl = time.Now()
531
562
532
563
return nil
533
564
}