Skip to content

Commit 99e1ae0

Browse files
committed
fix(corrosion): handle resubscription with from_change=0 that streams rows in new corrosion version
1 parent 5f07bc5 commit 99e1ae0

2 files changed

Lines changed: 183 additions & 2 deletions

File tree

‎internal/corrosion/subscribe.go‎

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -287,8 +287,12 @@ func (c *APIClient) resubscribeWithBackoffFn(id string) func(context.Context, ui
287287
}
288288
return func(ctx context.Context, fromChange uint64) (*Subscription, error) {
289289
return backoff.RetryWithData(func() (*Subscription, error) {
290-
slog.Debug("Retrying to resubscribe to Corrosion query.", "id", id, "from_change", fromChange)
291-
return c.ResubscribeContext(ctx, id, fromChange)
290+
sub, err := c.ResubscribeContext(ctx, id, fromChange)
291+
if err != nil {
292+
slog.Debug("Failed to resubscribe to Corrosion query. Retrying with backoff.",
293+
"id", id, "from_change", fromChange, "err", err)
294+
}
295+
return sub, err
292296
}, c.newResubBackoff())
293297
}
294298
}
@@ -320,5 +324,24 @@ func (c *APIClient) ResubscribeContext(ctx context.Context, id string, fromChang
320324
return nil, fmt.Errorf("unexpected status code %d: %s", resp.StatusCode, respBody)
321325
}
322326

