package corrosion import ( "context" "crypto/tls" "errors" "fmt" "log/slog" "net" "net/http" "net/netip" "net/url" "time" "github.com/cenkalti/backoff/v4" "golang.org/x/net/http2" ) const ( // http2ConnectTimeout is the maximum amount of time an HTTP2 client will wait for a connection to be established. http2ConnectTimeout = 3 * time.Second // http2MaxRetryTime is the maximum amount of time an HTTP2 client will retry a request. http2MaxRetryTime = 10 * time.Second // resubscribeMaxRetryTime is the maximum amount of time an API client will retry resubscribing to a query after // an error occurs. resubscribeMaxRetryTime = 60 * time.Second ) // APIClient is a client for the Corrosion API. type APIClient struct { baseURL *url.URL client *http.Client newResubBackoff func() backoff.BackOff } // NewAPIClient creates a new Corrosion API client. The bearerToken is sent in the Authorization header of every // request to authenticate against Corrosion API. // The client retries on network errors using an exponential backoff policy with a maximum interval of 1 second and // a maximum elapsed time of 10 seconds. // It automatically resubscribes to active subscriptions if an error occurs using an exponential backoff policy with a // maximum interval of 1 second and a maximum elapsed time of 60 seconds. // Use the WithHTTP2Client option to provide a custom HTTP client and the WithResubscribeBackoff option to change the // backoff policy for resubscribing to a query. func NewAPIClient(addr netip.AddrPort, bearerToken string, opts ...APIClientOption) (*APIClient, error) { baseURL, err := url.Parse(fmt.Sprintf("http://%s", addr)) if err != nil { return nil, fmt.Errorf("invalid URL: %w", err) } transport := &AuthRoundTripper{ Base: &RetryRoundTripper{ Base: &http2.Transport{ AllowHTTP: true, DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) { dialer := &net.Dialer{ Timeout: http2ConnectTimeout, } return dialer.DialContext(ctx, network, addr) }, }, NewBackoff: func() backoff.BackOff { return backoff.NewExponentialBackOff( backoff.WithInitialInterval(100*time.Millisecond), backoff.WithMaxInterval(1*time.Second), backoff.WithMaxElapsedTime(http2MaxRetryTime), ) }, }, Token: bearerToken, } c := &APIClient{ baseURL: baseURL, client: &http.Client{Transport: transport}, newResubBackoff: func() backoff.BackOff { return backoff.NewExponentialBackOff( backoff.WithInitialInterval(100*time.Millisecond), backoff.WithMaxInterval(1*time.Second), backoff.WithMaxElapsedTime(resubscribeMaxRetryTime), ) }, } for _, opt := range opts { opt(c) } return c, nil } type APIClientOption func(*APIClient) // WithHTTP2Client replaces the client's HTTP transport. The provided client bypasses the built-in bearer-token // injection, so the caller is responsible for setting the Authorization header. func WithHTTP2Client(client *http.Client) APIClientOption { return func(c *APIClient) { c.client = client } } // 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.newResubBackoff = newBackoff } } // AuthRoundTripper sets the Authorization header on every outgoing request. An empty Token leaves the header untouched. type AuthRoundTripper struct { Base http.RoundTripper Token string } func (rt *AuthRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { if rt.Token == "" { return rt.Base.RoundTrip(req) } // RoundTripper contract: must not mutate the caller's request. req = req.Clone(req.Context()) req.Header.Set("Authorization", "Bearer "+rt.Token) return rt.Base.RoundTrip(req) } type RetryRoundTripper struct { Base http.RoundTripper // NewBackoff creates a new backoff policy for each request. NewBackoff func() backoff.BackOff } func (rt *RetryRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { roundTrip := func() (*http.Response, error) { resp, err := rt.Base.RoundTrip(req) if err != nil { if _, ok := errors.AsType[*net.OpError](err); ok { // Not certain, but I expect operational errors should generally be retryable. slog.Debug("Retrying corrosion API request due to network error.", "error", err) return nil, err } // Don't retry on other errors. return nil, backoff.Permanent(err) } // Success, don't retry. return resp, err } 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() }