diff --git a/internal/corrosion/subscribe.go b/internal/corrosion/subscribe.go index 93b553a6..2e317cb7 100644 --- a/internal/corrosion/subscribe.go +++ b/internal/corrosion/subscribe.go @@ -287,8 +287,12 @@ func (c *APIClient) resubscribeWithBackoffFn(id string) func(context.Context, ui } return func(ctx context.Context, fromChange uint64) (*Subscription, error) { return backoff.RetryWithData(func() (*Subscription, error) { - slog.Debug("Retrying to resubscribe to Corrosion query.", "id", id, "from_change", fromChange) - return c.ResubscribeContext(ctx, id, fromChange) + sub, err := c.ResubscribeContext(ctx, id, fromChange) + if err != nil { + slog.Debug("Failed to resubscribe to Corrosion query. Retrying with backoff.", + "id", id, "from_change", fromChange, "err", err) + } + return sub, err }, c.newResubBackoff()) } } @@ -320,5 +324,24 @@ func (c *APIClient) ResubscribeContext(ctx context.Context, id string, fromChang return nil, fmt.Errorf("unexpected status code %d: %s", resp.StatusCode, respBody) } + // Since https://github.com/superfly/corrosion/pull/355, Corrosion treats a resubscription from change 0 like + // a fresh subscription: it replays the full query snapshot (a columns event, all rows, and an end-of-query event) + // before streaming changes. We don't expose rows in this case, so drain the snapshot here before consuming changes. + if fromChange == 0 { + rows, err := newRows(ctx, resp.Body, false) + if err != nil { + resp.Body.Close() + return nil, fmt.Errorf("parse resubscribe response: %w", err) + } + // Drain the replayed rows until the end-of-query event to reach the change stream. + for rows.Next() { + } + if err = rows.Err(); err != nil { + resp.Body.Close() + return nil, fmt.Errorf("drain resubscribe snapshot: %w", err) + } + return newSubscription(ctx, id, nil, rows.body, rows.decoder, c.resubscribeWithBackoffFn(id)), nil + } + return newSubscription(ctx, id, nil, resp.Body, nil, c.resubscribeWithBackoffFn(id)), nil } diff --git a/internal/corrosion/subscribe_test.go b/internal/corrosion/subscribe_test.go new file mode 100644 index 00000000..3d11c841 --- /dev/null +++ b/internal/corrosion/subscribe_test.go @@ -0,0 +1,158 @@ +package corrosion + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + "github.com/cenkalti/backoff/v4" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newTestClient builds an APIClient pointed at srv with a fast resubscribe backoff for tests. +func newTestClient(t *testing.T, srv *httptest.Server) *APIClient { + t.Helper() + + baseURL, err := url.Parse(srv.URL) + require.NoError(t, err) + + return &APIClient{ + baseURL: baseURL, + client: srv.Client(), + newResubBackoff: func() backoff.BackOff { + return backoff.NewExponentialBackOff( + backoff.WithInitialInterval(time.Millisecond), + backoff.WithMaxInterval(10*time.Millisecond), + backoff.WithMaxElapsedTime(5*time.Second), + ) + }, + } +} + +func flushString(t *testing.T, w http.ResponseWriter, s string) { + t.Helper() + _, err := w.Write([]byte(s)) + require.NoError(t, err) + w.(http.Flusher).Flush() +} + +// TestSubscription_ResubscribeFromZeroDrainsSnapshot verifies that when the connection drops before any +// change is seen (lastChangeID == 0) and Corrosion replays the full query snapshot on resubscription from +// change 0, the change handler drains the snapshot and delivers the subsequent change instead of looping on +// "expected change event, got: {Columns:...}". +// The change of behaviour in Corrosion: https://github.com/superfly/corrosion/pull/355 +func TestSubscription_ResubscribeFromZeroDrainsSnapshot(t *testing.T) { + t.Parallel() + + const snapshot = `{"columns":["id","info"]} +{"row":[1,["a","b"]]} +{"eoq":{"time":1e-7,"change_id":0}} +` + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodPost: + // Initial subscription: send the snapshot, then end the stream to force a resubscribe. + w.Header().Set("corro-query-id", "test-sub") + w.WriteHeader(http.StatusOK) + flushString(t, w, snapshot) + case http.MethodGet: + // Resubscription from change 0: Corrosion v1.0.0+ replays the full snapshot, then changes. + require.Equal(t, "0", r.URL.Query().Get("from")) + w.Header().Set("corro-query-id", "test-sub") + w.WriteHeader(http.StatusOK) + flushString(t, w, snapshot+`{"change":["insert",2,["c","d"],1]}`+"\n") + // Keep the connection open so the handler doesn't end the stream and trigger another resubscribe. + <-r.Context().Done() + default: + require.FailNow(t, fmt.Sprintf("unexpected request method: %s", r.Method)) + } + })) + defer srv.Close() + + client := newTestClient(t, srv) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sub, err := client.SubscribeContext(ctx, "SELECT id, info FROM machines", nil, false) + require.NoError(t, err) + + // Consume the initial rows so Changes can be called. + rows := sub.Rows() + for rows.Next() { + } + require.NoError(t, rows.Err()) + + changes, err := sub.Changes() + require.NoError(t, err) + + select { + case change := <-changes: + require.NotNil(t, change, "changes channel closed with error: %v", sub.Err()) + assert.Equal(t, ChangeTypeInsert, change.Type) + assert.Equal(t, uint64(2), change.RowID) + assert.Equal(t, uint64(1), change.ChangeID) + case <-time.After(5 * time.Second): + require.FailNow(t, "timed out waiting for a change after resubscribe", sub.Err()) + } +} + +// TestSubscription_ResubscribeFromNonZeroSkipsSnapshot verifies that a resubscription from a non-zero change +// streams changes directly, without a replayed snapshot, and the change handler delivers them. +func TestSubscription_ResubscribeFromNonZeroSkipsSnapshot(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodPost: + // Initial subscription with a non-zero last change id, then end the stream to force a resubscribe. + w.Header().Set("corro-query-id", "test-sub") + w.WriteHeader(http.StatusOK) + flushString(t, w, `{"columns":["id","info"]}`+"\n"+ + `{"row":[1,["a","b"]]}`+"\n"+ + `{"eoq":{"time":1e-7,"change_id":5}}`+"\n") + case http.MethodGet: + // Resubscription from change 5: only changes are streamed, no snapshot. + require.Equal(t, "5", r.URL.Query().Get("from")) + require.False(t, strings.Contains(r.URL.RawQuery, "skip_rows")) + w.WriteHeader(http.StatusOK) + flushString(t, w, `{"change":["update",1,["a","b2"],6]}`+"\n") + <-r.Context().Done() + default: + require.FailNow(t, fmt.Sprintf("unexpected request method: %s", r.Method)) + } + })) + defer srv.Close() + + client := newTestClient(t, srv) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sub, err := client.SubscribeContext(ctx, "SELECT id, info FROM machines", nil, false) + require.NoError(t, err) + + rows := sub.Rows() + for rows.Next() { + } + require.NoError(t, rows.Err()) + + changes, err := sub.Changes() + require.NoError(t, err) + + select { + case change := <-changes: + require.NotNil(t, change, "changes channel closed with error: %v", sub.Err()) + assert.Equal(t, ChangeTypeUpdate, change.Type) + assert.Equal(t, uint64(6), change.ChangeID) + case <-time.After(5 * time.Second): + require.FailNow(t, "timed out waiting for a change after resubscribe", sub.Err()) + } +}