From e123f7214e71982539c0bb2608c54f8c1cf21aaf Mon Sep 17 00:00:00 2001
From: Pavel Sviderski <me@psviderski.name>
Date: Tue, 1 Oct 2024 23:15:36 +1000
Subject: [PATCH] add exponential backoff for network errors in corrosion
 client

---
 go.mod                       |  4 +--
 internal/corrosion/client.go | 50 ++++++++++++++++++++++++++++++------
 2 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/go.mod b/go.mod
index dc2ab9ef..03ce2b73 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ toolchain go1.22.5
 
 require (
 	github.com/BurntSushi/toml v1.4.0
+	github.com/cenkalti/backoff/v4 v4.3.0
 	github.com/charmbracelet/huh v0.6.0
 	github.com/coreos/go-systemd v0.0.0-20191104093116-d3cd4ed1dbcf
 	github.com/dgraph-io/badger/v3 v3.2103.5
@@ -28,6 +29,7 @@ require (
 	go.uber.org/zap v1.27.0
 	go4.org/netipx v0.0.0-20231129151722-fdeea329fbba
 	golang.org/x/crypto v0.27.0
+	golang.org/x/net v0.29.0
 	golang.org/x/sync v0.8.0
 	golang.org/x/sys v0.25.0
 	golang.zx2c4.com/wireguard v0.0.0-20231211153847-12269c276173
@@ -48,7 +50,6 @@ require (
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/bgentry/speakeasy v0.2.0 // indirect
 	github.com/catppuccin/go v0.2.0 // indirect
-	github.com/cenkalti/backoff/v4 v4.3.0 // indirect
 	github.com/cespare/xxhash v1.1.0 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/charmbracelet/bubbles v0.20.0 // indirect
@@ -163,7 +164,6 @@ require (
 	go.uber.org/multierr v1.11.0 // indirect
 	golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56 // indirect
 	golang.org/x/mod v0.21.0 // indirect
-	golang.org/x/net v0.29.0 // indirect
 	golang.org/x/text v0.18.0 // indirect
 	golang.org/x/time v0.0.0-20220210224613-90d013bbcef8 // indirect
 	golang.org/x/tools v0.25.0 // indirect
diff --git a/internal/corrosion/client.go b/internal/corrosion/client.go
index f9b12cf8..c9a03174 100644
--- a/internal/corrosion/client.go
+++ b/internal/corrosion/client.go
@@ -7,8 +7,10 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/cenkalti/backoff/v4"
 	"golang.org/x/net/http2"
 	"io"
+	"log/slog"
 	"net"
 	"net/http"
 	"net/netip"
@@ -20,7 +22,7 @@ const (
 	// HTTP2ConnectTimeout is the maximum amount of time a client will wait for a connection to be established.
 	http2ConnectTimeout = 3 * time.Second
 	// HTTP2Timeout is the maximum amount of time a client will wait for a response.
-	http2Timeout = 30 * time.Second
+	http2Timeout = 20 * time.Second
 )
 
 // APIClient is a client for the Corrosion API.
@@ -38,19 +40,51 @@ func NewAPIClient(addr netip.AddrPort) (*APIClient, error) {
 		baseURL: baseURL,
 		client: &http.Client{
 			Timeout: http2Timeout,
-			Transport: &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)
+			Transport: &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)
+					},
 				},
+				Backoff: backoff.NewExponentialBackOff(
+					backoff.WithInitialInterval(100*time.Millisecond),
+					backoff.WithMaxInterval(1*time.Second),
+					backoff.WithMaxElapsedTime(10*time.Second),
+				),
 			},
 		},
 	}, nil
 }
 
+type RetryRoundTripper struct {
+	Base    http.RoundTripper
+	Backoff 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 {
+			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)
+				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.Backoff, req.Context())
+	return backoff.RetryWithData(roundTrip, boff)
+}
+
 type Statement struct {
 	Query  string `json:"query"`
 	Params []any  `json:"params"`