rta: keep reconnecting until the Conn is closed

rta/conn_test.go

@@ -9,7 +9,6 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
-	"strings"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -338,7 +337,10 @@ func TestReconnectRetriesInterruptedResubscribe(t *testing.T) {
 	}
 }
 
-func TestReconnectClosesAfterPersistentInterruptedResubscribe(t *testing.T) {
+// A socket that keeps dropping mid-handshake must not turn into a closed Conn;
+// the reconnect keeps going until a handshake lands.
+func TestReconnectOutlastsPersistentInterruptedResubscribe(t *testing.T) {
+	shortBackoff(t)
 	srv := newConnTestServer(t)
 	defer srv.Close()
 
@@ -359,20 +361,198 @@ func TestReconnectClosesAfterPersistentInterruptedResubscribe(t *testing.T) {
 		close(done)
 	}()
 
+	// Well past the handful of rounds a bounded budget would allow.
+	waitAtomicUint32(t, &srv.subscribeCount, 8, "subscribe count")
+	if conn.ctx.Err() != nil {
+		t.Fatalf("connection closed during interrupted resubscribes: %v", context.Cause(conn.ctx))
+	}
+	srv.closeSubscribesFrom(0)
+
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("reconnect did not finish once the handshake stopped being interrupted")
+	}
+	if !sub.Active() {
+		t.Fatal("subscription is inactive after the reconnect finally landed")
+	}
+	conn.subscriptionsMu.RLock()
+	_, tracked := conn.subscriptions[sub.ID()]
+	conn.subscriptionsMu.RUnlock()
+	if !tracked {
+		t.Fatal("subscription was not tracked after the reconnect landed")
+	}
+}
+
+// An outage longer than any fixed dial budget must be waited out, not turned
+// into a closed Conn with dead subscriptions.
+func TestReconnectKeepsDialingThroughOutage(t *testing.T) {
+	shortBackoff(t)
+	srv := newConnTestServer(t)
+	defer srv.Close()
+
+	conn := srv.Dial(t)
+	defer conn.Close()
+
+	sub := NewSubscription("test-resource", NopSubscriptionHandler{})
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	if err := conn.Subscribe(ctx, sub); err != nil {
+		t.Fatalf("Subscribe returned error: %v", err)
+	}
+
+	srv.rejectDials.Store(true)
+	done := make(chan struct{})
+	go func() {
+		conn.reconnect()
+		close(done)
+	}()
+
+	waitAtomicUint32(t, &srv.rejectedDials, 10, "rejected dial count")
+	if conn.ctx.Err() != nil {
+		t.Fatalf("connection closed during the outage: %v", context.Cause(conn.ctx))
+	}
+	srv.rejectDials.Store(false)
+
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("reconnect did not finish once dials were accepted again")
+	}
+	if got := srv.subscribeCount.Load(); got != 2 {
+		t.Fatalf("subscribe count = %d, want 2 (initial + one resubscribe)", got)
+	}
+	if !sub.Active() {
+		t.Fatal("subscription is inactive after the outage ended")
+	}
+}
+
+// Closing the Conn during an outage ends the retries and reports the close
+// cause to the subscriptions the reconnect was holding.
+func TestCloseDuringOutageStopsReconnectAndDeactivates(t *testing.T) {
+	shortBackoff(t)
+	srv := newConnTestServer(t)
+	defer srv.Close()
+
+	conn := srv.Dial(t)
+
+	sub := NewSubscription("test-resource", NopSubscriptionHandler{})
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	if err := conn.Subscribe(ctx, sub); err != nil {
+		t.Fatalf("Subscribe returned error: %v", err)
+	}
+
+	srv.rejectDials.Store(true)
+	done := make(chan struct{})
+	go func() {
+		conn.reconnect()
+		close(done)
+	}()
+	waitAtomicUint32(t, &srv.rejectedDials, 3, "rejected dial count")
+
+	if err := conn.Close(); err != nil {
+		t.Fatalf("Close returned error: %v", err)
+	}
 	select {
 	case <-done:
 	case <-time.After(2 * time.Second):
-		t.Fatal("reconnect did not finish after persistent interrupted resubscribe")
+		t.Fatal("reconnect kept running after Close")
 	}
-	if got, want := srv.subscribeCount.Load(), uint32(1+maxResubscribeAttempts); got != want {
-		t.Fatalf("subscribe count = %d, want %d", got, want)
+	if sub.Active() {
+		t.Fatal("subscription is still active after Close")
 	}
-	if err := context.Cause(conn.ctx); err == nil || !strings.Contains(err.Error(), "resubscribe interrupted") {
-		t.Fatalf("connection cause = %v, want resubscribe interrupted", err)
+}
+
+// A Close that lands while a resubscribe handshake is in flight must still
+// end with the subscription deactivated, even though the handshake re-tracks
+// it after Close's own deactivation loop has run.
+func TestCloseDuringResubscribeHandshakeDeactivates(t *testing.T) {
+	shortBackoff(t)
+	srv := newConnTestServer(t)
+	defer srv.Close()
+
+	conn := srv.Dial(t)
+	handler := newBlockingSubscribeHandler(2)
+	sub := NewSubscription("test-resource", handler)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	if err := conn.Subscribe(ctx, sub); err != nil {
+		t.Fatalf("Subscribe returned error: %v", err)
+	}
+
+	done := make(chan struct{})
+	go func() {
+		conn.reconnect()
+		close(done)
+	}()
+	select {
+	case <-handler.entered:
+	case <-time.After(2 * time.Second):
+		t.Fatal("resubscribe handshake did not reach the handler")
+	}
+	if err := conn.Close(); err != nil {
+		t.Fatalf("Close returned error: %v", err)
+	}
+	close(handler.unblock)
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("reconnect did not finish after Close")
 	}
 	if sub.Active() {
-		t.Fatal("subscription is still active after reconnect failure")
+		t.Fatal("subscription is still active on a closed Conn")
+	}
+	conn.subscriptionsMu.RLock()
+	_, tracked := conn.subscriptions[sub.ID()]
+	conn.subscriptionsMu.RUnlock()
+	if tracked {
+		t.Fatal("subscription is still tracked on a closed Conn")
+	}
+}
+
+// Close racing a reconnect must never leave a freshly dialed socket open:
+// every socket the server accepted is closed once Close has returned.
+func TestCloseRacingReconnectLeavesNoOpenSocket(t *testing.T) {
+	shortBackoff(t)
+	srv := newConnTestServer(t)
+	defer srv.Close()
+
+	for range 50 {
+		conn := srv.Dial(t)
+		sub := NewSubscription("test-resource", NopSubscriptionHandler{})
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		if err := conn.Subscribe(ctx, sub); err != nil {
+			cancel()
+			t.Fatalf("Subscribe returned error: %v", err)
+		}
+		cancel()
+		done := make(chan struct{})
+		go func() {
+			conn.reconnect()
+			close(done)
+		}()
+		if err := conn.Close(); err != nil {
+			t.Fatalf("Close returned error: %v", err)
+		}
+		select {
+		case <-done:
+		case <-time.After(2 * time.Second):
+			t.Fatal("reconnect did not finish after Close")
+		}
+		if sub.Active() {
+			t.Fatal("subscription is still active after Close")
+		}
 	}
+	waitAtomicUint32(t, &srv.closeCount, srv.dialCount.Load(), "closed socket count")
+}
+
+// shortBackoff makes reconnect attempts immediate for the rest of the test.
+func shortBackoff(t *testing.T) {
+	t.Helper()
+	old := reconnectBackoff
+	reconnectBackoff = func(int) time.Duration { return time.Millisecond }
+	t.Cleanup(func() { reconnectBackoff = old })
 }
 
 func TestZeroValueSubscriptionUsesNopHandler(t *testing.T) {
@@ -486,6 +666,9 @@ type connTestServer struct {
 	closeSubscribeMin atomic.Uint32
 	closeUnsubscribe  atomic.Bool
 	closeAfterUnsub   atomic.Bool
+	// rejectDials refuses WebSocket upgrades, simulating a service outage.
+	rejectDials   atomic.Bool
+	rejectedDials atomic.Uint32
 }
 
 func newConnTestServer(t *testing.T) *connTestServer {
@@ -543,6 +726,11 @@ func (s *connTestServer) closeAfterUnsubscribeResponse() {
 }
 
 func (s *connTestServer) handle(w http.ResponseWriter, r *http.Request) {
+	if s.rejectDials.Load() {
+		s.rejectedDials.Add(1)
+		http.Error(w, "service unavailable", http.StatusServiceUnavailable)
+		return
+	}
 	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
 		Subprotocols: []string{subprotocol},
 	})

rta/dial.go

@@ -2,7 +2,6 @@ package rta
 
 import (
 	"context"
-	"fmt"
 	"log/slog"
 	"math/rand"
 	"net/http"
@@ -46,6 +45,8 @@ func newConn(c *websocket.Conn, d *dialer) *Conn {
 type dialer struct {
 	log     *slog.Logger
 	options *websocket.DialOptions
+	// backoff returns the wait before reconnect attempt n. Tests shorten it.
+	backoff func(attempt int) time.Duration
 }
 
 func newDialer(client *http.Client, log *slog.Logger) *dialer {
@@ -58,9 +59,13 @@ func newDialer(client *http.Client, log *slog.Logger) *dialer {
 			Subprotocols: []string{subprotocol},
 			HTTPClient:   client,
 		},
+		backoff: reconnectBackoff,
 	}
 }
 
+// reconnectBackoff is the backoff schedule new dialers use; tests shorten it.
+var reconnectBackoff = backoffDuration
+
 // dial establishes a new WebSocket connection.
 func (d *dialer) dial(ctx context.Context) (*websocket.Conn, error) {
 	options := *d.options
@@ -72,43 +77,47 @@ func (d *dialer) dial(ctx context.Context) (*websocket.Conn, error) {
 	return c, nil
 }
 
-// reconnect attempts to establish a WebSocket connection with the RTA service.
-// It retries up to maxDialAttempts times, waiting between each attempt with
-// exponential backoff and jitter. If the context is canceled, it returns the
-// context error immediately.
+// reconnect re-establishes the WebSocket connection, retrying with capped
+// exponential backoff until it succeeds or ctx is done. A service outage can
+// outlast any fixed attempt budget, and a Conn that gave up would strand every
+// subscription until the caller noticed, so only ctx ends the retries.
 func (d *dialer) reconnect(ctx context.Context) (*websocket.Conn, error) {
-	for attempt := range maxDialAttempts {
+	for attempt := 0; ; attempt++ {
 		c, err := d.dial(ctx)
-		if err != nil {
-			sleep := backoffDuration(attempt)
-			d.log.Error("error re-establishing WebSocket connection",
-				slog.Int("attempt", attempt), slog.Int("maxAttempts", maxDialAttempts),
-				slog.Duration("sleep", sleep),
-			)
-			select {
-			case <-time.After(sleep):
-				continue
-			case <-ctx.Done():
-				return nil, ctx.Err()
-			}
+		if err == nil {
+			d.log.Debug("reconnected to RTA service", slog.Int("attempt", attempt))
+			return c, nil
+		}
+		sleep := d.backoff(attempt)
+		// The first failure is news; a long outage should not be an Error stream.
+		level := slog.LevelWarn
+		if attempt == 0 {
+			level = slog.LevelError
+		}
+		d.log.Log(ctx, level, "error re-establishing WebSocket connection",
+			slog.Any("error", err), slog.Int("attempt", attempt), slog.Duration("sleep", sleep),
+		)
+		select {
+		case <-time.After(sleep):
+		case <-ctx.Done():
+			return nil, ctx.Err()
 		}
-		d.log.Debug("reconnected to RTA service", slog.Int("attempt", attempt))
-		return c, nil
 	}
-	return nil, fmt.Errorf("max reconnect attempt (%d) reached", maxDialAttempts)
 }
 
-// backoffDuration returns the duration to wait before the next reconnect attempt.
-// The base duration doubles with each attempt with up to 50% additional jitter.
+// backoffDuration returns the wait before reconnect attempt n: one second
+// doubling per attempt up to maxReconnectBackoff, plus up to 50% jitter.
 func backoffDuration(attempt int) time.Duration {
-	base := time.Second << attempt
+	base := min(time.Second<<min(attempt, maxBackoffShift), maxReconnectBackoff)
 	jitter := time.Duration(rand.Int63n(int64(base / 2)))
 	return base + jitter
 }
 
-// maxDialAttempts is the maximum number of reconnect attempts before
-// [dialer.dialWithBackoff] gives up and returns an error.
-const maxDialAttempts = 4
+const (
+	maxReconnectBackoff = time.Minute
+	// maxBackoffShift bounds the doubling before the cap so the shift never overflows.
+	maxBackoffShift = 6
+)
 
 // subprotocol is the subprotocol used with connectURL, to establish a websocket connection.
 const subprotocol = "rta.xboxlive.com.V2"