package splitter import ( "context" "fmt" "net" "strconv" "sync" "time" "github.com/bluesky-social/indigo/events" "github.com/gorilla/websocket" "github.com/labstack/echo/v4" "github.com/prometheus/client_golang/prometheus" dto "github.com/prometheus/client_model/go" ) func (s *Splitter) HandleSubscribeRepos(c echo.Context) error { var since *int64 if sinceVal := c.QueryParam("cursor"); sinceVal != "" { sval, err := strconv.ParseInt(sinceVal, 10, 64) if err != nil { return err } since = &sval } // NOTE: the request context outlives the HTTP 101 response; it lives as long as the WebSocket is open, and then get cancelled. That is the behavior we want for this ctx, but should be careful if spawning goroutines which should outlive the WebSocket connection. // https://github.com/bluesky-social/indigo/pull/1023#pullrequestreview-2768335762 ctx, cancel := context.WithCancel(c.Request().Context()) defer cancel() // TODO: authhhh conn, err := websocket.Upgrade(c.Response(), c.Request(), c.Response().Header(), 10<<10, 10<<10) if err != nil { return fmt.Errorf("upgrading websocket: %w", err) } defer conn.Close() lastWriteLk := sync.Mutex{} lastWrite := time.Now() // Start a goroutine to ping the client every 30 seconds to check if it's // still alive. If the client doesn't respond to a ping within 5 seconds, // we'll close the connection and teardown the consumer. go func() { ticker := time.NewTicker(30 * time.Second) defer ticker.Stop() for { select { case <-ticker.C: lastWriteLk.Lock() lw := lastWrite lastWriteLk.Unlock() if time.Since(lw) < 30*time.Second { continue } if err := conn.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(5*time.Second)); err != nil { s.logger.Error("failed to ping client", "err", err) cancel() return } case <-ctx.Done(): return } } }() conn.SetPingHandler(func(message string) error { err := conn.WriteControl(websocket.PongMessage, []byte(message), time.Now().Add(time.Second*60)) if err == websocket.ErrCloseSent { return nil } else if e, ok := err.(net.Error); ok && e.Temporary() { return nil } return err }) // Start a goroutine to read messages from the client and discard them. go func() { for { _, _, err := conn.ReadMessage() if err != nil { s.logger.Error("failed to read message from client", "err", err) cancel() return } } }() ident := c.RealIP() + "-" + c.Request().UserAgent() evts, cleanup, err := s.events.Subscribe(ctx, ident, func(evt *events.XRPCStreamEvent) bool { return true }, since) if err != nil { return err } defer cleanup() // Keep track of the consumer for metrics and admin endpoints consumer := SocketConsumer{ RemoteAddr: c.RealIP(), UserAgent: c.Request().UserAgent(), ConnectedAt: time.Now(), } sentCounter := eventsSentCounter.WithLabelValues(consumer.RemoteAddr, consumer.UserAgent) consumer.EventsSent = sentCounter consumerID := s.registerConsumer(&consumer) defer s.cleanupConsumer(consumerID) s.logger.Info("new consumer", "remote_addr", consumer.RemoteAddr, "user_agent", consumer.UserAgent, "cursor", since, "consumer_id", consumerID, ) activeClientGauge.Inc() defer activeClientGauge.Dec() for { select { case evt, ok := <-evts: if !ok { s.logger.Error("event stream closed unexpectedly") return nil } wc, err := conn.NextWriter(websocket.BinaryMessage) if err != nil { s.logger.Error("failed to get next writer", "err", err) return err } if evt.Preserialized != nil { _, err = wc.Write(evt.Preserialized) } else { err = evt.Serialize(wc) } if err != nil { return fmt.Errorf("failed to write event: %w", err) } if err := wc.Close(); err != nil { s.logger.Warn("failed to flush-close our event write", "err", err) return nil } lastWriteLk.Lock() lastWrite = time.Now() lastWriteLk.Unlock() sentCounter.Inc() case <-ctx.Done(): return nil } } } type SocketConsumer struct { UserAgent string RemoteAddr string ConnectedAt time.Time EventsSent prometheus.Counter } func (s *Splitter) registerConsumer(c *SocketConsumer) uint64 { s.consumersLk.Lock() defer s.consumersLk.Unlock() id := s.nextConsumerID s.nextConsumerID++ s.consumers[id] = c return id } func (s *Splitter) cleanupConsumer(id uint64) { s.consumersLk.Lock() defer s.consumersLk.Unlock() c := s.consumers[id] var m = &dto.Metric{} if err := c.EventsSent.Write(m); err != nil { s.logger.Error("failed to get sent counter", "err", err) } s.logger.Info("consumer disconnected", "consumer_id", id, "remote_addr", c.RemoteAddr, "user_agent", c.UserAgent, "events_sent", m.Counter.GetValue()) delete(s.consumers, id) }