diff --git a/stream.go b/stream.go index eb60176..0a86bce 100644 --- a/stream.go +++ b/stream.go @@ -59,6 +59,14 @@ type StreamOptions struct { // MaxReconnectAttempts is the number of consecutive reconnect attempts // before the stream terminates with an error. Default: 10. A negative value // retries forever. + // + // Consecutive means what it says: the count and the backoff are both reset + // once a session reaches healthy progress — the subscription confirmed and + // at least one further server frame received on that connection. A stream + // that recovers, runs, and drops again hours later starts from a full + // budget. A server that accepts, or confirms, and then immediately drops + // never reaches healthy progress, so the cap still terminates a rapid + // flap. MaxReconnectAttempts int } @@ -98,6 +106,9 @@ func WithStreamMaxReconnectDelay(d time.Duration) StreamOption { // WithStreamMaxReconnectAttempts sets the maximum number of consecutive // reconnect attempts before the stream gives up. A negative value retries // forever. +// +// The budget is restored whenever a session reaches healthy progress; see +// StreamOptions.MaxReconnectAttempts for what that means. func WithStreamMaxReconnectAttempts(n int) StreamOption { return func(o *StreamOptions) { o.MaxReconnectAttempts = n @@ -350,7 +361,15 @@ func (s *PriceStream) run() { attempt := 0 for { - err := s.connectAndRead() + healthy, err := s.connectAndRead() + + // A session that got far enough to prove itself restores the budget. + // MaxReconnectAttempts is documented as consecutive attempts, so a + // recovery that worked must not leave the next disconnect closer to + // the cap than the first one was. + if healthy { + attempt = 0 + } // Context cancelled / Close() called: clean shutdown, no error. if s.ctx.Err() != nil { @@ -386,24 +405,34 @@ func (s *PriceStream) run() { } // connectAndRead dials, performs the ActionCable handshake, and reads frames -// until the connection drops or the context is cancelled. A successful -// subscription resets the caller's backoff counter via the returned nil-reset -// behaviour: callers treat any return as a disconnect and reconnect. -func (s *PriceStream) connectAndRead() error { +// until the connection drops or the context is cancelled. +// +// It reports whether the session reached healthy progress, which is what the +// run loop uses to restore the reconnect budget. Healthy means the +// subscription was confirmed AND at least one further server frame — an +// ActionCable ping or a channel message — arrived on the same connection. +// +// The boundary is deliberately not the TCP dial or the handshake alone. A +// server that accepts and hangs up, or that confirms and immediately drops, +// would reset the budget on every pass and turn MaxReconnectAttempts into an +// unbounded rapid-flap loop. Requiring one frame past the confirmation proves +// a live session rather than a reachable socket, and it is a frame the server +// sends on its own: ActionCable pings every few seconds. +func (s *PriceStream) connectAndRead() (healthy bool, err error) { header := http.Header{} header.Set("Authorization", "Token "+s.apiKey) header.Set("User-Agent", fmt.Sprintf("oilpriceapi-go/%s", Version)) conn, err := s.dial(s.ctx, s.url, header) if err != nil { - return err + return false, err } s.mu.Lock() if s.closed { s.mu.Unlock() _ = conn.Close() - return s.ctx.Err() + return false, s.ctx.Err() } s.conn = conn s.mu.Unlock() @@ -424,20 +453,47 @@ func (s *PriceStream) connectAndRead() error { "identifier": s.identifier, }) if err := conn.Write(s.ctx, subscribe); err != nil { - return err + return false, err } + confirmed := false + for { data, err := conn.Read(s.ctx) if err != nil { - return err + return healthy, err } - if fatal := s.handleFrame(data); fatal != nil { - return fatal + + kind, fatal := s.handleFrame(data) + if fatal != nil { + return healthy, fatal + } + + switch { + case kind == frameConfirmSubscription: + confirmed = true + case confirmed && (kind == frameServerKeepalive || kind == frameChannelMessage): + healthy = true } } } +// frameKind classifies an ActionCable transport frame for the reconnect +// accounting in connectAndRead. +type frameKind int + +const ( + // frameIgnored is a frame that proves nothing: malformed JSON, an empty + // channel message, or a server-initiated disconnect notice. + frameIgnored frameKind = iota + // frameServerKeepalive is a welcome or ping. + frameServerKeepalive + // frameConfirmSubscription is the subscription acknowledgement. + frameConfirmSubscription + // frameChannelMessage is a message delivered on the subscribed channel. + frameChannelMessage +) + // cableFrame is the ActionCable transport envelope. type cableFrame struct { Type string `json:"type"` @@ -446,35 +502,37 @@ type cableFrame struct { } // handleFrame decodes one ActionCable transport frame and dispatches channel -// messages. It returns a non-nil error only for fatal conditions (subscription -// rejected) that must terminate the stream without reconnecting. -func (s *PriceStream) handleFrame(data []byte) error { +// messages. It returns the frame's kind, which connectAndRead uses for its +// healthy-progress accounting, and a non-nil error only for fatal conditions +// (subscription rejected) that must terminate the stream without reconnecting. +func (s *PriceStream) handleFrame(data []byte) (frameKind, error) { var frame cableFrame if err := json.Unmarshal(data, &frame); err != nil { // Ignore malformed frames rather than tear down the stream. - return nil + return frameIgnored, nil } switch frame.Type { case "ping", "welcome": // Heartbeat / transport handshake — nothing to do. - return nil + return frameServerKeepalive, nil case "confirm_subscription": - return nil + return frameConfirmSubscription, nil case "reject_subscription": - return &StreamRejectedError{} + return frameIgnored, &StreamRejectedError{} case "disconnect": // Server-initiated disconnect; let the read loop's next Read fail and - // drive reconnect. - return nil + // drive reconnect. It is not healthy progress: the server is telling + // us to go away. + return frameIgnored, nil } // Channel message: payload lives under "message". if len(frame.Message) == 0 { - return nil + return frameIgnored, nil } s.dispatch(frame.Message) - return nil + return frameChannelMessage, nil } // messageType peeks at the channel message "type" field. diff --git a/stream_reconnect_budget_test.go b/stream_reconnect_budget_test.go new file mode 100644 index 0000000..c96b134 --- /dev/null +++ b/stream_reconnect_budget_test.go @@ -0,0 +1,261 @@ +package oilpriceapi + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "runtime" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/gorilla/websocket" +) + +// MaxReconnectAttempts is documented as "the number of consecutive reconnect +// attempts before the stream terminates". The run loop incremented one counter +// and never reset it, so it was really a lifetime disconnect count: a stream +// that recovered fully, ran healthily, and then dropped again hours later +// spent budget it had already earned back, and eventually terminated a working +// subscription. +// +// The boundary chosen for "recovered" is deliberately not the TCP dial — that +// resets on a server that accepts and hangs up, which is an unbounded rapid +// flap. A session counts as healthy only once the ActionCable subscription has +// been confirmed AND at least one further server frame (a ping, or a channel +// message) has arrived on it. That proves a live session beyond the handshake, +// and a server that confirms and immediately drops never reaches it. + +// scriptedCable is a cable server whose per-connection behaviour is scripted +// by connection index, so a test can stage healthy sessions followed by +// failures. +type scriptedCable struct { + upgrader websocket.Upgrader + + connections int32 + subscribes int32 + + // healthySessions is the number of leading connections that complete the + // handshake and deliver a post-confirmation ping before hanging up. + // Connections past that are accepted and immediately closed. + healthySessions int32 + + // postConfirmPing controls whether a confirmed session also sends the + // frame that makes it count as healthy. False models a flapping server + // that confirms and drops. + postConfirmPing bool +} + +func (s *scriptedCable) handler(w http.ResponseWriter, r *http.Request) { + n := atomic.AddInt32(&s.connections, 1) + + conn, err := s.upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + + if s.healthySessions >= 0 && n > s.healthySessions { + // Past the scripted healthy sessions: accept and hang up. + return + } + + _ = conn.WriteJSON(map[string]string{"type": "welcome"}) + + for { + _, data, err := conn.ReadMessage() + if err != nil { + return + } + var cmd struct { + Command string `json:"command"` + Identifier string `json:"identifier"` + } + if err := json.Unmarshal(data, &cmd); err != nil { + continue + } + if cmd.Command != "subscribe" { + continue + } + atomic.AddInt32(&s.subscribes, 1) + _ = conn.WriteJSON(map[string]string{"type": "confirm_subscription", "identifier": cmd.Identifier}) + + if s.postConfirmPing { + // The frame that makes this a healthy session rather than a bare + // handshake. + _ = conn.WriteJSON(map[string]string{"type": "ping"}) + } + // Drop the session. + return + } +} + +func newScriptedStream(t *testing.T, srv *scriptedCable, opts StreamOptions) (*PriceStream, *httptest.Server) { + t.Helper() + srv.upgrader.CheckOrigin = func(*http.Request) bool { return true } + httpSrv := httptest.NewServer(http.HandlerFunc(srv.handler)) + + client := NewClient("k", WithBaseURL(httpSrv.URL)) + stream, err := client.newStream(context.Background(), opts, defaultDialer) + if err != nil { + httpSrv.Close() + t.Fatalf("newStream: %v", err) + } + return stream, httpSrv +} + +// A healthy session must restore the budget. With MaxReconnectAttempts=1 and +// three healthy sessions staged, all three must be reached: each recovery +// earns the budget back. Before the fix the stream terminated after the +// second, having spent a budget that intervening healthy sessions never +// refunded. +func TestStreamHealthySessionResetsReconnectBudget(t *testing.T) { + srv := &scriptedCable{healthySessions: 3, postConfirmPing: true} + stream, httpSrv := newScriptedStream(t, srv, StreamOptions{ + AutoReconnect: true, + ReconnectDelay: time.Millisecond, + MaxReconnectDelay: 5 * time.Millisecond, + MaxReconnectAttempts: 1, + }) + defer httpSrv.Close() + defer stream.Close() + + // Drain until the stream terminates. + deadline := time.After(5 * time.Second) + for { + select { + case _, ok := <-stream.Updates(): + if !ok { + goto done + } + case <-deadline: + t.Fatal("timeout waiting for the stream to terminate") + } + } +done: + + if got := atomic.LoadInt32(&srv.subscribes); got < 3 { + t.Fatalf("confirmed subscriptions = %d, want the 3 staged healthy sessions; "+ + "the reconnect budget was not restored by a healthy session (terminating error: %v)", + got, stream.Err()) + } + if stream.Err() == nil { + t.Fatal("expected a terminating error once the server stopped recovering") + } + if !strings.Contains(stream.Err().Error(), "reconnect failed") { + t.Fatalf("unexpected terminating error: %v", stream.Err()) + } +} + +// The reset must not create an unbounded rapid-flap loop. A server that +// confirms the subscription and immediately drops has not delivered a healthy +// session, so the budget is not restored and the cap still terminates the +// stream. +func TestStreamConfirmThenImmediateDropDoesNotResetBudget(t *testing.T) { + srv := &scriptedCable{healthySessions: -1, postConfirmPing: false} + stream, httpSrv := newScriptedStream(t, srv, StreamOptions{ + AutoReconnect: true, + ReconnectDelay: time.Millisecond, + MaxReconnectDelay: 5 * time.Millisecond, + MaxReconnectAttempts: 2, + }) + defer httpSrv.Close() + defer stream.Close() + + deadline := time.After(5 * time.Second) + for { + select { + case _, ok := <-stream.Updates(): + if !ok { + goto done + } + case <-deadline: + t.Fatal("a confirm-then-drop server flapped without ever exhausting the reconnect budget") + } + } +done: + + if stream.Err() == nil { + t.Fatal("expected the reconnect budget to be exhausted by a flapping server") + } + // 1 initial session + at most MaxReconnectAttempts recoveries. + if got := atomic.LoadInt32(&srv.connections); got > 3 { + t.Fatalf("server saw %d connections, want at most 3 (1 initial + 2 attempts)", got) + } +} + +// A rejected subscription is fatal and must never be retried, healthy-session +// accounting or not. +func TestStreamRejectedSubscriptionStillTerminatesWithoutRetrying(t *testing.T) { + srv := &cableServer{rejectSubscribe: true} + stream, httpSrv := newTestStream(t, srv, "k", + WithStreamMaxReconnectAttempts(5), + WithStreamReconnectDelay(time.Millisecond), + ) + defer httpSrv.Close() + defer stream.Close() + + select { + case _, ok := <-stream.Updates(): + if ok { + t.Fatal("did not expect updates from a rejected subscription") + } + case <-time.After(3 * time.Second): + t.Fatal("timeout waiting for a rejected subscription to terminate the stream") + } + + var rejected *StreamRejectedError + if err := stream.Err(); err == nil || !asRejected(err, &rejected) { + t.Fatalf("got %v, want *StreamRejectedError", stream.Err()) + } + if got := atomic.LoadInt32(&srv.subscribeCount); got != 1 { + t.Fatalf("subscribe attempts = %d, want exactly 1 (a rejection must not be retried)", got) + } +} + +// Close() during the reconnect backoff must stop the loop and leave no +// goroutine behind. +func TestStreamCloseDuringBackoffLeavesNoGoroutine(t *testing.T) { + before := runtime.NumGoroutine() + + srv := &scriptedCable{healthySessions: 1, postConfirmPing: true} + stream, httpSrv := newScriptedStream(t, srv, StreamOptions{ + AutoReconnect: true, + ReconnectDelay: 500 * time.Millisecond, + MaxReconnectDelay: 500 * time.Millisecond, + MaxReconnectAttempts: 10, + }) + defer httpSrv.Close() + + // Wait until the first session has dropped and the loop is in backoff. + waitFor(t, func() bool { return atomic.LoadInt32(&srv.connections) >= 2 }, 2*time.Second) + + if err := stream.Close(); err != nil { + t.Fatalf("Close during backoff returned %v", err) + } + + if _, ok := <-stream.Updates(); ok { + t.Fatal("updates channel should be closed after Close") + } + + // Give the runtime a moment to reap. + for deadline := time.Now().Add(2 * time.Second); time.Now().Before(deadline); { + if runtime.NumGoroutine() <= before+2 { + break + } + time.Sleep(5 * time.Millisecond) + } + if after := runtime.NumGoroutine(); after > before+2 { + t.Fatalf("goroutines: before=%d after=%d", before, after) + } +} + +func asRejected(err error, target **StreamRejectedError) bool { + r, ok := err.(*StreamRejectedError) + if ok { + *target = r + } + return ok +}