Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
3.5 kB · 166 lines
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167package eventstream
import ( "context" "encoding/json" "errors" "log/slog" "net/http" "strconv" "time"
"github.com/gorilla/websocket" "tangled.org/core/notifier")
type Event struct { Rkey string `json:"rkey"` Nsid string `json:"nsid"` EventJson json.RawMessage `json:"event"` Created int64 `json:"created"`}
type Backend interface { GetEvents(cursor int64, limit int) ([]Event, error)}
const ( defaultBatchSize = 100 defaultMaxBatchesPerDrain = 1_000 keepAliveInterval = 30 * time.Second writeDeadline = 10 * time.Second)
var ErrDrainCap = errors.New("eventstream: drain cap reached, reconnect to continue")
var upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024,}
type StreamConfig struct { Backend Backend Notifier *notifier.Notifier Logger *slog.Logger
BatchSize int MaxBatchesPerDrain int}
func (c *StreamConfig) batchSize() int { if c.BatchSize > 0 { return c.BatchSize } return defaultBatchSize}
func (c *StreamConfig) maxBatchesPerDrain() int { if c.MaxBatchesPerDrain > 0 { return c.MaxBatchesPerDrain } return defaultMaxBatchesPerDrain}
func Stream(w http.ResponseWriter, r *http.Request, cfg StreamConfig) error { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { return err } defer conn.Close()
var cursor int64 if raw := r.URL.Query().Get("cursor"); raw != "" { parsed, perr := strconv.ParseInt(raw, 10, 64) if perr != nil { if cfg.Logger != nil { cfg.Logger.Warn("invalid cursor, starting from head", "cursor", raw, "err", perr) } } else { cursor = parsed } }
ch := cfg.Notifier.Subscribe() defer cfg.Notifier.Unsubscribe(ch)
ctx, cancel := context.WithCancel(r.Context()) defer cancel()
go func() { for { if _, _, err := conn.NextReader(); err != nil { cancel() return } } }()
drain := func() error { err := drainUntilShort(conn, cfg, &cursor) if errors.Is(err, ErrDrainCap) { _ = conn.WriteControl( websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseTryAgainLater, "drain cap reached, reconnect to continue"), time.Now().Add(writeDeadline), ) } return err }
if err := drain(); err != nil { return err }
for { select { case <-ctx.Done(): return nil case <-ch: if err := drain(); err != nil { return err } case <-time.After(keepAliveInterval): if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(writeDeadline)); err != nil { return err } } }}
func drainUntilShort(conn *websocket.Conn, cfg StreamConfig, cursor *int64) error { limit := cfg.batchSize() for range cfg.maxBatchesPerDrain() { n, err := streamBatch(conn, cfg, cursor) if err != nil { return err } if n < limit { return nil } } if cfg.Logger != nil { cfg.Logger.Warn("drain hit batch cap", "cursor", *cursor, "cap", cfg.maxBatchesPerDrain()) } return ErrDrainCap}
func streamBatch(conn *websocket.Conn, cfg StreamConfig, cursor *int64) (int, error) { events, err := cfg.Backend.GetEvents(*cursor, cfg.batchSize()) if err != nil { return 0, err } for _, ev := range events { msg, err := json.Marshal(ev) if err != nil { return 0, err } if err := conn.SetWriteDeadline(time.Now().Add(writeDeadline)); err != nil { return 0, err } if err := conn.WriteMessage(websocket.TextMessage, msg); err != nil { return 0, err } *cursor = ev.Created } return len(events), nil}