Merge pull request #1453 from pulumi/retry-post

Send request body on retries

Gopkg.lock

@@ -548,6 +548,6 @@
 [solve-meta]
   analyzer-name = "dep"
   analyzer-version = 1
-  inputs-digest = "3761aa14e049c63368b33a19c8e3b6d16bb63fd43d3f3f4738cdcabc2bd0f2b9"
+  inputs-digest = "981c1124c73b9ec3faa288afb4482a8376ccf9a3da5f24a4d136f345d6b95beb"
   solver-name = "gps-cdcl"
   solver-version = 1

pkg/backend/cloud/client/api.go

@@ -116,7 +116,7 @@ func pulumiAPICall(ctx context.Context, cloudAPI, method, path string, body []by
 	path = cleanPath(path)
 
 	url := fmt.Sprintf("%s%s", cloudAPI, path)
-	req, err := http.NewRequest(method, url, bytes.NewBuffer(body))
+	req, err := http.NewRequest(method, url, bytes.NewReader(body))
 	if err != nil {
 		return "", nil, errors.Wrapf(err, "creating new HTTP request")
 	}

pkg/util/httputil/http.go

@@ -28,12 +28,24 @@ const maxRetryCount = 5
 
 // DoWithRetry calls client.Do, and in the case of an error, retries the operation again after a slight delay.
 func DoWithRetry(req *http.Request, client *http.Client) (*http.Response, error) {
+	contract.Assertf(req.ContentLength == 0 || req.GetBody != nil,
+		"Retryable request must have no body or rewindable body")
+
 	inRange := func(test, lower, upper int) bool {
 		return lower <= test && test <= upper
 	}
 
 	_, res, err := retry.Until(context.Background(), retry.Acceptor{
 		Accept: func(try int, nextRetryTime time.Duration) (bool, interface{}, error) {
+			if try > 0 && req.GetBody != nil {
+				// Reset request body, if present, for retries.
+				rc, bodyErr := req.GetBody()
+				if bodyErr != nil {
+					return false, nil, bodyErr
+				}
+				req.Body = rc
+			}
+
 			res, resErr := client.Do(req)
 			if resErr == nil && !inRange(res.StatusCode, 500, 599) {
 				return true, res, nil

pkg/util/httputil/http_test.go

@@ -0,0 +1,94 @@
+// Copyright 2016-2018, Pulumi Corporation.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package httputil
+
+import (
+	"crypto/tls"
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"strings"
+	"testing"
+
+	"golang.org/x/net/http2"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func http2ServerAndClient(handler http.Handler) (*httptest.Server, *http.Client) {
+	// Create an HTTP/2 test server.
+	// httptest.StartTLS will set NextProtos to ["http/1.1"] if it's unset, so we need to add
+	// HTTP/2 eagerly before starting the server.
+	server := httptest.NewUnstartedServer(handler)
+	server.TLS = &tls.Config{
+		NextProtos: []string{http2.NextProtoTLS},
+	}
+	server.StartTLS()
+
+	// Create a client for the test server that will use HTTP/2.
+	// We need a client that will (a) upgrade to HTTP/2 and (b) trust the test server's certs.
+	// In order to satisfy (b), httptest sets Transport to an `http.Transport`, breaking (a),
+	// so we have to manually create an `http2.Transport` and copy over the `tls.Config`.
+	tlsConfig := server.Client().Transport.(*http.Transport).TLSClientConfig
+	client := &http.Client{
+		Transport: &http2.Transport{
+			TLSClientConfig: tlsConfig,
+		},
+	}
+
+	return server, client
+}
+
+// Test that DoWithRetry rewinds and resends the request body when retrying POSTs over HTTP/2.
+func TestRetryPostHTTP2(t *testing.T) {
+	tries := 0
+	handler := func(w http.ResponseWriter, r *http.Request) {
+		tries++
+		t.Logf("try %d", tries)
+
+		assert.Equal(t, "HTTP/2.0", r.Proto)
+
+		// Check that the body's content length matches the sent data.
+		defer r.Body.Close()
+		content, err := ioutil.ReadAll(r.Body)
+		assert.NoError(t, err)
+		assert.Equal(t, strconv.Itoa(len(content)), r.Header.Get("Content-Length"))
+
+		// Check the message matches.
+		assert.Equal(t, string(content), "hello, server")
+
+		// Fail the first try with 500, which will prompt a retry.
+		switch tries {
+		case 1:
+			w.WriteHeader(500)
+		default:
+			w.WriteHeader(200)
+		}
+	}
+
+	server, client := http2ServerAndClient(http.HandlerFunc(handler))
+
+	req, err := http.NewRequest("POST", server.URL, strings.NewReader("hello, server"))
+	assert.NoError(t, err)
+
+	res, err := DoWithRetry(req, client)
+	assert.NoError(t, err)
+	defer res.Body.Close()
+
+	// Check that the request succeeded on the second try.
+	assert.Equal(t, 2, tries)
+	assert.Equal(t, 200, res.StatusCode)
+}