auto resubscribe a corrosion subscription if an error occurs

internal/corrosion/client.go

@@ -29,7 +29,7 @@ const (
 type APIClient struct {
 	baseURL         *url.URL
 	client          *http.Client
-	newResubsribeBackoff func() backoff.BackOff
+	newResubBackoff func() backoff.BackOff
 }
 
 // NewAPIClient creates a new Corrosion API client. The client retries on network errors using an exponential backoff
@@ -65,7 +65,7 @@ func NewAPIClient(addr netip.AddrPort, opts ...APIClientOption) (*APIClient, err
 				},
 			},
 		},
-		newResubsribeBackoff: func() backoff.BackOff {
+		newResubBackoff: func() backoff.BackOff {
 			return backoff.NewExponentialBackOff(
 				backoff.WithInitialInterval(100*time.Millisecond),
 				backoff.WithMaxInterval(1*time.Second),
@@ -88,10 +88,11 @@ func WithHTTP2Client(client *http.Client) APIClientOption {
 	}
 }
 
-// WithResubscribeBackoff sets the backoff policy for resubscribing to a query.
+// WithResubscribeBackoff sets the backoff policy for resubscribing to a query if an error occurs.
+// Pass nil to disable resubscribing.
 func WithResubscribeBackoff(newBackoff func() backoff.BackOff) APIClientOption {
 	return func(c *APIClient) {
-		c.newResubsribeBackoff = newBackoff
+		c.newResubBackoff = newBackoff
 	}
 }
 
@@ -108,7 +109,7 @@ func (rt *RetryRoundTripper) RoundTrip(req *http.Request) (*http.Response, error
 			var opErr *net.OpError
 			if errors.As(err, &opErr) {
 				// Not certain, but I expect operational errors should generally be retryable.
-				slog.Debug("Retrying corrosion API request due to network error", "error", err)
+				slog.Debug("Retrying corrosion API request due to network error.", "error", err)
 				return nil, err
 			}
 			// Don't retry on other errors.
@@ -120,3 +121,113 @@ func (rt *RetryRoundTripper) RoundTrip(req *http.Request) (*http.Response, error
 	boff := backoff.WithContext(rt.NewBackoff(), req.Context())
 	return backoff.RetryWithData(roundTrip, boff)
 }
+
+type RetrySubscription struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+
+	client       *APIClient
+	sub          *Subscription
+	changes      chan *ChangeEvent
+	lastChangeID uint64
+	err          error
+	backoff      backoff.BackOff
+}
+
+func NewRetrySubscription(ctx context.Context, client *APIClient, sub *Subscription, boff backoff.BackOff) *RetrySubscription {
+	ctx, cancel := context.WithCancel(ctx)
+	if boff == nil {
+		boff = backoff.NewExponentialBackOff(
+			backoff.WithInitialInterval(100*time.Millisecond),
+			backoff.WithMaxInterval(1*time.Second),
+			backoff.WithMaxElapsedTime(60*time.Second),
+		)
+	}
+	return &RetrySubscription{
+		ctx:     ctx,
+		cancel:  cancel,
+		client:  client,
+		sub:     sub,
+		backoff: backoff.WithContext(boff, ctx),
+	}
+}
+
+func (rs *RetrySubscription) ID() string {
+	return rs.sub.ID()
+}
+
+func (rs *RetrySubscription) Changes() (<-chan *ChangeEvent, error) {
+	if rs.changes != nil {
+		return rs.changes, nil
+	}
+
+	// Ensure the rows are consumed if they have been requested.
+	changes, err := rs.sub.Changes()
+	if err != nil {
+		return nil, err
+	}
+	rs.lastChangeID = rs.sub.lastChangeID
+
+	go func() {
+		defer close(rs.changes)
+
+		var change *ChangeEvent
+		for {
+			select {
+			case change = <-changes:
+			case <-rs.ctx.Done():
+				return
+			}
+
+			if change != nil {
+				select {
+				case rs.changes <- change:
+				case <-rs.ctx.Done():
+					return
+				}
+				rs.lastChangeID = change.ChangeID
+				continue
+			}
+
+			// The underlying subscription has been closed due to an error or context cancellation.
+			// Return if the context is done or try to resubscribe otherwise.
+			if rs.ctx.Err() != nil {
+				return
+			}
+
+			if err = rs.resubscribe(); err != nil {
+				// resubscribe returns a permanent error after unsuccessful retries.
+				rs.err = fmt.Errorf("resubscribe to query: %w", err)
+				return
+			}
+			changes, err = rs.sub.Changes()
+			if err != nil {
+				// Unexpected error but report it anyway.
+				rs.err = err
+				return
+			}
+		}
+	}()
+
+	return rs.changes, nil
+}
+
+func (rs *RetrySubscription) resubscribe() error {
+	return backoff.Retry(func() error {
+		slog.Debug("Resubscribing to Corrosion query.", "id", rs.ID(), "from_change", rs.lastChangeID)
+		sub, err := rs.client.ResubscribeContext(rs.ctx, rs.ID(), rs.lastChangeID)
+		if err != nil {
+			return err
+		}
+		rs.sub = sub
+		return nil
+	}, rs.backoff)
+}
+
+func (rs *RetrySubscription) Err() error {
+	return rs.err
+}
+
+func (rs *RetrySubscription) Close() {
+	rs.cancel()
+}

internal/corrosion/subscribe.go

@@ -6,71 +6,13 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/cenkalti/backoff/v4"
 	"io"
+	"log/slog"
 	"net/http"
+	"strconv"
 )
 
-// SubscribeContext creates a subscription to receive updates for a desired SQL query. If skipRows is false,
-// Subscription.Rows must be consumed before Subscription.Changes can be called. If skipRows is true, Subscription.Rows
-// will be nil.
-func (c *APIClient) SubscribeContext(
-	ctx context.Context, query string, args []any, skipRows bool,
-) (*Subscription, error) {
-	statement := Statement{
-		Query:  query,
-		Params: args,
-	}
-	body, err := json.Marshal(statement)
-	if err != nil {
-		return nil, fmt.Errorf("marshal query: %w", err)
-	}
-
-	subURL := c.baseURL.JoinPath("/v1/subscriptions")
-	if skipRows {
-		q := subURL.Query()
-		q.Set("skip_rows", "true")
-		subURL.RawQuery = q.Encode()
-	}
-
-	req, err := http.NewRequestWithContext(ctx, "POST", subURL.String(), bytes.NewReader(body))
-	if err != nil {
-		return nil, fmt.Errorf("create request: %w", err)
-	}
-	req.Header.Set("Content-Type", "application/json")
-	req.Header.Set("Accept", "application/json")
-
-	resp, err := c.client.Do(req)
-	if err != nil {
-		return nil, fmt.Errorf("send request: %w", err)
-	}
-
-	if resp.StatusCode != http.StatusOK {
-		respBody, err := io.ReadAll(resp.Body)
-		resp.Body.Close()
-		if err != nil {
-			return nil, fmt.Errorf("read response body: %w", err)
-		}
-		return nil, fmt.Errorf("unexpected status code %d: %s", resp.StatusCode, respBody)
-	}
-
-	id := resp.Header.Get("corro-query-id")
-	if id == "" {
-		resp.Body.Close()
-		return nil, errors.New("missing corro-query-id header in response")
-	}
-
-	if skipRows {
-		return newSubscription(ctx, id, nil, resp.Body, nil), nil
-	}
-
-	rows, err := newRows(ctx, resp.Body, false)
-	if err != nil {
-		resp.Body.Close()
-		return nil, fmt.Errorf("parse query response: %w", err)
-	}
-	return newSubscription(ctx, id, rows, rows.body, rows.decoder), nil
-}
-
 type ChangeType string
 
 var (
@@ -138,25 +80,32 @@ type Subscription struct {
 	rows         *Rows
 	body         io.ReadCloser
 	decoder      *json.Decoder
+	resubscribe  func(ctx context.Context, fromChange uint64) (*Subscription, error)
 	changes      chan *ChangeEvent
 	lastChangeID uint64
 	err          error
 }
 
 func newSubscription(
-	ctx context.Context, id string, rows *Rows, body io.ReadCloser, decoder *json.Decoder,
+	ctx context.Context,
+	id string,
+	rows *Rows,
+	body io.ReadCloser,
+	decoder *json.Decoder,
+	resubscribe func(ctx context.Context, fromChange uint64) (*Subscription, error),
 ) *Subscription {
 	ctx, cancel := context.WithCancel(ctx)
 	if decoder == nil {
 		decoder = json.NewDecoder(body)
 	}
 	return &Subscription{
-		id:      id,
-		rows:    rows,
 		ctx:         ctx,
 		cancel:      cancel,
+		id:          id,
+		rows:        rows,
 		body:        body,
 		decoder:     decoder,
+		resubscribe: resubscribe,
 	}
 }
 
@@ -193,8 +142,12 @@ func (s *Subscription) Changes() (<-chan *ChangeEvent, error) {
 		<-s.ctx.Done()
 		s.body.Close()
 	}()
+	go s.handleChangeEvents()
 
-	go func() {
+	return s.changes, nil
+}
+
+func (s *Subscription) handleChangeEvents() {
 	defer s.cancel()
 	defer close(s.changes)
 
@@ -206,39 +159,53 @@ func (s *Subscription) Changes() (<-chan *ChangeEvent, error) {
 		}
 
 		var e QueryEvent
-			if err := s.decoder.Decode(&e); err != nil {
+		var err error
-				// Do not report an error that occurred due to context cancellation.
+		if err = s.decoder.Decode(&e); err != nil {
-				if s.ctx.Err() == nil {
+			// Do not report an error that occurred due to context cancellation, just return.
-					s.err = fmt.Errorf("decode query event: %w", err)
+			if s.ctx.Err() != nil {
-				}
-				return
-			}
-
-			if e.Error != nil {
-				s.err = fmt.Errorf("query error: %s", *e.Error)
-				return
-			}
-			if e.Change == nil {
-				s.err = fmt.Errorf("expected change event, got: %+v", e)
 				return
 			}
+			err = fmt.Errorf("decode query event: %w", err)
+		} else if e.Error != nil {
+			err = fmt.Errorf("query error: %s", *e.Error)
+		} else if e.Change == nil {
+			err = fmt.Errorf("expected change event, got: %+v", e)
+		} else if s.lastChangeID != 0 && e.Change.ChangeID != s.lastChangeID+1 {
 			// If skipRows is true, the last change ID is unknown.
-			if s.lastChangeID != 0 && e.Change.ChangeID != s.lastChangeID+1 {
+			err = fmt.Errorf("missed a change: expected change ID %d, got %d",
-				s.err = fmt.Errorf("missed a change: expected change ID %d, got %d",
 				s.lastChangeID+1, e.Change.ChangeID)
-				return
 		}
 
+		if err == nil {
 			s.lastChangeID = e.Change.ChangeID
 			select {
 			case s.changes <- e.Change:
 			case <-s.ctx.Done():
 				return
 			}
+		} else {
+			// Report the error if resubscribing is disabled.
+			if s.resubscribe == nil {
+				s.err = err
+				return
 			}
-	}()
 
-	return s.changes, nil
+			slog.Info("Resubscribing to Corrosion query due to an error.",
+				"error", err, "id", s.id, "from_change", s.lastChangeID)
+			sub, sErr := s.resubscribe(s.ctx, s.lastChangeID)
+			if sErr != nil {
+				// resubscribe returns a permanent error after unsuccessful retries.
+				s.err = fmt.Errorf("resubscribe to query with backoff: %w", sErr)
+				return
+			}
+			// Reset the subscription to the new one.
+			s.rows = nil
+			s.body = sub.body
+			s.decoder = sub.decoder
+			// Do not close the sub to not close the body.
+			sub.cancel()
+		}
+	}
 }
 
 // Err returns the error, if any, that was encountered during fetching changes.
@@ -251,3 +218,106 @@ func (s *Subscription) Close() error {
 	s.cancel()
 	return s.body.Close()
 }
+
+// SubscribeContext creates a subscription to receive updates for a desired SQL query. If skipRows is false,
+// Subscription.Rows must be consumed before Subscription.Changes can be called. If skipRows is true, Subscription.Rows
+// will return nil.
+func (c *APIClient) SubscribeContext(
+	ctx context.Context, query string, args []any, skipRows bool,
+) (*Subscription, error) {
+	statement := Statement{
+		Query:  query,
+		Params: args,
+	}
+	body, err := json.Marshal(statement)
+	if err != nil {
+		return nil, fmt.Errorf("marshal query: %w", err)
+	}
+
+	subURL := c.baseURL.JoinPath("/v1/subscriptions")
+	if skipRows {
+		q := subURL.Query()
+		q.Set("skip_rows", "true")
+		subURL.RawQuery = q.Encode()
+	}
+
+	req, err := http.NewRequestWithContext(ctx, "POST", subURL.String(), bytes.NewReader(body))
+	if err != nil {
+		return nil, fmt.Errorf("create request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Accept", "application/json")
+
+	resp, err := c.client.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("send request: %w", err)
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		respBody, err := io.ReadAll(resp.Body)
+		resp.Body.Close()
+		if err != nil {
+			return nil, fmt.Errorf("read response body: %w", err)
+		}
+		return nil, fmt.Errorf("unexpected status code %d: %s", resp.StatusCode, respBody)
+	}
+
+	id := resp.Header.Get("corro-query-id")
+	if id == "" {
+		resp.Body.Close()
+		return nil, errors.New("missing corro-query-id header in response")
+	}
+
+	if skipRows {
+		return newSubscription(ctx, id, nil, resp.Body, nil, c.resubscribeWithBackoffFn(id)), nil
+	}
+
+	rows, err := newRows(ctx, resp.Body, false)
+	if err != nil {
+		resp.Body.Close()
+		return nil, fmt.Errorf("parse query response: %w", err)
+	}
+	return newSubscription(ctx, id, rows, rows.body, rows.decoder, c.resubscribeWithBackoffFn(id)), nil
+}
+
+func (c *APIClient) resubscribeWithBackoffFn(id string) func(context.Context, uint64) (*Subscription, error) {
+	if c.newResubBackoff == nil {
+		return nil
+	}
+	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)
+		}, c.newResubBackoff())
+	}
+}
+
+func (c *APIClient) ResubscribeContext(ctx context.Context, id string, fromChange uint64) (*Subscription, error) {
+	subURL := c.baseURL.JoinPath("/v1/subscriptions", id)
+	q := subURL.Query()
+	q.Set("from", strconv.FormatUint(fromChange, 10))
+	subURL.RawQuery = q.Encode()
+
+	req, err := http.NewRequestWithContext(ctx, "GET", subURL.String(), nil)
+	if err != nil {
+		return nil, fmt.Errorf("create request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Accept", "application/json")
+
+	resp, err := c.client.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("send request: %w", err)
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		respBody, err := io.ReadAll(resp.Body)
+		resp.Body.Close()
+		if err != nil {
+			return nil, fmt.Errorf("read response body: %w", err)
+		}
+		return nil, fmt.Errorf("unexpected status code %d: %s", resp.StatusCode, respBody)
+	}
+
+	return newSubscription(ctx, id, nil, resp.Body, nil, c.resubscribeWithBackoffFn(id)), nil
+}