Something went wrong. Try again.
Monorepo for Tangled
Something went wrong. Try again.
7.0 kB · 326 lines
Go
at eti/checkbox
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327package eventconsumer
import ( "context" "encoding/json" "log/slog" "math/rand" "net/url" "sync" "time"
"tangled.org/core/eventconsumer/cursor" "tangled.org/core/log"
"github.com/avast/retry-go/v4" "github.com/gorilla/websocket")
type ProcessFunc func(ctx context.Context, source Source, message Message) error
type Message struct { Rkey string Nsid string Created int64 `json:"created"` EventJson json.RawMessage `json:"event"`}
type ConsumerConfig struct { Sources map[Source]struct{} ProcessFunc ProcessFunc RetryInterval time.Duration MaxRetryInterval time.Duration ConnectionTimeout time.Duration WorkerCount int QueueSize int Logger *slog.Logger Dev bool CursorStore cursor.Store}
func NewConsumerConfig() *ConsumerConfig { return &ConsumerConfig{ Sources: make(map[Source]struct{}), }}
type Source interface { // url to start streaming events from Url(cursor int64, dev bool) (*url.URL, error) // cache key for cursor storage Key() string}
type Consumer struct { wg sync.WaitGroup dialer *websocket.Dialer jobQueue chan job logger *slog.Logger randSource *rand.Rand
// sourcesMu guards sources. It must only be held for short, non-blocking // map operations; never across a blocking call (dial, read, close). sourcesMu sync.Mutex sources map[Source]*sourceState
cfg ConsumerConfig}
type sourceState struct { cancel context.CancelFunc conn *websocket.Conn}
type job struct { source Source message []byte}
func NewConsumer(cfg ConsumerConfig) *Consumer { if cfg.RetryInterval == 0 { cfg.RetryInterval = 15 * time.Minute } if cfg.ConnectionTimeout == 0 { cfg.ConnectionTimeout = 10 * time.Second } if cfg.WorkerCount <= 0 { cfg.WorkerCount = 5 } if cfg.MaxRetryInterval == 0 { cfg.MaxRetryInterval = 1 * time.Hour } if cfg.Logger == nil { cfg.Logger = log.New("consumer") } if cfg.QueueSize == 0 { cfg.QueueSize = 100 } if cfg.CursorStore == nil { cfg.CursorStore = &cursor.MemoryStore{} } return &Consumer{ cfg: cfg, dialer: websocket.DefaultDialer, jobQueue: make(chan job, cfg.QueueSize), // buffered job queue logger: cfg.Logger, randSource: rand.New(rand.NewSource(time.Now().UnixNano())), sources: make(map[Source]*sourceState), }}
func (c *Consumer) Start(ctx context.Context) { c.cfg.Logger.Info("starting consumer", "config", c.cfg)
// start workers for range c.cfg.WorkerCount { c.wg.Add(1) go c.worker(ctx) }
// start streaming for source := range c.cfg.Sources { c.AddSource(ctx, source) }}
func (c *Consumer) Stop() { // snapshot conns under lock so we don't hold sourcesMu across Close c.sourcesMu.Lock() conns := make([]*websocket.Conn, 0, len(c.sources)) for _, st := range c.sources { if st.conn != nil { conns = append(conns, st.conn) } } c.sourcesMu.Unlock()
for _, conn := range conns { conn.Close() }
c.wg.Wait() close(c.jobQueue)}
func (c *Consumer) AddSource(ctx context.Context, s Source) { c.sourcesMu.Lock() if _, ok := c.sources[s]; ok { c.sourcesMu.Unlock() c.logger.Info("source already present", "source", s) return } srcCtx, cancel := context.WithCancel(ctx) c.sources[s] = &sourceState{cancel: cancel} c.sourcesMu.Unlock()
c.wg.Add(1) go c.startConnectionLoop(srcCtx, s)}
func (c *Consumer) RemoveSource(s Source) { c.sourcesMu.Lock() st, ok := c.sources[s] if !ok { c.sourcesMu.Unlock() c.logger.Info("source not present", "source", s) return } delete(c.sources, s) cancel := st.cancel conn := st.conn c.sourcesMu.Unlock()
// release lock before any potentially blocking call if cancel != nil { cancel() } if conn != nil { conn.Close() }}
func (c *Consumer) worker(ctx context.Context) { defer c.wg.Done() for { select { case <-ctx.Done(): return case j, ok := <-c.jobQueue: if !ok { return }
var msg Message err := json.Unmarshal(j.message, &msg) if err != nil { c.logger.Error("error deserializing message", "source", j.source.Key(), "err", err) return }
if err := c.cfg.ProcessFunc(ctx, j.source, msg); err != nil { c.logger.Error("error processing message", "source", j.source, "err", err) }
cursorVal := msg.Created if cursorVal == 0 { cursorVal = time.Now().UnixNano() } c.cfg.CursorStore.Set(j.source.Key(), cursorVal) } }}
func (c *Consumer) startConnectionLoop(ctx context.Context, source Source) { defer c.wg.Done()
// attempt connection initially err := c.runConnection(ctx, source) if err != nil { c.logger.Error("failed to run connection", "err", err) }
timer := time.NewTimer(1 * time.Minute) defer timer.Stop()
// every subsequent attempt is delayed by 1 minute for { select { case <-ctx.Done(): return case <-timer.C: err := c.runConnection(ctx, source) if err != nil { c.logger.Error("failed to run connection", "err", err) } timer.Reset(1 * time.Minute) } }}
func (c *Consumer) runConnection(ctx context.Context, source Source) error { cursor := c.cfg.CursorStore.Get(source.Key())
u, err := source.Url(cursor, c.cfg.Dev) if err != nil { return err }
c.logger.Info("connecting", "url", u.String())
retryOpts := []retry.Option{ retry.Attempts(0), // infinite attempts retry.DelayType(retry.BackOffDelay), retry.Delay(c.cfg.RetryInterval), retry.MaxDelay(c.cfg.MaxRetryInterval), retry.MaxJitter(c.cfg.RetryInterval / 5), retry.OnRetry(func(n uint, err error) { c.logger.Info("retrying connection", "source", source, "url", u.String(), "attempt", n+1, "err", err, ) }), retry.Context(ctx), }
var conn *websocket.Conn
err = retry.Do(func() error { connCtx, cancel := context.WithTimeout(ctx, c.cfg.ConnectionTimeout) defer cancel() conn, _, err = c.dialer.DialContext(connCtx, u.String(), nil) return err }, retryOpts...) if err != nil { return err }
// Register the conn. If the source was removed (or our ctx cancelled) // while we were dialing, drop this conn instead of installing it. c.sourcesMu.Lock() st, ok := c.sources[source] if !ok || ctx.Err() != nil { c.sourcesMu.Unlock() conn.Close() if ctx.Err() != nil { return ctx.Err() } return nil } st.conn = conn c.sourcesMu.Unlock()
defer func() { // Clear the conn from state, but only if it's still our conn (a // concurrent RemoveSource may have already done it). c.sourcesMu.Lock() if st, ok := c.sources[source]; ok && st.conn == conn { st.conn = nil } c.sourcesMu.Unlock() conn.Close() }()
c.logger.Info("connected", "source", source)
for { select { case <-ctx.Done(): return nil default: msgType, msg, err := conn.ReadMessage() if err != nil { return err } if msgType != websocket.TextMessage { continue } select { case c.jobQueue <- job{source: source, message: msg}: case <-ctx.Done(): return nil } } }}