327+
// Since https://github.com/superfly/corrosion/pull/355, Corrosion treats a resubscription from change 0 like
328+
// a fresh subscription: it replays the full query snapshot (a columns event, all rows, and an end-of-query event)
329+
// before streaming changes. We don't expose rows in this case, so drain the snapshot here before consuming changes.
330+
if fromChange == 0 {
331+
rows, err := newRows(ctx, resp.Body, false)
332+
if err != nil {
333+
resp.Body.Close()
334+
return nil, fmt.Errorf("parse resubscribe response: %w", err)
335+
}
336+
// Drain the replayed rows until the end-of-query event to reach the change stream.
337+
for rows.Next() {
338+
}
339+
if err = rows.Err(); err != nil {
340+
resp.Body.Close()
341+
return nil, fmt.Errorf("drain resubscribe snapshot: %w", err)
342+
}
343+
return newSubscription(ctx, id, nil, rows.body, rows.decoder, c.resubscribeWithBackoffFn(id)), nil
344+
}
345+
323346
return newSubscription(ctx, id, nil, resp.Body, nil, c.resubscribeWithBackoffFn(id)), nil
324347
}
Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
1+
package corrosion
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"net/http"
7+
"net/http/httptest"
8+
"net/url"
9+
"strings"
10+
"testing"
11+
"time"
12+
13+
"github.com/cenkalti/backoff/v4"
14+
"github.com/stretchr/testify/assert"
15+
"github.com/stretchr/testify/require"
16+
)
17+
18+
// newTestClient builds an APIClient pointed at srv with a fast resubscribe backoff for tests.
19+
func newTestClient(t *testing.T, srv *httptest.Server) *APIClient {
20+
t.Helper()
21+
22+
baseURL, err := url.Parse(srv.URL)
23+
require.NoError(t, err)
24+
25+
return &APIClient{
26+
baseURL: baseURL,
27+
client: srv.Client(),
28+
newResubBackoff: func() backoff.BackOff {
29+
return backoff.NewExponentialBackOff(
30+
backoff.WithInitialInterval(time.Millisecond),
31+
backoff.WithMaxInterval(10*time.Millisecond),
32+
backoff.WithMaxElapsedTime(5*time.Second),
33+
)
34+
},
35+
}
36+
}
37+
38+
func flushString(t *testing.T, w http.ResponseWriter, s string) {
39+
t.Helper()
40+
_, err := w.Write([]byte(s))
41+
require.NoError(t, err)
42+
w.(http.Flusher).Flush()
43+
}
44+
45+
// TestSubscription_ResubscribeFromZeroDrainsSnapshot verifies that when the connection drops before any
46+
// change is seen (lastChangeID == 0) and Corrosion replays the full query snapshot on resubscription from
47+
// change 0, the change handler drains the snapshot and delivers the subsequent change instead of looping on
48+
// "expected change event, got: {Columns:...}".
49+
// The change of behaviour in Corrosion: https://github.com/superfly/corrosion/pull/355
50+
func TestSubscription_ResubscribeFromZeroDrainsSnapshot(t *testing.T) {
51+
t.Parallel()
52+
53+
const snapshot = `{"columns":["id","info"]}
54+
{"row":[1,["a","b"]]}
55+
{"eoq":{"time":1e-7,"change_id":0}}
56+
`
57+
58+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
59+
switch r.Method {
60+
case http.MethodPost:
61+
// Initial subscription: send the snapshot, then end the stream to force a resubscribe.
62+
w.Header().Set("corro-query-id", "test-sub")
63+
w.WriteHeader(http.StatusOK)
64+
flushString(t, w, snapshot)
65+
case http.MethodGet:
66+
// Resubscription from change 0: Corrosion v1.0.0+ replays the full snapshot, then changes.
67+
require.Equal(t, "0", r.URL.Query().Get("from"))
68+
w.Header().Set("corro-query-id", "test-sub")
69+
w.WriteHeader(http.StatusOK)
70+
flushString(t, w, snapshot+`{"change":["insert",2,["c","d"],1]}`+"\n")
71+
// Keep the connection open so the handler doesn't end the stream and trigger another resubscribe.
72+
<-r.Context().Done()
73+
default:
74+
require.FailNow(t, fmt.Sprintf("unexpected request method: %s", r.Method))
75+
}
76+
}))
77+
defer srv.Close()
78+
79+
client := newTestClient(t, srv)
80+
81+
ctx, cancel := context.WithCancel(context.Background())
82+
defer cancel()
83+
84+
sub, err := client.SubscribeContext(ctx, "SELECT id, info FROM machines", nil, false)
85+
require.NoError(t, err)
86+
87+
// Consume the initial rows so Changes can be called.
88+
rows := sub.Rows()
89+
for rows.Next() {
90+
}
91+
require.NoError(t, rows.Err())
92+
93+
changes, err := sub.Changes()
94+
require.NoError(t, err)
95+
96+
select {
97+
case change := <-changes:
98+
require.NotNil(t, change, "changes channel closed with error: %v", sub.Err())
99+
assert.Equal(t, ChangeTypeInsert, change.Type)
100+
assert.Equal(t, uint64(2), change.RowID)
101+
assert.Equal(t, uint64(1), change.ChangeID)
102+
case <-time.After(5 * time.Second):
103+
require.FailNow(t, "timed out waiting for a change after resubscribe", sub.Err())
104+
}
105+
}
106+
107+
// TestSubscription_ResubscribeFromNonZeroSkipsSnapshot verifies that a resubscription from a non-zero change
108+
// streams changes directly, without a replayed snapshot, and the change handler delivers them.
109+
func TestSubscription_ResubscribeFromNonZeroSkipsSnapshot(t *testing.T) {
110+
t.Parallel()
111+
112+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
113+
switch r.Method {
114+
case http.MethodPost:
115+
// Initial subscription with a non-zero last change id, then end the stream to force a resubscribe.
116+
w.Header().Set("corro-query-id", "test-sub")
117+
w.WriteHeader(http.StatusOK)
118+
flushString(t, w, `{"columns":["id","info"]}`+"\n"+
119+
`{"row":[1,["a","b"]]}`+"\n"+
120+
`{"eoq":{"time":1e-7,"change_id":5}}`+"\n")
121+
case http.MethodGet:
122+
// Resubscription from change 5: only changes are streamed, no snapshot.
123+
require.Equal(t, "5", r.URL.Query().Get("from"))
124+
require.False(t, strings.Contains(r.URL.RawQuery, "skip_rows"))
125+
w.WriteHeader(http.StatusOK)
126+
flushString(t, w, `{"change":["update",1,["a","b2"],6]}`+"\n")
127+
<-r.Context().Done()
128+
default:
129+
require.FailNow(t, fmt.Sprintf("unexpected request method: %s", r.Method))
130+
}
131+
}))
132+
defer srv.Close()
133+
134+
client := newTestClient(t, srv)
135+
136+
ctx, cancel := context.WithCancel(context.Background())
137+
defer cancel()
138+
139+
sub, err := client.SubscribeContext(ctx, "SELECT id, info FROM machines", nil, false)
140+
require.NoError(t, err)
141+
142+
rows := sub.Rows()
143+
for rows.Next() {
144+
}
145+
require.NoError(t, rows.Err())
146+
147+
changes, err := sub.Changes()
148+
require.NoError(t, err)
149+
150+
select {
151+
case change := <-changes:
152+
require.NotNil(t, change, "changes channel closed with error: %v", sub.Err())
153+
assert.Equal(t, ChangeTypeUpdate, change.Type)
154+
assert.Equal(t, uint64(6), change.ChangeID)
155+
case <-time.After(5 * time.Second):
156+
require.FailNow(t, "timed out waiting for a change after resubscribe", sub.Err())
157+
}
158+
}

0 commit comments

Comments
 (0)