diff --git a/cmd/root.go b/cmd/root.go index 4148672..406c87e 100644 --- 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 index e7b76b5..8401db2 100644 --- 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 index 977b312..7238f0f 100644 --- 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 {