use retryableHTTP in client, fix memory usage of restore

instead of awkwardly reading an entire file into a buffer, we use
retryablehttp's reader func to open the file fresh if we need to
retry, so a small fixed-size buffer can be used internally for copying
the contents onto the network.

client.go

@@ -80,8 +80,14 @@ type InternalClient interface {
 	GetNodeUsage(ctx context.Context, uri *pnet.URI) (map[string]NodeUsage, error)
 	GetPastQueries(ctx context.Context, uri *pnet.URI) ([]PastQueryStatus, error)
 
-	ImportFieldKeys(ctx context.Context, uri *pnet.URI, index, field string, remote bool, rddbdata io.Reader) error
-	ImportIndexKeys(ctx context.Context, uri *pnet.URI, index string, partitionID int, remote bool, rddbdata io.Reader) error
+	// ImportFieldKeys and ImportIndexKeys are mainly used when
+	// restoring a backup. They take a readerFunc which returns a
+	// reader rather than taking an io.Reader directly to allow for
+	// efficient retries (rather than reading the entire request body
+	// into a buffer and reusing it). Reader returned from the func
+	// must be properly closed by the implementation.
+	ImportFieldKeys(ctx context.Context, uri *pnet.URI, index, field string, remote bool, readerFunc func() (io.Reader, error)) error
+	ImportIndexKeys(ctx context.Context, uri *pnet.URI, index string, partitionID int, remote bool, readerFunc func() (io.Reader, error)) error
 
 	// SetInternalAPI tells the client the API it should use for internal/loopback ops
 	// where applicable.
@@ -277,11 +283,11 @@ func (n nopInternalClient) GetNodeUsage(ctx context.Context, uri *pnet.URI) (map
 func (n nopInternalClient) GetPastQueries(ctx context.Context, uri *pnet.URI) ([]PastQueryStatus, error) {
 	return nil, nil
 }
-func (c nopInternalClient) ImportFieldKeys(ctx context.Context, uri *pnet.URI, index, field string, remote bool, rddbdata io.Reader) error {
+func (c nopInternalClient) ImportFieldKeys(ctx context.Context, uri *pnet.URI, index, field string, remote bool, readerFunc func() (io.Reader, error)) error {
 	return nil
 }
 
-func (c nopInternalClient) ImportIndexKeys(ctx context.Context, uri *pnet.URI, index string, partitionID int, remote bool, rddbdata io.Reader) error {
+func (c nopInternalClient) ImportIndexKeys(ctx context.Context, uri *pnet.URI, index string, partitionID int, remote bool, readerFunc func() (io.Reader, error)) error {
 	return nil
 }
 

cmd/slurp/slurp.go

@@ -92,8 +92,10 @@ func (r *stateMachine) NewHeader(h *tar.Header, tr *tar.Reader) error {
 
 			byteData, err := ioutil.ReadAll(tr)
 			vprint.PanicOn(err)
-			br := bytes.NewReader(byteData)
-			err = r.client.ImportFieldKeys(context.Background(), uri, index, fieldName, false, br)
+			readerFunc := func() (io.Reader, error) {
+				return bytes.NewReader(byteData), nil
+			}
+			err = r.client.ImportFieldKeys(context.Background(), uri, index, fieldName, false, readerFunc)
 			if err != nil {
 				return err
 			}
@@ -106,9 +108,11 @@ func (r *stateMachine) NewHeader(h *tar.Header, tr *tar.Reader) error {
 			}
 			byteData, err := ioutil.ReadAll(tr)
 			vprint.PanicOn(err)
+			readerFunc := func() (io.Reader, error) {
+				return bytes.NewReader(byteData), nil
+			}
 
-			br := bytes.NewReader(byteData)
-			err = r.client.ImportIndexKeys(context.Background(), uri, index, int(partition), false, br)
+			err = r.client.ImportIndexKeys(context.Background(), uri, index, int(partition), false, readerFunc)
 			if err != nil {
 				return err
 			}

ctl/restore.go

@@ -178,7 +178,7 @@ func (cmd *RestoreCommand) restoreSchema(ctx context.Context, primary *topology.
 	return err
 }
 
-func RetryWith400(ctx context.Context, resp *http.Response, err error) (bool, error) {
+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
 	}
@@ -202,7 +202,7 @@ func (cmd *RestoreCommand) restoreIDAlloc(ctx context.Context, primary *topology
 
 	client := retryablehttp.NewClient()
 	client.RetryWaitMax = cmd.RetryPeriod
-	client.CheckRetry = RetryWith400
+	client.CheckRetry = retryWith400
 	_, err = client.Post(url, "application/octet-stream", f)
 	return err
 }
@@ -281,7 +281,7 @@ func (cmd *RestoreCommand) restoreShard(ctx context.Context, filename string) er
 
 		client := retryablehttp.NewClient()
 		client.RetryWaitMax = cmd.RetryPeriod
-		client.CheckRetry = RetryWith400
+		client.CheckRetry = retryWith400
 		resp, err := client.Do(req)
 		if err != nil {
 			return err
@@ -349,13 +349,11 @@ func (cmd *RestoreCommand) restoreIndexTranslationFile(ctx context.Context, file
 
 	for _, node := range nodes {
 		if err := func() error {
-			f, err := os.Open(filename)
-			if err != nil {
-				return err
+			readerFunc := func() (io.Reader, error) {
+				return os.Open(filename) // gets used as an HTTP request body and closed by http library
 			}
-			defer f.Close()
 
-			return cmd.client.ImportIndexKeys(ctx, &node.URI, indexName, partitionID, false, f)
+			return cmd.client.ImportIndexKeys(ctx, &node.URI, indexName, partitionID, false, readerFunc)
 		}(); err != nil {
 			return err
 		}
@@ -410,13 +408,11 @@ func (cmd *RestoreCommand) restoreFieldTranslationFile(ctx context.Context, node
 
 	for _, node := range nodes {
 		if err := func() error {
-			f, err := os.Open(filename)
-			if err != nil {
-				return err
+			readerFunc := func() (io.Reader, error) {
+				return os.Open(filename)
 			}
-			defer f.Close()
 
-			return cmd.client.ImportFieldKeys(ctx, &node.URI, indexName, fieldName, false, f)
+			return cmd.client.ImportFieldKeys(ctx, &node.URI, indexName, fieldName, false, readerFunc)
 		}(); err != nil {
 			return err
 		}

http/client.go

@@ -18,6 +18,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/hashicorp/go-retryablehttp"
 	pilosa "github.com/molecula/featurebase/v2"
 	"github.com/molecula/featurebase/v2/encoding/proto"
 	"github.com/molecula/featurebase/v2/ingest"
@@ -35,10 +36,9 @@ type InternalClient struct {
 
 	log logger.Logger
 
-	retryPeriod time.Duration
-
 	// The client to use for HTTP communication.
-	httpClient *http.Client
+	httpClient      *http.Client
+	retryableClient *retryablehttp.Client
 	// the local node's API, used for operations that we can short-circuit that way
 	api *pilosa.API
 }
@@ -64,9 +64,13 @@ type InternalClientOption func(c *InternalClient)
 
 // WithClientRetryPeriod is the max amount of total time the client will
 // retry failed requests using exponential backoff.
-func WithClientRetryPeriod(period time.Duration) InternalClientOption {
+func WithClientRetryPeriod(waitMax time.Duration) InternalClientOption {
 	return func(c *InternalClient) {
-		c.retryPeriod = period
+		c.retryableClient = &retryablehttp.Client{
+			HTTPClient:   c.httpClient,
+			RetryWaitMax: waitMax,
+			CheckRetry:   retryWith400Policy,
+		}
 	}
 }
 
@@ -76,6 +80,21 @@ func WithClientLogger(log logger.Logger) InternalClientOption {
 	}
 }
 
+func noRetryPolicy(ctx context.Context, resp *http.Response, err error) (bool, error) {
+	return false, nil
+}
+
+// retryWith400Policy wraps retryablehttp's default retry policy to
+// also retry on 4XX errors which *should* be client errors and
+// therefore useless to retry, but we have some incorrect status codes.
+// TODO: fix the incorrect status codes so we can get rid of this.
+func retryWith400Policy(ctx context.Context, resp *http.Response, err error) (bool, error) {
+	if resp != nil && resp.StatusCode > 400 {
+		return true, nil
+	}
+	return retryablehttp.DefaultRetryPolicy(ctx, resp, err)
+}
+
 func NewInternalClientFromURI(defaultURI *pnet.URI, remoteClient *http.Client, opts ...InternalClientOption) *InternalClient {
 	ic := &InternalClient{
 		defaultURI: defaultURI,
@@ -87,6 +106,13 @@ func NewInternalClientFromURI(defaultURI *pnet.URI, remoteClient *http.Client, o
 	for _, opt := range opts {
 		opt(ic)
 	}
+
+	if ic.retryableClient == nil {
+		ic.retryableClient = &retryablehttp.Client{
+			HTTPClient: ic.httpClient,
+			CheckRetry: noRetryPolicy,
+		}
+	}
 	return ic
 }
 
@@ -1753,57 +1779,28 @@ func (n nopCloser) Close() error {
 	return nil
 }
 
-func (c *InternalClient) doWithRetry(req *http.Request) (*http.Response, error) {
-	sleepDuration := time.Second
-	var bod []byte
-	var err error
-	if req.Body != nil {
-		bod, err = ioutil.ReadAll(req.Body)
-		if err != nil {
-			return nil, errors.Wrap(err, "reading body")
-		}
-		req.Body = nopCloser{bytes.NewReader(bod)}
-	}
-	resp, err := c.httpClient.Do(req)
-	// start timer after first request, so if retryPeriod > 0 we
-	// pretty much always do at least one retry
-	start := time.Now()
-	for ; err != nil || resp.StatusCode < 200 || resp.StatusCode >= 300; resp, err = c.httpClient.Do(req) {
-		if req.Body != nil {
-			req.Body = nopCloser{bytes.NewReader(bod)} // can't seek due to races with http lib internals
-		}
-		if time.Since(start) > c.retryPeriod {
-			break
-		}
-		if err != nil {
-			c.log.Printf("retrying request due to error: '%v'", err)
-		} else {
-			if bod, readErr := ioutil.ReadAll(resp.Body); readErr != nil {
-				c.log.Printf("retrying request due to status: %d, error reading body: '%v', body: '%s'", resp.StatusCode, readErr, bod)
-			} else {
-				c.log.Printf("retrying request due to status: %d, body: '%s'", resp.StatusCode, bod)
-			}
-		}
-
-		time.Sleep(sleepDuration)
-		sleepDuration *= 2
-	}
-	return resp, err
-}
-
 // executeRequest executes the given request and checks the Response. For
 // responses with non-2XX status, the body is read and closed, and an error is
 // returned. If the error is nil, the caller must ensure that the response body
 // is closed.
 func (c *InternalClient) executeRequest(req *http.Request, opts ...executeRequestOption) (*http.Response, error) {
+	return c.executeRetryableRequest(&retryablehttp.Request{Request: req}, opts...)
+}
+
+func (c *InternalClient) executeRetryableRequest(req *retryablehttp.Request, opts ...executeRequestOption) (*http.Response, error) {
+	tracing.GlobalTracer.InjectHTTPHeaders(req.Request)
+	req.Close = false
 	eo := &executeOpts{}
 	for _, opt := range opts {
 		opt(eo)
 	}
 
-	tracing.GlobalTracer.InjectHTTPHeaders(req)
-	req.Close = false
-	resp, err := c.doWithRetry(req)
+	resp, err := c.retryableClient.Do(req)
+
+	return c.handleResponse(req.Request, eo, resp, err)
+}
+
+func (c *InternalClient) handleResponse(req *http.Request, eo *executeOpts, resp *http.Response, err error) (*http.Response, error) {
 	if err != nil {
 		if resp != nil {
 			resp.Body.Close()
@@ -2083,7 +2080,7 @@ func (c *InternalClient) RetrieveTranslatePartitionFromURI(ctx context.Context,
 
 	return resp.Body, nil
 }
-func (c *InternalClient) ImportIndexKeys(ctx context.Context, uri *pnet.URI, index string, partitionID int, remote bool, rddbdata io.Reader) error {
+func (c *InternalClient) ImportIndexKeys(ctx context.Context, uri *pnet.URI, index string, partitionID int, remote bool, readerFunc func() (io.Reader, error)) error {
 	span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.ImportIndexKeys")
 	defer span.Finish()
 
@@ -2100,14 +2097,14 @@ func (c *InternalClient) ImportIndexKeys(ctx context.Context, uri *pnet.URI, ind
 	url := fmt.Sprintf("%s/internal/translate/index/%s/%d", uri, index, partitionID)
 
 	// Generate HTTP request.
-	httpReq, err := http.NewRequest("POST", url, rddbdata)
+	httpReq, err := retryablehttp.NewRequest("POST", url, readerFunc)
 	if err != nil {
 		return errors.Wrap(err, "creating request")
 	}
 	httpReq.Header.Set("User-Agent", "pilosa/"+pilosa.Version)
 
 	// Execute request against the host.
-	resp, err := c.executeRequest(httpReq.WithContext(ctx))
+	resp, err := c.executeRetryableRequest(httpReq.WithContext(ctx))
 	if err != nil {
 		return err
 	}
@@ -2115,7 +2112,7 @@ func (c *InternalClient) ImportIndexKeys(ctx context.Context, uri *pnet.URI, ind
 	return nil
 }
 
-func (c *InternalClient) ImportFieldKeys(ctx context.Context, uri *pnet.URI, index, field string, remote bool, rddbdata io.Reader) error {
+func (c *InternalClient) ImportFieldKeys(ctx context.Context, uri *pnet.URI, index, field string, remote bool, readerFunc func() (io.Reader, error)) error {
 	span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.ImportFieldKeys")
 	defer span.Finish()
 
@@ -2132,14 +2129,14 @@ func (c *InternalClient) ImportFieldKeys(ctx context.Context, uri *pnet.URI, ind
 	url := fmt.Sprintf("%s/internal/translate/field/%s/%s", uri, index, field)
 
 	// Generate HTTP request.
-	httpReq, err := http.NewRequest("POST", url, rddbdata)
+	httpReq, err := retryablehttp.NewRequest("POST", url, readerFunc)
 	if err != nil {
 		return errors.Wrap(err, "creating request")
 	}
 	httpReq.Header.Set("User-Agent", "pilosa/"+pilosa.Version)
 
 	// Execute request against the host.
-	resp, err := c.executeRequest(httpReq.WithContext(ctx))
+	resp, err := c.executeRetryableRequest(httpReq.WithContext(ctx))
 	if err != nil {
 		return err
 	}