Resolves #436 - Handle Connection Errors

elasticpath · Jan 7, 2024 · fc49f7c · fc49f7c
1 parent ef94eca
commit fc49f7c
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 11 deletions.
diff --git a/cmd/root.go b/cmd/root.go
@@ -124,6 +124,10 @@ func InitializeCmd() {
 	RootCmd.PersistentFlags().Uint16VarP(&rateLimit, "rate-limit", "", 10, "Request limit per second")
 	RootCmd.PersistentFlags().BoolVarP(&httpclient.Retry5xx, "retry-5xx", "", false, "Whether we should retry requests with HTTP 5xx response code")
 	RootCmd.PersistentFlags().BoolVarP(&httpclient.Retry429, "retry-429", "", false, "Whether we should retry requests with HTTP 429 response code")
+	RootCmd.PersistentFlags().BoolVarP(&httpclient.RetryConnectionErrors, "retry-connection-errors", "", false, "Whether we should retry requests with connection errors")
+	RootCmd.PersistentFlags().UintVarP(&httpclient.RetryDelay, "retry-delay", "", 500, "When retrying how long should we delay")
+	RootCmd.PersistentFlags().BoolVarP(&httpclient.RetryAllErrors, "retry-all-errors", "", false, "When enable retries on all errors (i.e., the same as --retry-5xx --retry-429 and --retry-connection-errors")
+
 	RootCmd.PersistentFlags().BoolVarP(&httpclient.DontLog2xxs, "silence-2xx", "", false, "Whether we should silence HTTP 2xx response code logging")
 
 	RootCmd.PersistentFlags().Float32VarP(&requestTimeout, "timeout", "", 60, "Request timeout in seconds (fractional values allowed)")

diff --git a/external/httpclient/httpclient.go b/external/httpclient/httpclient.go
@@ -97,6 +97,11 @@ func Initialize(rateLimit uint16, requestTimeout float32, statisticsFrequency in
 var Retry429 = false
 var Retry5xx = false
 
+var RetryConnectionErrors = false
+
+var RetryAllErrors = false
+var RetryDelay uint = 500
+
 var statsLock = &sync.Mutex{}
 
 var HttpClient = &http.Client{}
@@ -179,9 +184,16 @@ func doRequestInternal(ctx context.Context, method string, contentType string, p
 	reqURL.Path = path
 	reqURL.RawQuery = query
 
-	var bodyBuf bytes.Buffer
+	var bodyBuf []byte
 	if payload != nil {
-		payload = io.TeeReader(payload, &bodyBuf)
+		buf := new(bytes.Buffer)
+		_, err := buf.ReadFrom(payload)
+		if err != nil {
+			log.Warnf("Error reading payload, %s", err)
+		}
+		bodyBuf = buf.Bytes()
+
+		payload = bytes.NewReader(bodyBuf)
 	}
 
 	req, err := http.NewRequest(method, reqURL.String(), payload)
@@ -270,9 +282,25 @@ func doRequestInternal(ctx context.Context, method string, contentType string, p
 	statsLock.Unlock()
 
 	log.Tracef("Stats processing complete")
-
-	if err != nil {
-		return nil, err
+	requestError := err
+	if requestError != nil {
+
+		resp = &http.Response{
+			Status:           "CONNECTION_ERROR",
+			StatusCode:       0,
+			Proto:            "",
+			ProtoMajor:       0,
+			ProtoMinor:       0,
+			Header:           map[string][]string{},
+			Body:             io.NopCloser(strings.NewReader(fmt.Sprintf("%v", err))),
+			ContentLength:    0,
+			TransferEncoding: nil,
+			Close:            true,
+			Uncompressed:     false,
+			Trailer:          nil,
+			Request:          nil,
+			TLS:              nil,
+		}
 	}
 
 	var logf func(string, ...interface{})
@@ -322,7 +350,7 @@ func doRequestInternal(ctx context.Context, method string, contentType string, p
 
 	if displayLongFormRequestAndResponse {
 		if payload != nil {
-			body, _ := io.ReadAll(&bodyBuf)
+			body := bodyBuf
 			if len(body) > 0 {
 				logf("(%0.4d) %s %s%s", requestNumber, method, reqURL.String(), requestHeaders)
 				if contentType == "application/json" {
@@ -352,10 +380,24 @@ func doRequestInternal(ctx context.Context, method string, contentType string, p
 	log.Tracef("Starting log to disk")
 	profiles.LogRequestToDisk(method, path, dumpReq, dumpRes, resp.StatusCode)
 	log.Tracef("Done log to disk")
-	if resp.StatusCode == 429 && Retry429 {
-		return doRequestInternal(ctx, method, contentType, path, query, &bodyBuf)
-	} else if resp.StatusCode >= 500 && Retry5xx {
-		return doRequestInternal(ctx, method, contentType, path, query, &bodyBuf)
+	if resp.StatusCode == 429 && (Retry429 || RetryAllErrors) {
+		if RetryDelay > 0 {
+			log.Debugf("Retrying request in %d ms", RetryDelay)
+			time.Sleep(time.Duration(RetryDelay) * time.Millisecond)
+		}
+		return doRequestInternal(ctx, method, contentType, path, query, bytes.NewReader(bodyBuf))
+	} else if resp.StatusCode >= 500 && (Retry5xx || RetryAllErrors) {
+		if RetryDelay > 0 {
+			log.Debugf("Retrying request in %d ms", RetryDelay)
+			time.Sleep(time.Duration(RetryDelay) * time.Millisecond)
+		}
+		return doRequestInternal(ctx, method, contentType, path, query, bytes.NewReader(bodyBuf))
+	} else if requestError != nil && (RetryConnectionErrors || RetryAllErrors) {
+		if RetryDelay > 0 {
+			log.Debugf("Retrying request in %d ms", RetryDelay)
+			time.Sleep(time.Duration(RetryDelay) * time.Millisecond)
+		}
+		return doRequestInternal(ctx, method, contentType, path, query, bytes.NewReader(bodyBuf))
 	} else {
 		return resp, err
 	}

diff --git a/external/profiles/requestlogs.go b/external/profiles/requestlogs.go
@@ -86,7 +86,12 @@ func LogRequestToDisk(requestMethod string, requestPath string, requestBytes []b
 		responseBytes = regex1.ReplaceAll(responseBytes, []byte("client_secret=*****"))
 	}
 
-	return SaveRequest(fmt.Sprintf("%s %s ==> %d", requestMethod, requestPath, responseCode), requestBytes, responseBytes)
+	statusCode := fmt.Sprintf("%d", responseCode)
+
+	if responseCode == 0 {
+		statusCode = "ERROR"
+	}
+	return SaveRequest(fmt.Sprintf("%s %s ==> %s", requestMethod, requestPath, statusCode), requestBytes, responseBytes)
 }
 
 func SaveRequest(title string, requestBytes []byte, responseBytes []byte) error {