From 4c4606655ad34889b58ad05bd525cfdabda40ca2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luk=C3=A1=C5=A1=20Dvo=C5=99=C3=A1k?= <dvorakluk@gmail.com>
Date: Sat, 28 Oct 2023 14:36:51 +0200
Subject: [PATCH] proxy: add retries with exponential backoff

---
 cmd/templ/generatecmd/proxy/proxy.go | 36 ++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/cmd/templ/generatecmd/proxy/proxy.go b/cmd/templ/generatecmd/proxy/proxy.go
index 01571566d..07998e286 100644
--- a/cmd/templ/generatecmd/proxy/proxy.go
+++ b/cmd/templ/generatecmd/proxy/proxy.go
@@ -4,12 +4,14 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"math"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
 	"os"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/a-h/templ/cmd/templ/generatecmd/sse"
 
@@ -31,6 +33,11 @@ type Handler struct {
 func New(port int, target *url.URL) *Handler {
 	p := httputil.NewSingleHostReverseProxy(target)
 	p.ErrorLog = log.New(os.Stderr, "Proxy to target error: ", 0)
+	p.Transport = &roundTripper{
+		maxRetries:      10,
+		initialDelay:    100 * time.Millisecond,
+		backoffExponent: 1.5,
+	}
 	p.ModifyResponse = func(r *http.Response) error {
 		if contentType := r.Header.Get("Content-Type"); !strings.HasPrefix(contentType, "text/html") {
 			return nil
@@ -74,3 +81,32 @@ func (p *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 func (p *Handler) SendSSE(eventType string, data string) {
 	p.sse.Send(eventType, data)
 }
+
+type roundTripper struct {
+	maxRetries      int
+	initialDelay    time.Duration
+	backoffExponent float64
+}
+
+func (rt *roundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+	var err error
+
+	for retries := 0; retries < rt.maxRetries; retries++ {
+		req := r.Clone(r.Context())
+
+		req.Body, err = r.GetBody()
+		if err != nil {
+			return nil, err
+		}
+
+		resp, err := http.DefaultTransport.RoundTrip(req)
+		if err != nil {
+			time.Sleep(rt.initialDelay * time.Duration(math.Pow(rt.backoffExponent, float64(retries))))
+			continue
+		}
+
+		return resp, nil
+	}
+
+	return nil, fmt.Errorf("max retries reached")
+}