feat: retries & timeouts

Author: danielgtaylor

README.md

@@ -33,6 +33,7 @@ Features include:
   - MessagePack (<https://msgpack.org/>)
   - Amazon Ion (<http://amzn.github.io/ion-docs/>)
   - Gzip ([RFC 1952](https://tools.ietf.org/html/rfc1952)), Deflate ([RFC 1951](https://datatracker.ietf.org/doc/html/rfc1951)), and Brotli ([RFC 7932](https://tools.ietf.org/html/rfc7932)) content encoding
+- Automatic retries with support for [`Retry-After`](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After) and `X-Retry-In` headers when APIs are rate-limited.
 - Standardized [hypermedia](https://smartbear.com/learn/api-design/what-is-hypermedia/) parsing into queryable/followable response links:
   - HTTP Link relation headers ([RFC 5988](https://tools.ietf.org/html/rfc5988#section-6.2.2))
   - [HAL](http://stateless.co/hal_specification.html)

cli/cli.go

@@ -551,6 +551,8 @@ Not after (expires): %s (%s)
 	AddGlobalFlag("rsh-client-key", "", "Path to a PEM encoded private key", "", false)
 	AddGlobalFlag("rsh-ca-cert", "", "Path to a PEM encoded CA cert", "", false)
 	AddGlobalFlag("rsh-ignore-status-code", "", "Do not set exit code from HTTP status code", false, false)
+	AddGlobalFlag("rsh-retry", "", "Number of times to retry on certain failures", 2, false)
+	AddGlobalFlag("rsh-timeout", "t", "Timeout for HTTP requests", time.Duration(0), false)
 
 	Root.RegisterFlagCompletionFunc("rsh-output-format", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 		return []string{"auto", "json", "yaml"}, cobra.ShellCompDirectiveNoFileComp
@@ -770,6 +772,12 @@ func Run() (returnErr error) {
 	}
 	profile, _ := GlobalFlags.GetString("rsh-profile")
 	viper.Set("rsh-profile", profile)
+	if retries, _ := GlobalFlags.GetInt("rsh-retry"); retries > 0 {
+		viper.Set("rsh-retry", retries)
+	}
+	if timeout, _ := GlobalFlags.GetDuration("rsh-timeout"); timeout > 0 {
+		viper.Set("rsh-timeout", timeout)
+	}
 
 	// Now that global flags are parsed we can enable verbose mode if requested.
 	if viper.GetBool("rsh-verbose") {

cli/cli_test.go

@@ -23,6 +23,9 @@ func reset(color bool) {
 		viper.Set("nocolor", true)
 	}
 
+	// Most tests are easier to write without retries.
+	viper.Set("rsh-retry", 0)
+
 	Init("test", "1.0.0'")
 	Defaults()
 }

cli/flag.go

@@ -3,6 +3,7 @@ package cli
 import (
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/spf13/viper"
 )
@@ -21,6 +22,13 @@ func AddGlobalFlag(name, short, description string, defaultValue interface{}, mu
 			flags.BoolP(name, short, viper.GetBool(name), description)
 			GlobalFlags.BoolP(name, short, viper.GetBool(name), description)
 		}
+	case time.Duration:
+		if multi {
+			panic(fmt.Errorf("unsupported float slice param"))
+		} else {
+			flags.DurationP(name, short, viper.GetDuration(name), description)
+			GlobalFlags.DurationP(name, short, viper.GetDuration(name), description)
+		}
 	case int, int16, int32, int64, uint16, uint32, uint64:
 		if multi {
 			flags.IntSliceP(name, short, viper.Get(name).([]int), description)

cli/request.go

@@ -1,8 +1,11 @@
 package cli
 
 import (
+	"bytes"
+	"context"
 	"crypto/tls"
 	"crypto/x509"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -100,8 +103,6 @@ func IgnoreStatus() requestOption {
 // before sending it out on the wire. If verbose mode is enabled, it will
 // print out both the request and response.
 func MakeRequest(req *http.Request, options ...requestOption) (*http.Response, error) {
-	start := time.Now()
-
 	name, config := findAPI(req.URL.String())
 
 	if config == nil {
@@ -258,11 +259,7 @@ func MakeRequest(req *http.Request, options ...requestOption) (*http.Response, e
 		}
 	}
 
-	if log {
-		LogDebugRequest(req)
-	}
-
-	resp, err := client.Do(req)
+	resp, err := doRequestWithRetry(log, client, req)
 	if err != nil {
 		return nil, err
 	}
@@ -271,11 +268,109 @@ func MakeRequest(req *http.Request, options ...requestOption) (*http.Response, e
 		lastStatus = resp.StatusCode
 	}
 
-	if log {
-		LogDebugResponse(start, resp)
+	return resp, nil
+}
+
+// isRetryable returns true if a request should be retried.
+func isRetryable(code int) bool {
+	if code == /* 408 */ http.StatusRequestTimeout ||
+		code == /*  425 */ http.StatusTooEarly ||
+		code == /*  429 */ http.StatusTooManyRequests ||
+		code == /*  500 */ http.StatusInternalServerError ||
+		code == /*  502 */ http.StatusBadGateway ||
+		code == /*  503 */ http.StatusServiceUnavailable ||
+		code == /*  504 */ http.StatusGatewayTimeout {
+		return true
+	}
+	return false
+}
+
+// doRequestWithRetry logs and makes a request, retrying as needed (if
+// configured) and returning the last response.
+func doRequestWithRetry(log bool, client *http.Client, req *http.Request) (*http.Response, error) {
+	retries := viper.GetInt("rsh-retry")
+
+	if retries == 0 {
+		return client.Do(req)
 	}
 
-	return resp, nil
+	var bodyContents []byte
+	if req.Body != nil {
+		bodyContents, _ = io.ReadAll(req.Body)
+	}
+
+	var resp *http.Response
+	var err error
+	triesLeft := 1 + retries
+	for triesLeft > 0 {
+		triesLeft--
+
+		if len(bodyContents) > 0 {
+			// Reset the body reader for each retry.
+			req.Body = io.NopCloser(bytes.NewReader(bodyContents))
+		}
+
+		if log {
+			LogDebugRequest(req)
+		}
+
+		if timeout := viper.GetDuration("rsh-timeout"); timeout > 0 {
+			ctx, cancel := context.WithTimeout(req.Context(), timeout)
+			defer cancel()
+			req = req.WithContext(ctx)
+		}
+
+		start := time.Now()
+		resp, err = client.Do(req)
+		if err != nil {
+			if errors.Is(err, context.DeadlineExceeded) {
+				if triesLeft > 0 {
+					// Try again after letting the user know.
+					LogWarning("Got request timeout after %s, retrying", viper.GetDuration("rsh-timeout").Truncate(time.Millisecond))
+					continue
+				} else {
+					// Add a human-friendly error before the original (context deadline
+					// exceeded).
+					err = fmt.Errorf("Request timed out after %s: %w", viper.GetDuration("rsh-timeout"), err)
+				}
+			}
+			return resp, err
+		}
+
+		if log {
+			LogDebugResponse(start, resp)
+		}
+
+		if triesLeft > 0 && isRetryable(resp.StatusCode) {
+			// Attempt to parse when to retry! Default is 1 second.
+			retryAfter := 1 * time.Second
+
+			if v := resp.Header.Get("Retry-After"); v != "" {
+				// Could be either an integer number of seconds, or an HTTP date.
+				if d, err := strconv.ParseInt(v, 10, 64); err == nil {
+					retryAfter = time.Duration(d) * time.Second
+				}
+
+				if d, err := http.ParseTime(v); err == nil {
+					retryAfter = time.Until(d)
+				}
+			}
+
+			if v := resp.Header.Get("X-Retry-In"); v != "" {
+				if d, err := time.ParseDuration(v); err == nil {
+					retryAfter = d
+				}
+			}
+
+			LogWarning("Got %s, retrying in %s", resp.Status, retryAfter.Truncate(time.Millisecond))
+			time.Sleep(retryAfter)
+
+			continue
+		}
+		break
+	}
+
+	return resp, err
 }
 
 // Response describes a parsed HTTP response which can be marshalled to enable

cli/request_test.go

@@ -1,10 +1,13 @@
 package cli
 
 import (
+	"bytes"
 	"errors"
 	"net/http"
 	"testing"
+	"time"
 
+	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
 	"gopkg.in/h2non/gock.v1"
 )
@@ -123,3 +126,83 @@ func TestIgnoreStatus(t *testing.T) {
 
 	assert.Equal(t, 0, GetLastStatus())
 }
+
+func TestRequestRetryIn(t *testing.T) {
+	defer gock.Off()
+
+	reset(false)
+	viper.Set("rsh-retry", 1)
+
+	// Duration string value (with units)
+	gock.New("http://example.com").
+		Get("/").
+		Times(1).
+		Reply(http.StatusTooManyRequests).
+		SetHeader("X-Retry-In", "1ms")
+
+	gock.New("http://example.com").
+		Get("/").
+		Times(1).
+		Reply(http.StatusOK)
+
+	req, _ := http.NewRequest(http.MethodGet, "http://example.com/", nil)
+	resp, err := MakeRequest(req)
+
+	assert.NoError(t, err)
+	assert.Equal(t, resp.StatusCode, http.StatusOK)
+}
+
+func TestRequestRetryAfter(t *testing.T) {
+	defer gock.Off()
+
+	reset(false)
+	viper.Set("rsh-retry", 2)
+
+	// Seconds value
+	gock.New("http://example.com").
+		Put("/").
+		Times(1).
+		Reply(http.StatusTooManyRequests).
+		SetHeader("Retry-After", "0")
+
+	// HTTP date value
+	gock.New("http://example.com").
+		Put("/").
+		Times(1).
+		Reply(http.StatusTooManyRequests).
+		SetHeader("Retry-After", time.Now().Format(http.TimeFormat))
+
+	gock.New("http://example.com").
+		Put("/").
+		Times(1).
+		Reply(http.StatusOK)
+
+	req, _ := http.NewRequest(http.MethodPut, "http://example.com/", bytes.NewReader([]byte("hello")))
+	resp, err := MakeRequest(req)
+
+	assert.NoError(t, err)
+	assert.Equal(t, resp.StatusCode, http.StatusOK)
+}
+
+func TestRequestRetryTimeout(t *testing.T) {
+	defer gock.Off()
+
+	reset(false)
+	viper.Set("rsh-retry", 1)
+	viper.Set("rsh-timeout", 1*time.Millisecond)
+
+	// Duration string value (with units)
+	gock.New("http://example.com").
+		Get("/").
+		Times(2).
+		Reply(http.StatusOK).
+		Delay(2 * time.Millisecond)
+		// Note: delay seems to have a bug where subsequent requests without the
+		// delay are still delayed... For now just have it reply twice.
+
+	req, _ := http.NewRequest(http.MethodGet, "http://example.com/", nil)
+	_, err := MakeRequest(req)
+
+	assert.Error(t, err)
+	assert.ErrorContains(t, err, "timed out")
+}

docs/README.md

@@ -41,6 +41,7 @@ Start with the [guide](/guide.md) to learn how to install and configure Restish
   - MessagePack (<https://msgpack.org/>)
   - Amazon Ion (<http://amzn.github.io/ion-docs/>)
   - Gzip ([RFC 1952](https://tools.ietf.org/html/rfc1952)), Deflate ([RFC 1951](https://datatracker.ietf.org/doc/html/rfc1951)), and Brotli ([RFC 7932](https://tools.ietf.org/html/rfc7932)) content encoding
+- Automatic retries with support for [`Retry-After`](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After) and `X-Retry-In` headers when APIs are rate-limited.
 - Standardized [hypermedia](https://smartbear.com/learn/api-design/what-is-hypermedia/) parsing into queryable/followable response links:
   - HTTP Link relation headers ([RFC 5988](https://tools.ietf.org/html/rfc5988#section-6.2.2))
   - [HAL](http://stateless.co/hal_specification.html)

docs/_sidebar.md

@@ -6,5 +6,6 @@
 - [Input](input.md "Restish Input")
 - [CLI Shorthand](shorthand.md "CLI Shorthand")
 - [Output](output.md "Restish Output")
+- [Retries & Timeouts](retries.md "Retries & Timeouts")
 - [Hypermedia](hypermedia.md "Hypermedia Linking in Restish")
 - [Bulk Management](bulk.md "Bulk Resource Management")

docs/index.html

@@ -169,6 +169,10 @@
       .token.diffRemoved {
         color: #ff5f87;
       }
+
+      .token.warning {
+        color: #d78700;
+      }
     </style>
   </head>
   <body>
@@ -236,7 +240,7 @@
         comment: /^\s*#.*/m,
         redirect: /2>\/dev\/null/,
         response: {
-          pattern: /^(HTTP\/[12]|\{|\[)(.|\n)+(\]|\}(\n|$))/gm,
+          pattern: /^(HTTP\/[12].*$)|((\{|\[)(.|\n)+(\]|\}(\n|$)))/gm,
           greedy: false,
           inside: Prism.languages.readable
         },
@@ -259,6 +263,13 @@
           alias: "variable",
           pattern: /\$[A-Z0-9_]+/,
         },
+        log: {
+          pattern: /(INFO|WARN|ERROR): .*/,
+          inside: {
+            warning: /WARN:/,
+            error: /ERROR:/,
+          },
+        },
         header: {
           pattern: /-H [A-Z][a-zA-Z0-9-]+:\S+/,
           inside: {