logcli: Add retries to unsuccessful log queries (#3879)

* Add retries to unsuccessful log queries.

cmd/logcli/main.go

@@ -211,6 +211,7 @@ func newQueryClient(app *kingpin.Application) client.Client {
 	app.Flag("org-id", "adds X-Scope-OrgID to API requests for representing tenant ID. Useful for requesting tenant data when bypassing an auth gateway.").Default("").Envar("LOKI_ORG_ID").StringVar(&client.OrgID)
 	app.Flag("bearer-token", "adds the Authorization header to API requests for authentication purposes. Can also be set using LOKI_BEARER_TOKEN env var.").Default("").Envar("LOKI_BEARER_TOKEN").StringVar(&client.BearerToken)
 	app.Flag("bearer-token-file", "adds the Authorization header to API requests for authentication purposes. Can also be set using LOKI_BEARER_TOKEN_FILE env var.").Default("").Envar("LOKI_BEARER_TOKEN_FILE").StringVar(&client.BearerTokenFile)
+	app.Flag("retries", "How many times to retry each query when getting an error response from Loki. Can also be set using LOKI_CLIENT_RETRIES").Default("0").Envar("LOKI_CLIENT_RETRIES").IntVar(&client.Retries)
 
 	return client
 }

pkg/logcli/client/client.go

@@ -56,6 +56,7 @@ type DefaultClient struct {
 	Tripperware     Tripperware
 	BearerToken     string
 	BearerTokenFile string
+	Retries         int
 }
 
 // Query uses the /api/v1/query endpoint to execute an instant query
@@ -216,21 +217,39 @@ func (c *DefaultClient) doRequest(path, query string, quiet bool, out interface{
 	if c.Tripperware != nil {
 		client.Transport = c.Tripperware(client.Transport)
 	}
-	resp, err := client.Do(req)
-	if err != nil {
-		return err
+
+	var resp *http.Response
+	attempts := c.Retries + 1
+	success := false
+
+	for attempts > 0 {
+		attempts--
+
+		resp, err = client.Do(req)
+		if err != nil {
+			log.Println("error sending request", err)
+			continue
+		}
+		if resp.StatusCode/100 != 2 {
+			buf, _ := ioutil.ReadAll(resp.Body) // nolint
+			log.Printf("Error response from server: %s (%v) attempts remaining: %d", string(buf), err, attempts)
+			if err := resp.Body.Close(); err != nil {
+				log.Println("error closing body", err)
+			}
+			continue
+		}
+		success = true
+		break
+	}
+	if !success {
+		return fmt.Errorf("Run out of attempts while querying the server")
 	}
+
 	defer func() {
 		if err := resp.Body.Close(); err != nil {
 			log.Println("error closing body", err)
 		}
 	}()
-
-	if resp.StatusCode/100 != 2 {
-		buf, _ := ioutil.ReadAll(resp.Body) // nolint
-		return fmt.Errorf("Error response from server: %s (%v)", string(buf), err)
-	}
-
 	return json.NewDecoder(resp.Body).Decode(out)
 }