willow.sh
forked from
tangled.org/core
Monorepo for Tangled
8.2 kB
366 lines
1package eventconsumer
2
3import (
4 "context"
5 "encoding/json"
6 "log/slog"
7 "net/http"
8 "sync"
9 "time"
10
11 "tangled.org/core/eventconsumer/cursor"
12 "tangled.org/core/eventstream"
13 "tangled.org/core/log"
14
15 "github.com/avast/retry-go/v4"
16 "github.com/gorilla/websocket"
17)
18
19type ProcessFunc func(ctx context.Context, source Source, event eventstream.Event) error
20
21// server sends a ping every 30s, so any silence longer than this means the
22// connection is half-open, dropped without a close frame.
23// so without a read deadline, ReadMessage only notices when the kernel's
24// tcp keepalive gives up.
25const livenessTimeout = 90 * time.Second
26
27type ConsumerConfig struct {
28 Sources map[Source]struct{}
29 ProcessFunc ProcessFunc
30 RetryInterval time.Duration
31 MaxRetryInterval time.Duration
32 ConnectionTimeout time.Duration
33 WorkerCount int
34 QueueSize int
35 Logger *slog.Logger
36 CursorStore cursor.Store
37
38 Dialer *websocket.Dialer
39 RequestHeader http.Header
40 MaxRetryAttempts uint
41 OnConnectExceeded func(Source, error)
42}
43
44func NewConsumerConfig() *ConsumerConfig {
45 return &ConsumerConfig{
46 Sources: make(map[Source]struct{}),
47 }
48}
49
50type Consumer struct {
51 sourceWg sync.WaitGroup
52 workerWg sync.WaitGroup
53 dialer *websocket.Dialer
54 jobQueue chan job
55 logger *slog.Logger
56
57 // sourcesMu guards sources. It must only be held for short, non-blocking
58 // map operations; never across a blocking call (dial, read, close).
59 sourcesMu sync.Mutex
60 sources map[Source]*sourceState
61
62 cfg ConsumerConfig
63}
64
65type sourceState struct {
66 cancel context.CancelFunc
67 conn *websocket.Conn
68
69 cursorMu sync.Mutex
70 cursorMax int64
71}
72
73type job struct {
74 source Source
75 message []byte
76}
77
78func NewConsumer(cfg ConsumerConfig) *Consumer {
79 if cfg.RetryInterval == 0 {
80 cfg.RetryInterval = 15 * time.Minute
81 }
82 if cfg.ConnectionTimeout == 0 {
83 cfg.ConnectionTimeout = 10 * time.Second
84 }
85 if cfg.WorkerCount <= 0 {
86 cfg.WorkerCount = 5
87 }
88 if cfg.MaxRetryInterval == 0 {
89 cfg.MaxRetryInterval = 1 * time.Hour
90 }
91 if cfg.Logger == nil {
92 cfg.Logger = log.New("consumer")
93 }
94 if cfg.QueueSize == 0 {
95 cfg.QueueSize = 100
96 }
97 if cfg.CursorStore == nil {
98 cfg.CursorStore = &cursor.MemoryStore{}
99 }
100 dialer := cfg.Dialer
101 if dialer == nil {
102 dialer = websocket.DefaultDialer
103 }
104 return &Consumer{
105 cfg: cfg,
106 dialer: dialer,
107 jobQueue: make(chan job, cfg.QueueSize),
108 logger: cfg.Logger,
109 sources: make(map[Source]*sourceState),
110 }
111}
112
113func (c *Consumer) Start(ctx context.Context) {
114 c.cfg.Logger.Info("starting consumer", "config", c.cfg)
115
116 for range c.cfg.WorkerCount {
117 c.workerWg.Add(1)
118 go c.worker(ctx)
119 }
120
121 for source := range c.cfg.Sources {
122 c.AddSource(ctx, source)
123 }
124}
125
126func (c *Consumer) Stop() {
127 // snapshot cancels and conns under lock so we don't hold sourcesMu across Close
128 c.sourcesMu.Lock()
129 cancels := make([]context.CancelFunc, 0, len(c.sources))
130 conns := make([]*websocket.Conn, 0, len(c.sources))
131 for _, st := range c.sources {
132 if st.cancel != nil {
133 cancels = append(cancels, st.cancel)
134 }
135 if st.conn != nil {
136 conns = append(conns, st.conn)
137 }
138 }
139 c.sourcesMu.Unlock()
140
141 for _, cancel := range cancels {
142 cancel()
143 }
144 for _, conn := range conns {
145 conn.Close()
146 }
147
148 c.sourceWg.Wait()
149 close(c.jobQueue)
150 c.workerWg.Wait()
151}
152
153func (c *Consumer) AddSource(ctx context.Context, s Source) {
154 c.sourcesMu.Lock()
155 if _, ok := c.sources[s]; ok {
156 c.sourcesMu.Unlock()
157 c.logger.Info("source already present", "source", s)
158 return
159 }
160 srcCtx, cancel := context.WithCancel(ctx)
161 c.sources[s] = &sourceState{cancel: cancel}
162 c.sourcesMu.Unlock()
163
164 c.sourceWg.Add(1)
165 go c.startConnectionLoop(srcCtx, s)
166}
167
168func (c *Consumer) RemoveSource(s Source) {
169 c.sourcesMu.Lock()
170 st, ok := c.sources[s]
171 if !ok {
172 c.sourcesMu.Unlock()
173 c.logger.Info("source not present", "source", s)
174 return
175 }
176 delete(c.sources, s)
177 cancel := st.cancel
178 conn := st.conn
179 c.sourcesMu.Unlock()
180
181 // release lock before any potentially blocking call
182 if cancel != nil {
183 cancel()
184 }
185 if conn != nil {
186 conn.Close()
187 }
188}
189
190func (c *Consumer) worker(ctx context.Context) {
191 defer c.workerWg.Done()
192 for {
193 select {
194 case <-ctx.Done():
195 return
196 case j, ok := <-c.jobQueue:
197 if !ok {
198 return
199 }
200
201 var ev eventstream.Event
202 err := json.Unmarshal(j.message, &ev)
203 if err != nil {
204 c.logger.Error("error deserializing message", "source", j.source.Key(), "err", err)
205 continue
206 }
207
208 if err := c.cfg.ProcessFunc(ctx, j.source, ev); err != nil {
209 c.logger.Error("error processing message", "source", j.source, "err", err)
210 }
211
212 c.advanceCursor(j.source, ev.Created)
213 }
214 }
215}
216
217func (c *Consumer) advanceCursor(s Source, newCursor int64) {
218 if newCursor == 0 {
219 return
220 }
221 c.sourcesMu.Lock()
222 st, ok := c.sources[s]
223 c.sourcesMu.Unlock()
224 if !ok {
225 return
226 }
227
228 st.cursorMu.Lock()
229 defer st.cursorMu.Unlock()
230 if newCursor <= st.cursorMax {
231 return
232 }
233 st.cursorMax = newCursor
234 c.cfg.CursorStore.Set(s.Key(), newCursor)
235}
236
237func (c *Consumer) startConnectionLoop(ctx context.Context, source Source) {
238 defer c.sourceWg.Done()
239
240 // attempt connection initially
241 err := c.runConnection(ctx, source)
242 if err != nil {
243 c.logger.Error("failed to run connection", "err", err)
244 }
245
246 timer := time.NewTimer(1 * time.Minute)
247 defer timer.Stop()
248
249 // every subsequent attempt is delayed by 1 minute
250 for {
251 select {
252 case <-ctx.Done():
253 return
254 case <-timer.C:
255 err := c.runConnection(ctx, source)
256 if err != nil {
257 c.logger.Error("failed to run connection", "err", err)
258 }
259 timer.Reset(1 * time.Minute)
260 }
261 }
262}
263
264func (c *Consumer) runConnection(ctx context.Context, source Source) error {
265 cursor := c.cfg.CursorStore.Get(source.Key())
266
267 u, err := source.URL(cursor)
268 if err != nil {
269 return err
270 }
271
272 c.logger.Info("connecting", "url", u.String())
273
274 retryOpts := []retry.Option{
275 retry.Attempts(c.cfg.MaxRetryAttempts),
276 retry.DelayType(retry.BackOffDelay),
277 retry.Delay(c.cfg.RetryInterval),
278 retry.MaxDelay(c.cfg.MaxRetryInterval),
279 retry.MaxJitter(c.cfg.RetryInterval / 5),
280 retry.OnRetry(func(n uint, err error) {
281 c.logger.Info("retrying connection",
282 "source", source,
283 "url", u.String(),
284 "attempt", n+1,
285 "err", err,
286 )
287 }),
288 retry.Context(ctx),
289 }
290
291 var conn *websocket.Conn
292
293 err = retry.Do(func() error {
294 connCtx, cancel := context.WithTimeout(ctx, c.cfg.ConnectionTimeout)
295 defer cancel()
296 conn, _, err = c.dialer.DialContext(connCtx, u.String(), c.cfg.RequestHeader)
297 return err
298 }, retryOpts...)
299 if err != nil {
300 if c.cfg.OnConnectExceeded != nil {
301 c.cfg.OnConnectExceeded(source, err)
302 }
303 return err
304 }
305
306 // Register the conn. If the source was removed (or our ctx cancelled)
307 // while we were dialing, drop this conn instead of installing it.
308 c.sourcesMu.Lock()
309 st, ok := c.sources[source]
310 if !ok || ctx.Err() != nil {
311 c.sourcesMu.Unlock()
312 conn.Close()
313 if ctx.Err() != nil {
314 return ctx.Err()
315 }
316 return nil
317 }
318 st.conn = conn
319 c.sourcesMu.Unlock()
320
321 defer func() {
322 // Clear the conn from state, but only if it's still our conn (a
323 // concurrent RemoveSource may have already done it).
324 c.sourcesMu.Lock()
325 if st, ok := c.sources[source]; ok && st.conn == conn {
326 st.conn = nil
327 }
328 c.sourcesMu.Unlock()
329 conn.Close()
330 }()
331
332 c.logger.Info("connected", "source", source)
333
334 conn.SetReadDeadline(time.Now().Add(livenessTimeout))
335 conn.SetPongHandler(func(string) error {
336 return conn.SetReadDeadline(time.Now().Add(livenessTimeout))
337 })
338 conn.SetPingHandler(func(appData string) error {
339 err := conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(10*time.Second))
340 if err != nil {
341 return err
342 }
343 return conn.SetReadDeadline(time.Now().Add(livenessTimeout))
344 })
345
346 for {
347 select {
348 case <-ctx.Done():
349 return nil
350 default:
351 msgType, msg, err := conn.ReadMessage()
352 if err != nil {
353 return err
354 }
355 if msgType != websocket.TextMessage {
356 continue
357 }
358 conn.SetReadDeadline(time.Now().Add(livenessTimeout))
359 select {
360 case c.jobQueue <- job{source: source, message: msg}:
361 case <-ctx.Done():
362 return nil
363 }
364 }
365 }
366}