diff --git a/internal/corrosion/client.go b/internal/corrosion/client.go index bee5d2d1..fcbf8b78 100644 --- a/internal/corrosion/client.go +++ b/internal/corrosion/client.go @@ -27,9 +27,9 @@ const ( // APIClient is a client for the Corrosion API. type APIClient struct { - baseURL *url.URL - client *http.Client - newResubsribeBackoff func() backoff.BackOff + baseURL *url.URL + client *http.Client + 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() +} diff --git a/internal/corrosion/subscribe.go b/internal/corrosion/subscribe.go index ae9ec45e..820127cf 100644 --- a/internal/corrosion/subscribe.go +++ b/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, - body: body, - decoder: decoder, + ctx: ctx, + cancel: cancel, + id: id, + rows: rows, + body: body, + decoder: decoder, + resubscribe: resubscribe, } } @@ -193,52 +142,70 @@ func (s *Subscription) Changes() (<-chan *ChangeEvent, error) { <-s.ctx.Done() s.body.Close() }() + go s.handleChangeEvents() - go func() { - defer s.cancel() - defer close(s.changes) + return s.changes, nil +} - for { - select { - case <-s.ctx.Done(): - return - default: - } +func (s *Subscription) handleChangeEvents() { + defer s.cancel() + defer close(s.changes) - var e QueryEvent - if err := s.decoder.Decode(&e); err != nil { - // Do not report an error that occurred due to context cancellation. - if s.ctx.Err() == nil { - s.err = fmt.Errorf("decode query event: %w", err) - } - return - } + for { + select { + case <-s.ctx.Done(): + return + default: + } - 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) + var e QueryEvent + var err error + if err = s.decoder.Decode(&e); err != nil { + // Do not report an error that occurred due to context cancellation, just return. + if s.ctx.Err() != nil { 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 { - s.err = fmt.Errorf("missed a change: expected change ID %d, got %d", - s.lastChangeID+1, e.Change.ChangeID) - return - } + err = fmt.Errorf("missed a change: expected change ID %d, got %d", + s.lastChangeID+1, e.Change.ChangeID) + } + 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 +}