tangled.org
Monorepo for Tangled
tangled.org
1package jetstream
2
3import (
4 "context"
5 "fmt"
6 "log/slog"
7 "os"
8 "os/signal"
9 "sync"
10 "sync/atomic"
11 "syscall"
12 "time"
13
14 "github.com/bluesky-social/jetstream/pkg/client"
15 "github.com/bluesky-social/jetstream/pkg/client/schedulers/sequential"
16 "github.com/bluesky-social/jetstream/pkg/models"
17 "tangled.org/core/log"
18)
19
20type DB interface {
21 GetLastTimeUs() (int64, error)
22 SaveLastTimeUs(int64) error
23}
24
25type Set[T comparable] map[T]struct{}
26
27type JetstreamClient struct {
28 cfg *client.ClientConfig
29 client *client.Client
30 ident string
31 l *slog.Logger
32
33 logDids bool
34 wantedDids Set[string]
35 unfilteredNsids Set[string]
36 db DB
37 waitForDid bool
38 mu sync.RWMutex
39
40 lastSeenUs atomic.Int64
41
42 cancel context.CancelFunc
43 cancelMu sync.Mutex
44}
45
46func (j *JetstreamClient) AddDid(did string) {
47 if did == "" {
48 return
49 }
50
51 if j.logDids {
52 j.l.Info("adding did to in-memory filter", "did", did)
53 }
54 j.mu.Lock()
55 j.wantedDids[did] = struct{}{}
56 j.mu.Unlock()
57}
58
59func (j *JetstreamClient) ExemptCollection(nsid string) {
60 j.mu.Lock()
61 j.unfilteredNsids[nsid] = struct{}{}
62 j.mu.Unlock()
63}
64
65func (j *JetstreamClient) RemoveDid(did string) {
66 if did == "" {
67 return
68 }
69
70 if j.logDids {
71 j.l.Info("removing did from in-memory filter", "did", did)
72 }
73 j.mu.Lock()
74 delete(j.wantedDids, did)
75 j.mu.Unlock()
76}
77
78type processor func(context.Context, *models.Event) error
79
80func (j *JetstreamClient) withDidFilter(processFunc processor) processor {
81 // since this closure references j.WantedDids; it should auto-update
82 // existing instances of the closure when j.WantedDids is mutated
83 return func(ctx context.Context, evt *models.Event) error {
84 j.mu.RLock()
85 // empty filter => all dids allowed
86 matches := len(j.wantedDids) == 0
87 if !matches {
88 if _, ok := j.wantedDids[evt.Did]; ok {
89 matches = true
90 }
91 }
92 if !matches && evt.Commit != nil {
93 if _, ok := j.unfilteredNsids[evt.Commit.Collection]; ok {
94 matches = true
95 }
96 }
97 j.mu.RUnlock()
98
99 var err error
100 if matches {
101 err = processFunc(ctx, evt)
102 }
103
104 j.lastSeenUs.Store(evt.TimeUS + 1)
105 return err
106 }
107}
108
109func NewJetstreamClient(endpoint, ident string, collections []string, cfg *client.ClientConfig, logger *slog.Logger, db DB, waitForDid, logDids bool) (*JetstreamClient, error) {
110 if cfg == nil {
111 cfg = client.DefaultClientConfig()
112 cfg.WebsocketURL = endpoint
113 cfg.WantedCollections = collections
114 }
115
116 return &JetstreamClient{
117 cfg: cfg,
118 ident: ident,
119 db: db,
120 l: logger,
121 wantedDids: make(map[string]struct{}),
122 unfilteredNsids: make(map[string]struct{}),
123
124 logDids: logDids,
125
126 // This will make the goroutine in StartJetstream wait until
127 // j.wantedDids has been populated, typically using addDids.
128 waitForDid: waitForDid,
129 }, nil
130}
131
132// StartJetstream starts the jetstream client and processes events using the provided processFunc.
133// The client persists the last time_us cursor itself via the DB it was constructed with.
134func (j *JetstreamClient) StartJetstream(ctx context.Context, processFunc func(context.Context, *models.Event) error) error {
135 logger := j.l
136
137 sched := sequential.NewScheduler(j.ident, logger, j.withDidFilter(processFunc))
138
139 client, err := client.NewClient(j.cfg, logger, sched)
140 if err != nil {
141 return fmt.Errorf("failed to create jetstream client: %w", err)
142 }
143 j.client = client
144
145 go func() {
146 if j.waitForDid {
147 for {
148 j.mu.RLock()
149 hasDid := len(j.wantedDids) != 0
150 j.mu.RUnlock()
151 if hasDid {
152 break
153 }
154 time.Sleep(time.Second)
155 }
156 }
157 logger.Info("done waiting for did")
158
159 go j.periodicLastTimeSave(ctx)
160 j.saveIfKilled(ctx)
161
162 j.connectAndRead(ctx)
163 }()
164
165 return nil
166}
167
168func (j *JetstreamClient) connectAndRead(ctx context.Context) {
169 l := log.FromContext(ctx)
170 for {
171 cursor := j.resumeCursor(ctx)
172
173 connCtx, cancel := context.WithCancel(ctx)
174 j.cancelMu.Lock()
175 j.cancel = cancel
176 j.cancelMu.Unlock()
177
178 if err := j.client.ConnectAndRead(connCtx, cursor); err != nil {
179 l.Error("error reading jetstream", "error", err)
180 cancel()
181 continue
182 }
183
184 select {
185 case <-ctx.Done():
186 l.Info("context done, stopping jetstream")
187 return
188 case <-connCtx.Done():
189 l.Info("connection context done, reconnecting")
190 continue
191 }
192 }
193}
194
195// save cursor periodically
196func (j *JetstreamClient) periodicLastTimeSave(ctx context.Context) {
197 ticker := time.NewTicker(time.Minute)
198 defer ticker.Stop()
199
200 for {
201 select {
202 case <-ctx.Done():
203 return
204 case <-ticker.C:
205 if seen := j.lastSeenUs.Load(); seen != 0 {
206 if err := j.db.SaveLastTimeUs(seen); err != nil {
207 log.FromContext(ctx).Error("failed to save cursor", "error", err)
208 }
209 }
210 }
211 }
212}
213
214func (j *JetstreamClient) resumeCursor(ctx context.Context) *int64 {
215 if seen := j.lastSeenUs.Load(); seen != 0 {
216 return &seen
217 }
218 return j.getLastTimeUs(ctx)
219}
220
221func (j *JetstreamClient) getLastTimeUs(ctx context.Context) *int64 {
222 l := log.FromContext(ctx)
223 lastTimeUs, err := j.db.GetLastTimeUs()
224 if err != nil {
225 l.Warn("couldn't get last time us, starting from now", "error", err)
226 lastTimeUs = time.Now().UnixMicro()
227 if err = j.db.SaveLastTimeUs(lastTimeUs); err != nil {
228 l.Error("failed to save last time us", "error", err)
229 }
230 }
231
232 l.Info("found last time_us", "time_us", lastTimeUs)
233 return &lastTimeUs
234}
235
236func (j *JetstreamClient) saveIfKilled(ctx context.Context) context.Context {
237 ctxWithCancel, cancel := context.WithCancel(ctx)
238
239 sigChan := make(chan os.Signal, 1)
240
241 signal.Notify(sigChan,
242 syscall.SIGINT,
243 syscall.SIGTERM,
244 syscall.SIGQUIT,
245 syscall.SIGHUP,
246 syscall.SIGKILL,
247 syscall.SIGSTOP,
248 )
249
250 go func() {
251 sig := <-sigChan
252 j.l.Info("Received signal, initiating graceful shutdown", "signal", sig)
253
254 if seen := j.lastSeenUs.Load(); seen != 0 {
255 if err := j.db.SaveLastTimeUs(seen); err != nil {
256 j.l.Error("Failed to save last time during shutdown", "error", err)
257 }
258 j.l.Info("Saved lastTimeUs before shutdown", "lastTimeUs", seen)
259 }
260
261 j.cancelMu.Lock()
262 if j.cancel != nil {
263 j.cancel()
264 }
265 j.cancelMu.Unlock()
266
267 cancel()
268
269 os.Exit(0)
270 }()
271
272 return ctxWithCancel
273}