uncloud/internal/corrosion/subscribe_test.go

package corrosion

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync/atomic"
	"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())
	}
}

// TestSubscription_ResubscribeNotFoundFailsFast verifies that when Corrosion returns 404 to a resubscription
// (the subscription no longer exists, e.g. after a restart that dropped it), the change handler stops retrying
// immediately and closes the changes channel with ErrSubscriptionNotFound, instead of retrying for the full
// backoff window.
func TestSubscription_ResubscribeNotFoundFailsFast(t *testing.T) {
	t.Parallel()

	var getCount atomic.Int32

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case http.MethodPost:
			// Initial subscription, 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: Corrosion no longer knows this subscription.
			getCount.Add(1)
			http.Error(w, "", http.StatusNotFound)
		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:
		// The channel must close (nil change) rather than deliver anything.
		require.Nil(t, change, "expected the changes channel to close on a gone subscription")
	case <-time.After(5 * time.Second):
		require.FailNow(t, "timed out waiting for the changes channel to close")
	}

	require.ErrorIs(t, sub.Err(), ErrSubscriptionNotFound)
	// The 404 must not be retried: exactly one GET resubscription request was made.
	assert.Equal(t, int32(1), getCount.Load())
}