chore: Improve retryablehttp client (#931)

* chore: Improve retryablehttp client * new request using NewRetryableRequestWithContext * Update internal/http/client.go Co-authored-by: Alex Plischke <[email protected]> --------- Co-authored-by: Alex Plischke <[email protected]>
saucelabs · Jul 30, 2024 · f203ddc · f203ddc
1 parent 390e015
commit f203ddc
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 82 deletions.
diff --git a/internal/http/client.go b/internal/http/client.go
@@ -1,6 +1,7 @@
 package http
 
 import (
+	"context"
 	"net/http"
 	"time"
 
@@ -17,7 +18,16 @@ func NewRetryableClient(timeout time.Duration) *retryablehttp.Client {
 		RetryWaitMin: 1 * time.Second,
 		RetryWaitMax: 30 * time.Second,
 		RetryMax:     3,
-		CheckRetry:   retryablehttp.DefaultRetryPolicy,
+		CheckRetry: func(ctx context.Context, resp *http.Response, err error) (bool, error) {
+			ok, e := retryablehttp.DefaultRetryPolicy(ctx, resp, err)
+			if e != nil {
+				return ok, e
+			}
+			if !ok && resp.StatusCode == http.StatusNotFound {
+				return true, nil
+			}
+			return ok, e
+		},
 		Backoff:      retryablehttp.DefaultBackoff,
 		ErrorHandler: retryablehttp.PassthroughErrorHandler,
 	}

diff --git a/internal/http/rdcservice.go b/internal/http/rdcservice.go
@@ -178,19 +178,14 @@ func (c *RDCService) StopJob(ctx context.Context, id string, realDevice bool) (j
 		return job.Job{}, errors.New("the RDC client does not support virtual device jobs")
 	}
 
-	req, err := NewRequestWithContext(ctx, http.MethodPut,
+	req, err := NewRetryableRequestWithContext(ctx, http.MethodPut,
 		fmt.Sprintf("%s/v1/rdc/jobs/%s/stop", c.URL, id), nil)
 	if err != nil {
 		return job.Job{}, err
 	}
 	req.SetBasicAuth(c.Username, c.AccessKey)
 
-	r, err := retryablehttp.FromRequest(req)
-	if err != nil {
-		return job.Job{}, err
-	}
-
-	resp, err := c.Client.Do(r)
+	resp, err := c.Client.Do(req)
 	if err != nil {
 		return job.Job{}, err
 	}
@@ -220,19 +215,14 @@ func (c *RDCService) ReadJob(ctx context.Context, id string, realDevice bool) (j
 		return job.Job{}, errors.New("the RDC client does not support virtual device jobs")
 	}
 
-	req, err := NewRequestWithContext(ctx, http.MethodGet,
+	req, err := NewRetryableRequestWithContext(ctx, http.MethodGet,
 		fmt.Sprintf("%s/v1/rdc/jobs/%s", c.URL, id), nil)
 	if err != nil {
 		return job.Job{}, err
 	}
 	req.SetBasicAuth(c.Username, c.AccessKey)
 
-	r, err := retryablehttp.FromRequest(req)
-	if err != nil {
-		return job.Job{}, err
-	}
-
-	resp, err := c.Client.Do(r)
+	resp, err := c.Client.Do(req)
 	if err != nil {
 		return job.Job{}, err
 	}
@@ -297,19 +287,14 @@ func (c *RDCService) GetJobAssetFileNames(ctx context.Context, jobID string, rea
 		return nil, errors.New("the RDC client does not support virtual device jobs")
 	}
 
-	req, err := NewRequestWithContext(ctx, http.MethodGet,
+	req, err := NewRetryableRequestWithContext(ctx, http.MethodGet,
 		fmt.Sprintf("%s/v1/rdc/jobs/%s", c.URL, jobID), nil)
 	if err != nil {
 		return []string{}, err
 	}
 	req.SetBasicAuth(c.Username, c.AccessKey)
 
-	r, err := retryablehttp.FromRequest(req)
-	if err != nil {
-		return []string{}, err
-	}
-
-	resp, err := c.Client.Do(r)
+	resp, err := c.Client.Do(req)
 	if err != nil {
 		return []string{}, err
 	}
@@ -368,7 +353,7 @@ func (c *RDCService) GetJobAssetFileContent(ctx context.Context, jobID, fileName
 		acceptHeader = "text/plain"
 	}
 
-	req, err := NewRequestWithContext(ctx, http.MethodGet,
+	req, err := NewRetryableRequestWithContext(ctx, http.MethodGet,
 		fmt.Sprintf("%s/v1/rdc/jobs/%s/%s", c.URL, jobID, URIFileName), nil)
 	if err != nil {
 		return nil, err
@@ -378,12 +363,7 @@ func (c *RDCService) GetJobAssetFileContent(ctx context.Context, jobID, fileName
 	if acceptHeader != "" {
 		req.Header.Set("Accept", acceptHeader)
 	}
-
-	rreq, err := retryablehttp.FromRequest(req)
-	if err != nil {
-		return nil, err
-	}
-	resp, err := c.Client.Do(rreq)
+	resp, err := c.Client.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -445,7 +425,7 @@ func (c *RDCService) downloadArtifact(targetDir, jobID, fileName string, realDev
 
 // GetDevices returns the list of available devices using a specific operating system.
 func (c *RDCService) GetDevices(ctx context.Context, OS string) ([]devices.Device, error) {
-	req, err := NewRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("%s/v1/rdc/devices/filtered", c.URL), nil)
+	req, err := NewRetryableRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("%s/v1/rdc/devices/filtered", c.URL), nil)
 	if err != nil {
 		return nil, err
 	}
@@ -455,12 +435,7 @@ func (c *RDCService) GetDevices(ctx context.Context, OS string) ([]devices.Devic
 	req.URL.RawQuery = q.Encode()
 	req.SetBasicAuth(c.Username, c.AccessKey)
 
-	r, err := retryablehttp.FromRequest(req)
-	if err != nil {
-		return nil, err
-	}
-
-	res, err := c.Client.Do(r)
+	res, err := c.Client.Do(req)
 	if err != nil {
 		return []devices.Device{}, err
 	}

diff --git a/internal/http/resto.go b/internal/http/resto.go
@@ -84,7 +84,7 @@ func (c *Resto) ReadJob(ctx context.Context, id string, realDevice bool) (job.Jo
 		return job.Job{}, errors.New("the VDC client does not support real device jobs")
 	}
 
-	req, err := NewRequestWithContext(ctx, http.MethodGet,
+	req, err := NewRetryableRequestWithContext(ctx, http.MethodGet,
 		fmt.Sprintf("%s/rest/v1.1/%s/jobs/%s", c.URL, c.Username, id), nil)
 	if err != nil {
 		return job.Job{}, err
@@ -93,11 +93,7 @@ func (c *Resto) ReadJob(ctx context.Context, id string, realDevice bool) (job.Jo
 	req.Header.Set("Content-Type", "application/json")
 	req.SetBasicAuth(c.Username, c.AccessKey)
 
-	rreq, err := retryablehttp.FromRequest(req)
-	if err != nil {
-		return job.Job{}, err
-	}
-	resp, err := c.Client.Do(rreq)
+	resp, err := c.Client.Do(req)
 	if err != nil {
 		return job.Job{}, err
 	}
@@ -163,19 +159,15 @@ func (c *Resto) GetJobAssetFileNames(ctx context.Context, jobID string, realDevi
 		return nil, errors.New("the VDC client does not support real device jobs")
 	}
 
-	req, err := NewRequestWithContext(ctx, http.MethodGet,
+	req, err := NewRetryableRequestWithContext(ctx, http.MethodGet,
 		fmt.Sprintf("%s/rest/v1/%s/jobs/%s/assets", c.URL, c.Username, jobID), nil)
 	if err != nil {
 		return nil, err
 	}
 
 	req.SetBasicAuth(c.Username, c.AccessKey)
 
-	rreq, err := retryablehttp.FromRequest(req)
-	if err != nil {
-		return nil, err
-	}
-	resp, err := c.Client.Do(rreq)
+	resp, err := c.Client.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -219,20 +211,15 @@ func (c *Resto) GetJobAssetFileContent(ctx context.Context, jobID, fileName stri
 		return nil, errors.New("the VDC client does not support real device jobs")
 	}
 
-	req, err := NewRequestWithContext(ctx, http.MethodGet,
+	req, err := NewRetryableRequestWithContext(ctx, http.MethodGet,
 		fmt.Sprintf("%s/rest/v1/%s/jobs/%s/assets/%s", c.URL, c.Username, jobID, fileName), nil)
 	if err != nil {
 		return nil, err
 	}
 
 	req.SetBasicAuth(c.Username, c.AccessKey)
 
-	rreq, err := retryablehttp.FromRequest(req)
-	if err != nil {
-		return nil, err
-	}
-
-	resp, err := c.Client.Do(rreq)
+	resp, err := c.Client.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -270,7 +257,7 @@ func (c *Resto) IsTunnelRunning(ctx context.Context, id, owner string, filter tu
 }
 
 func (c *Resto) isTunnelRunning(ctx context.Context, id, owner string, filter tunnels.Filter) error {
-	req, err := NewRequestWithContext(ctx, http.MethodGet,
+	req, err := NewRetryableRequestWithContext(ctx, http.MethodGet,
 		fmt.Sprintf("%s/rest/v1/%s/tunnels", c.URL, c.Username), nil)
 	if err != nil {
 		return err
@@ -286,12 +273,7 @@ func (c *Resto) isTunnelRunning(ctx context.Context, id, owner string, filter tu
 	}
 	req.URL.RawQuery = q.Encode()
 
-	r, err := retryablehttp.FromRequest(req)
-	if err != nil {
-		return err
-	}
-
-	res, err := c.Client.Do(r)
+	res, err := c.Client.Do(req)
 	if err != nil {
 		return err
 	}
@@ -334,7 +316,7 @@ func (c *Resto) StopJob(ctx context.Context, jobID string, realDevice bool) (job
 		return job.Job{}, errors.New("the VDC client does not support real device jobs")
 	}
 
-	req, err := NewRequestWithContext(ctx, http.MethodPut,
+	req, err := NewRetryableRequestWithContext(ctx, http.MethodPut,
 		fmt.Sprintf("%s/rest/v1/%s/jobs/%s/stop", c.URL, c.Username, jobID), nil)
 	if err != nil {
 		return job.Job{}, err
@@ -343,11 +325,7 @@ func (c *Resto) StopJob(ctx context.Context, jobID string, realDevice bool) (job
 	req.Header.Set("Content-Type", "application/json")
 	req.SetBasicAuth(c.Username, c.AccessKey)
 
-	rreq, err := retryablehttp.FromRequest(req)
-	if err != nil {
-		return job.Job{}, err
-	}
-	resp, err := c.Client.Do(rreq)
+	resp, err := c.Client.Do(req)
 	if err != nil {
 		return job.Job{}, err
 	}
@@ -409,18 +387,13 @@ func (c *Resto) downloadArtifact(targetDir, jobID, fileName string) error {
 
 // GetVirtualDevices returns the list of available virtual devices.
 func (c *Resto) GetVirtualDevices(ctx context.Context, kind string) ([]vmd.VirtualDevice, error) {
-	req, err := NewRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("%s/rest/v1.1/info/platforms/all", c.URL), nil)
+	req, err := NewRetryableRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("%s/rest/v1.1/info/platforms/all", c.URL), nil)
 	if err != nil {
 		return nil, err
 	}
 	req.SetBasicAuth(c.Username, c.AccessKey)
 
-	r, err := retryablehttp.FromRequest(req)
-	if err != nil {
-		return []vmd.VirtualDevice{}, err
-	}
-
-	res, err := c.Client.Do(r)
+	res, err := c.Client.Do(req)
 	if err != nil {
 		return []vmd.VirtualDevice{}, err
 	}
@@ -465,18 +438,13 @@ func (c *Resto) GetVirtualDevices(ctx context.Context, kind string) ([]vmd.Virtu
 }
 
 func (c *Resto) GetBuildID(ctx context.Context, jobID string, buildSource build.Source) (string, error) {
-	req, err := NewRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("%s/v2/builds/%s/jobs/%s/build/", c.URL, buildSource, jobID), nil)
+	req, err := NewRetryableRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("%s/v2/builds/%s/jobs/%s/build/", c.URL, buildSource, jobID), nil)
 	if err != nil {
 		return "", err
 	}
 	req.SetBasicAuth(c.Username, c.AccessKey)
 
-	r, err := retryablehttp.FromRequest(req)
-	if err != nil {
-		return "", err
-	}
-
-	resp, err := c.Client.Do(r)
+	resp, err := c.Client.Do(req)
 	if err != nil {
 		return "", err
 	}