byre/internal/server/sse.go

package server

import (
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"g.fc.run/misc/byre/internal/api"
)

const backfillPage = 500

func (s *Server) StreamEvents(w http.ResponseWriter, r *http.Request, params api.StreamEventsParams) {
	last := int64(0)
	if params.After != nil {
		last = *params.After
	}
	if params.LastEventID != nil {
		last = *params.LastEventID
	}

	rc := http.NewResponseController(w)
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.WriteHeader(http.StatusOK)

	live := s.hub.subscribe()
	defer s.hub.unsubscribe(live)

	for {
		events, err := s.store.EventsAfter(r.Context(), last, backfillPage)
		if err != nil {
			return
		}
		for _, ev := range events {
			if !writeEvent(w, ev) {
				return
			}
			last = ev.Seq
		}
		if len(events) < backfillPage {
			break
		}
	}

	if rc.Flush() != nil {
		return
	}

	ping := time.NewTicker(30 * time.Second)
	defer ping.Stop()

	for {
		select {
		case ev, ok := <-live:
			if !ok {
				return
			}
			if ev.Seq != 0 && ev.Seq <= last {
				continue
			}
			if !writeEvent(w, ev) {
				return
			}
			if ev.Seq != 0 {
				last = ev.Seq
			}

		case <-ping.C:
			_, err := fmt.Fprint(w, ": ping\n\n")
			if err != nil {
				return
			}

		case <-r.Context().Done():
			return
		}

		if rc.Flush() != nil {
			return
		}
	}
}

func writeEvent(w http.ResponseWriter, ev api.Event) bool {
	data, err := json.Marshal(ev)
	if err != nil {
		return false
	}
	if ev.Seq == 0 {
		_, err = fmt.Fprintf(w, "data: %s\n\n", data)
		return err == nil
	}
	_, err = fmt.Fprintf(w, "id: %d\ndata: %s\n\n", ev.Seq, data)
	return err == nil
}