tracker/live.go

package main

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Hub fans out change notifications to connected browsers.
type Hub struct {
	mu      sync.Mutex
	clients map[chan string]struct{}
	version int
}

func NewHub() *Hub { return &Hub{clients: map[chan string]struct{}{}} }

func (h *Hub) subscribe() (chan string, int) {
	ch := make(chan string, 1)
	h.mu.Lock()
	h.clients[ch] = struct{}{}
	v := h.version
	h.mu.Unlock()
	return ch, v
}

func (h *Hub) unsubscribe(ch chan string) {
	h.mu.Lock()
	delete(h.clients, ch)
	h.mu.Unlock()
}

func (h *Hub) Broadcast() {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.version++
	msg := strconv.Itoa(h.version)
	for ch := range h.clients {
		select {
		case ch <- msg:
		default: // client is still processing the previous update; it will refetch anyway
		}
	}
}

// ServeHTTP streams Server-Sent Events.
func (h *Hub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming unsupported", http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no") // disable proxy buffering (nginx ingress)

	ch, v := h.subscribe()
	defer h.unsubscribe(ch)

	fmt.Fprintf(w, "retry: 3000\nevent: hello\ndata: %d\n\n", v)
	flusher.Flush()

	heartbeat := time.NewTicker(25 * time.Second)
	defer heartbeat.Stop()
	for {
		select {
		case <-r.Context().Done():
			return
		case msg := <-ch:
			fmt.Fprintf(w, "event: update\ndata: %s\n\n", msg)
			flusher.Flush()
		case <-heartbeat.C:
			fmt.Fprint(w, ": ping\n\n")
			flusher.Flush()
		}
	}
}

// Watch polls the data directory, reloads on change and notifies clients.
func Watch(ctx context.Context, store *Store, hub *Hub, dataDir string, every time.Duration) {
	last := dirSignature(dataDir)
	t := time.NewTicker(every)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			d := dirSignature(dataDir)
			if d == last {
				continue
			}
			// Editors and os.WriteFile truncate before writing. Wait until the
			// files stop changing so a half-written file is not parsed as an error.
			last = settle(ctx, dataDir, d)
			store.Reload()
			hub.Broadcast()
		}
	}
}

// settle re-reads the directory signature until two reads 100 ms apart match
// (at most about one second), and returns the final signature.
func settle(ctx context.Context, dir, sig string) string {
	for range 10 {
		select {
		case <-ctx.Done():
			return sig
		case <-time.After(100 * time.Millisecond):
		}
		next := dirSignature(dir)
		if next == sig {
			return sig
		}
		sig = next
	}
	return sig
}

func dirSignature(dir string) string {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return ""
	}
	var parts []string
	for _, e := range entries {
		if e.IsDir() {
			if e.Name() == "diagrams" {
				parts = append(parts, "diagrams:"+dirSignature(filepath.Join(dir, e.Name())))
			}
			continue
		}
		if !strings.HasSuffix(e.Name(), ".json") {
			continue
		}
		fi, err := os.Stat(filepath.Join(dir, e.Name()))
		if err != nil {
			continue
		}
		parts = append(parts, e.Name()+fi.ModTime().String()+strconv.FormatInt(fi.Size(), 10))
	}
	sort.Strings(parts)
	return strings.Join(parts, "|")
}