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