From e1c67e33c11b36aa84963ab155e211378a032faf Mon Sep 17 00:00:00 2001
From: hidu <613972+hidu@users.noreply.github.com>
Date: Fri, 12 Apr 2024 10:02:13 +0800
Subject: [PATCH] feat: upstream/remote support custom http client (#101)

---
 upstream/remote/remote.go | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/upstream/remote/remote.go b/upstream/remote/remote.go
index 33b81d7..287ff75 100644
--- a/upstream/remote/remote.go
+++ b/upstream/remote/remote.go
@@ -26,7 +26,7 @@ const cloudHostnameSuffix = "pyroscope.cloud"
 type Remote struct {
 	cfg    Config
 	jobs   chan *upstream.UploadJob
-	client *http.Client
+	client HTTPClient
 	logger Logger
 
 	done chan struct{}
@@ -35,6 +35,10 @@ type Remote struct {
 	flushWG sync.WaitGroup
 }
 
+type HTTPClient interface {
+	Do(req *http.Request) (*http.Response, error)
+}
+
 type Config struct {
 	AuthToken         string
 	BasicAuthUser     string // http basic auth user
@@ -45,6 +49,7 @@ type Config struct {
 	Address           string
 	Timeout           time.Duration
 	Logger            Logger
+	HTTPClient        HTTPClient // optional, custom client
 }
 
 type Logger interface {
@@ -74,6 +79,9 @@ func NewRemote(cfg Config) (*Remote, error) {
 		logger: cfg.Logger,
 		done:   make(chan struct{}),
 	}
+	if cfg.HTTPClient != nil {
+		r.client = cfg.HTTPClient
+	}
 
 	// parse the upstream address
 	u, err := url.Parse(cfg.Address)