From 2bce39644554b3a7eee8ec6af23511ebd55f64d2 Mon Sep 17 00:00:00 2001
From: Matthew Jaffee <jaffee@pilosa.com>
Date: Mon, 20 Dec 2021 22:16:06 -0600
Subject: [PATCH] add retry restore test and custom retry policy

---
 ctl/restore.go                        | 10 ++++++++
 internal/clustertests/cluster_test.go | 35 +++++++++++++++++++++++----
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/ctl/restore.go b/ctl/restore.go
index 0804c3fa2..cf99619cb 100644
--- a/ctl/restore.go
+++ b/ctl/restore.go
@@ -49,6 +49,7 @@ type RestoreCommand struct {
 func NewRestoreCommand(stdin io.Reader, stdout, stderr io.Writer) *RestoreCommand {
 	return &RestoreCommand{
 		CmdIO:       pilosa.NewCmdIO(stdin, stdout, stderr),
+		RetryPeriod: time.Second * 30,
 		Concurrency: 1,
 	}
 }
@@ -168,6 +169,13 @@ func (cmd *RestoreCommand) restoreSchema(ctx context.Context, primary *topology.
 	return err
 }
 
+func RetryWith400(ctx context.Context, resp *http.Response, err error) (bool, error) {
+	if resp != nil && resp.StatusCode > 400 { // we have some dumb status codes
+		return true, nil
+	}
+	return retryablehttp.DefaultRetryPolicy(ctx, resp, err)
+}
+
 func (cmd *RestoreCommand) restoreIDAlloc(ctx context.Context, primary *topology.Node) error {
 	logger := cmd.Logger()
 
@@ -185,6 +193,7 @@ func (cmd *RestoreCommand) restoreIDAlloc(ctx context.Context, primary *topology
 
 	client := retryablehttp.NewClient()
 	client.RetryWaitMax = cmd.RetryPeriod
+	client.CheckRetry = RetryWith400
 	_, err = client.Post(url, "application/octet-stream", f)
 	return err
 }
@@ -263,6 +272,7 @@ func (cmd *RestoreCommand) restoreShard(ctx context.Context, filename string) er
 
 		client := retryablehttp.NewClient()
 		client.RetryWaitMax = cmd.RetryPeriod
+		client.CheckRetry = RetryWith400
 		resp, err := client.Do(req)
 		if err != nil {
 			return err
diff --git a/internal/clustertests/cluster_test.go b/internal/clustertests/cluster_test.go
index 00967ac07..af89abc2a 100644
--- a/internal/clustertests/cluster_test.go
+++ b/internal/clustertests/cluster_test.go
@@ -4,7 +4,7 @@ package clustertest
 import (
 	"context"
 	"fmt"
-	"io/ioutil"
+	"net/http"
 	"os"
 	"os/exec"
 	"testing"
@@ -104,10 +104,7 @@ func TestClusterStuff(t *testing.T) {
 			t.Fatalf("sending stop command: %v", err)
 		}
 		var backupCmd *exec.Cmd
-		tmpdir, err := ioutil.TempDir("", "")
-		if err != nil {
-			t.Fatalf("getting tmp dir: %v", err)
-		}
+		tmpdir := t.TempDir()
 		if backupCmd, err = startCmd(
 			"featurebase", "backup", "--host=pilosa1:10101", fmt.Sprintf("--output=%s", tmpdir+"/backuptest")); err != nil {
 			t.Fatalf("sending backup command: %v", err)
@@ -121,6 +118,34 @@ func TestClusterStuff(t *testing.T) {
 			t.Fatalf("waiting on backup to finish: %v", err)
 		}
 
+		fmt.Println("STARTING RESTORE")
+
+		client := http.Client{}
+		if req, err := http.NewRequest(http.MethodDelete, "http://pilosa1:10101/index/testidx", nil); err != nil {
+			t.Fatalf("getting req: %v", err)
+		} else if resp, err := client.Do(req); err != nil {
+			t.Fatalf("doing request: %v", err)
+		} else if resp.StatusCode >= 400 {
+			t.Fatalf("bad response: %v", resp)
+		}
+
+		var restoreCmd *exec.Cmd
+		if restoreCmd, err = startCmd("featurebase", "restore", "-s", tmpdir+"/backuptest", "--host", "pilosa1:10101"); err != nil {
+			t.Fatalf("starting restore: %v", err)
+		}
+		time.Sleep(time.Millisecond * 50)
+		if err = sendCmd("docker", "stop", "clustertests_pilosa2_1"); err != nil {
+			t.Fatalf("sending stop command: %v", err)
+		}
+
+		time.Sleep(time.Second * 10)
+		if err = sendCmd("docker", "start", "clustertests_pilosa2_1"); err != nil {
+			t.Fatalf("sending stop command: %v", err)
+		}
+		if err := restoreCmd.Wait(); err != nil {
+			t.Fatalf("restore failed: %v", err)
+		}
+
 		// now do backup with all nodes down and too short a timeout
 		// so it fails. Has be to be all 3 because the cluster has
 		// replicas=3 and the backup command will retry on replicas.