From 290ff085e268e9bacc95479cefdb663a84cd5e27 Mon Sep 17 00:00:00 2001 From: Luna Date: Thu, 23 Oct 2025 22:21:40 -0300 Subject: [PATCH] add watchdog for inactive jetstream --- main.go | 193 +++++++++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 157 insertions(+), 36 deletions(-) diff --git a/main.go b/main.go index 69af466..5bfd18b 100644 --- a/main.go +++ b/main.go @@ -1,6 +1,7 @@ package main import ( + "context" "fmt" "log/slog" "net/http" @@ -25,9 +26,10 @@ var DEFAULT_POOL = []string{ // Broadcaster manages subscribers to Jetstream events type Broadcaster struct { - listeners []chan []byte - mu sync.Mutex - connected atomic.Bool + listeners []chan []byte + mu sync.Mutex + connected atomic.Bool + lastMessageTime atomic.Int64 // Unix timestamp in seconds } // Subscribe returns a new channel that will receive Jetstream events @@ -56,6 +58,8 @@ func (b *Broadcaster) Unsubscribe(ch chan []byte) { } func (b *Broadcaster) Broadcast(message []byte) { + b.lastMessageTime.Store(time.Now().Unix()) + b.mu.Lock() defer b.mu.Unlock() @@ -145,28 +149,123 @@ func handleSubscribe(broadcaster *Broadcaster) http.HandlerFunc { } } +// raceUpstreams connects to all upstreams simultaneously and returns the first one to deliver a message +func raceUpstreams(pool []string) (string, error) { + slog.Info("Racing upstreams to find fastest message delivery") + + type result struct { + url string + duration time.Duration + } + + results := make(chan result, len(pool)) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + for _, url := range pool { + go func(u string) { + start := time.Now() + conn, _, err := websocket.DefaultDialer.DialContext(ctx, u+"/subscribe", nil) + if err != nil { + slog.Debug("Failed to connect during race", slog.String("url", u), slog.Any("error", err)) + return + } + defer conn.Close() + + // Wait for first message + _, _, err = conn.ReadMessage() + if err != nil { + slog.Debug("Failed to read message during race", slog.String("url", u), slog.Any("error", err)) + return + } + + duration := time.Since(start) + select { + case results <- result{url: u, duration: duration}: + case <-ctx.Done(): + } + }(url) + } + + select { + case res := <-results: + slog.Info("Race winner", slog.String("url", res.url), slog.Duration("time_to_first_message", res.duration)) + return res.url, nil + case <-ctx.Done(): + return "", fmt.Errorf("no upstream delivered a message within timeout") + } +} + +// watchdog monitors message activity and triggers reconnection if stalled +func watchdog(broadcaster *Broadcaster, trigger chan struct{}) { + ticker := time.NewTicker(5 * time.Second) + defer ticker.Stop() + + for range ticker.C { + if !broadcaster.connected.Load() { + continue + } + + lastMsg := broadcaster.lastMessageTime.Load() + if lastMsg == 0 { + // No messages received yet + continue + } + + timeSinceLastMsg := time.Since(time.Unix(lastMsg, 0)) + if timeSinceLastMsg > 20*time.Second { + slog.Warn("No messages received", slog.Duration("duration", timeSinceLastMsg)) + select { + case trigger <- struct{}{}: + // Trigger sent + default: + // Trigger already pending + } + } + } +} + // connectToUpstream maintains a connection to the upstream websocket and broadcasts messages -func connectToUpstream(pool []string, broadcaster *Broadcaster) { +func connectToUpstream(pool []string, broadcaster *Broadcaster, watchdogTrigger <-chan struct{}) { backoff := 50 * time.Millisecond maxBackoff := 20 * time.Second var currentUpstream string + var raceTriggered bool for { - // Find best upstream (re-evaluate on each connection attempt) - bestUpstream, err := findBestUpstream(pool) - if err != nil { - slog.Error("Failed to find best upstream", slog.Any("error", err)) - time.Sleep(backoff) - backoff *= 2 - if backoff > maxBackoff { - backoff = maxBackoff + if raceTriggered { + slog.Info("Watchdog triggered, racing upstreams") + // Watchdog triggered - race all upstreams + bestUpstream, err := raceUpstreams(pool) + if err != nil { + slog.Error("Failed to race upstreams", slog.Any("error", err)) + time.Sleep(backoff) + backoff *= 2 + if backoff > maxBackoff { + backoff = maxBackoff + } + continue } - continue - } - - if bestUpstream != currentUpstream { - slog.Info("Switching to new upstream", slog.String("url", bestUpstream)) currentUpstream = bestUpstream + backoff = 50 * time.Millisecond // Reset backoff + raceTriggered = false + } else { + // Find best upstream (re-evaluate on each connection attempt) + bestUpstream, err := findBestUpstream(pool) + if err != nil { + slog.Error("Failed to find best upstream", slog.Any("error", err)) + time.Sleep(backoff) + backoff *= 2 + if backoff > maxBackoff { + backoff = maxBackoff + } + continue + } + + if bestUpstream != currentUpstream { + slog.Info("Switching to new upstream", slog.String("url", bestUpstream)) + currentUpstream = bestUpstream + } } slog.Info("Connecting to upstream", slog.String("url", currentUpstream)) @@ -185,30 +284,49 @@ func connectToUpstream(pool []string, broadcaster *Broadcaster) { slog.Info("Connected to upstream", slog.String("url", currentUpstream)) broadcaster.connected.Store(true) - backoff = time.Second // Reset backoff on successful connection + broadcaster.lastMessageTime.Store(time.Now().Unix()) + backoff = 50 * time.Millisecond // Reset backoff on successful connection // Read messages from upstream and broadcast them - for { - messageType, message, err := conn.ReadMessage() - if err != nil { - slog.Error("Error reading from upstream", slog.Any("error", err)) - broadcaster.connected.Store(false) - conn.Close() - break + readDone := make(chan struct{}) + go func() { + defer close(readDone) + for { + messageType, message, err := conn.ReadMessage() + if err != nil { + slog.Error("Error reading from upstream", slog.Any("error", err)) + return + } + + // Only broadcast text/binary messages + if messageType == websocket.TextMessage || messageType == websocket.BinaryMessage { + broadcaster.Broadcast(message) + } } + }() - // Only broadcast text/binary messages - if messageType == websocket.TextMessage || messageType == websocket.BinaryMessage { - broadcaster.Broadcast(message) - } + // Wait for either read error or watchdog trigger + select { + case <-readDone: + // Normal disconnection + case <-watchdogTrigger: + // Watchdog triggered disconnection + slog.Info("Watchdog triggered disconnection") + raceTriggered = true + conn.Close() + <-readDone // Wait for read goroutine to finish } - // Connection lost, will re-evaluate best upstream on next iteration - slog.Info("Connection lost, finding new upstream", slog.Duration("backoff", backoff)) - time.Sleep(backoff) - backoff *= 2 - if backoff > maxBackoff { - backoff = maxBackoff + broadcaster.connected.Store(false) + + if !raceTriggered { + // Connection lost, will re-evaluate best upstream on next iteration + slog.Info("Connection lost, finding new upstream", slog.Duration("backoff", backoff)) + time.Sleep(backoff) + backoff *= 2 + if backoff > maxBackoff { + backoff = maxBackoff + } } } } @@ -288,7 +406,10 @@ func main() { // Create broadcaster and start upstream connection // connectToUpstream will continuously find the best upstream and reconnect on failures broadcaster := &Broadcaster{} - go connectToUpstream(pool, broadcaster) + watchdogTrigger := make(chan struct{}, 1) + + go watchdog(broadcaster, watchdogTrigger) + go connectToUpstream(pool, broadcaster, watchdogTrigger) // Setup HTTP server http.HandleFunc("/", handleHealth(broadcaster)) -- 2.51.2