From 6dcdb8bf59973991657f2d2daee4545e4fd78553 Mon Sep 17 00:00:00 2001
From: rare-magma <rare-magma@posteo.eu>
Date: Sun, 19 Jan 2025 13:34:00 +0100
Subject: [PATCH] feat: setup backoff retry strategy

Signed-off-by: rare-magma <rare-magma@posteo.eu>
---
 main.go | 63 ++++++++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 54 insertions(+), 9 deletions(-)

diff --git a/main.go b/main.go
index e9012a9..28ae7af 100644
--- a/main.go
+++ b/main.go
@@ -8,7 +8,7 @@ import (
 	"fmt"
 	"io"
 	"log"
-	"net"
+	"math"
 	"net/http"
 	"os"
 	"time"
@@ -39,6 +39,52 @@ type Config struct {
 	Org              string `json:"Org"`
 }
 
+type retryableTransport struct {
+	transport             http.RoundTripper
+	TLSHandshakeTimeout   time.Duration
+	ResponseHeaderTimeout time.Duration
+}
+
+const retryCount = 3
+
+func shouldRetry(err error, resp *http.Response) bool {
+	if err != nil {
+		return true
+	}
+	switch resp.StatusCode {
+	case http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout:
+		return true
+	default:
+		return false
+	}
+}
+
+func (t *retryableTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	var bodyBytes []byte
+	if req.Body != nil {
+		bodyBytes, _ = io.ReadAll(req.Body)
+		req.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
+	}
+	resp, err := t.transport.RoundTrip(req)
+	retries := 0
+	for shouldRetry(err, resp) && retries < retryCount {
+		backoff := time.Duration(math.Pow(2, float64(retries))) * time.Second
+		time.Sleep(backoff)
+		if resp.Body != nil {
+			io.Copy(io.Discard, resp.Body)
+			resp.Body.Close()
+		}
+		if req.Body != nil {
+			req.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
+		}
+		log.Printf("Previous request failed with %s", resp.Status)
+		log.Printf("Retry %d of request to: %s", retries+1, req.URL)
+		resp, err = t.transport.RoundTrip(req)
+		retries++
+	}
+	return resp, err
+}
+
 func main() {
 	confFilePath := "pvpc_exporter.json"
 	confData, err := os.Open(confFilePath)
@@ -68,15 +114,14 @@ func main() {
 	flag.IntVar(&days, "days", 0, "Number of days in the past to fetch")
 	flag.Parse()
 
+	transport := &retryableTransport{
+		transport:             &http.Transport{},
+		TLSHandshakeTimeout:   30 * time.Second,
+		ResponseHeaderTimeout: 30 * time.Second,
+	}
 	client := &http.Client{
-		Timeout: 30 * time.Second,
-		Transport: &http.Transport{
-			DialContext: (&net.Dialer{
-				Timeout: 30 * time.Second,
-			}).DialContext,
-			TLSHandshakeTimeout:   30 * time.Second,
-			ResponseHeaderTimeout: 30 * time.Second,
-		},
+		Timeout:   30 * time.Second,
+		Transport: transport,
 	}
 
 	date := time.Now().AddDate(0, 0, -days).Format("2006-01-02")