Fix AzDo Throttling Timeouts

cmd/server.go

@@ -35,12 +35,15 @@ import (
 // 1. Add a const with the flag name (in alphabetic order).
 // 2. Add a new field to server.UserConfig and set the mapstructure tag equal to the flag name.
 // 3. Add your flag's description etc. to the stringFlags, intFlags, or boolFlags slices.
+// 4. If appropriate, modify the setDefaults function to assign teh default value
 const (
 	// Flag names.
-	ADWebhookPasswordFlag      = "azuredevops-webhook-password" // nolint: gosec
-	ADWebhookUserFlag          = "azuredevops-webhook-user"
+	ADHonorRetryAfter          = "azuredevops-honor-retry-after"
+	ADTimeoutSeconds           = "azuredevops-timeout-seconds"
 	ADTokenFlag                = "azuredevops-token" // nolint: gosec
 	ADUserFlag                 = "azuredevops-user"
+	ADWebhookPasswordFlag      = "azuredevops-webhook-password" // nolint: gosec
+	ADWebhookUserFlag          = "azuredevops-webhook-user"
 	AllowForkPRsFlag           = "allow-fork-prs"
 	AllowRepoConfigFlag        = "allow-repo-config"
 	AtlantisURLFlag            = "atlantis-url"
@@ -104,6 +107,7 @@ const (
 	// NOTE: Must manually set these as defaults in the setDefaults function.
 	DefaultADBasicUser      = ""
 	DefaultADBasicPassword  = ""
+	DefaultADTimeoutSeconds = 10
 	DefaultAutoplanFileList = "**/*.tf,**/*.tfvars,**/*.tfvars.json,**/terragrunt.hcl"
 	DefaultCheckoutStrategy = "branch"
 	DefaultBitbucketBaseURL = bitbucketcloud.BaseURL
@@ -278,6 +282,10 @@ var stringFlags = map[string]stringFlag{
 }
 
 var boolFlags = map[string]boolFlag{
+	ADHonorRetryAfter: {
+		description:  "If set to true, the Azure DevOps client will sleep according to the Retry-After response header",
+		defaultValue: false,
+	},
 	AllowForkPRsFlag: {
 		description:  "Allow Atlantis to run on pull requests from forks. A security issue for public repos.",
 		defaultValue: false,
@@ -371,6 +379,10 @@ var boolFlags = map[string]boolFlag{
 	},
 }
 var intFlags = map[string]intFlag{
+	ADTimeoutSeconds: {
+		description:  "Sets timeout seconds for the Azure DevOps api client",
+		defaultValue: DefaultADTimeoutSeconds,
+	},
 	ParallelPoolSize: {
 		description:  "Max size of the wait group that runs parallel plans and applies (if enabled).",
 		defaultValue: DefaultParallelPoolSize,
@@ -582,6 +594,9 @@ func (s *ServerCmd) setDefaults(c *server.UserConfig) {
 	if c.AutoplanFileList == "" {
 		c.AutoplanFileList = DefaultAutoplanFileList
 	}
+	if c.AzureDevopsTimeoutSeconds == 0 {
+		c.AzureDevopsTimeoutSeconds = DefaultADTimeoutSeconds
+	}
 	if c.CheckoutStrategy == "" {
 		c.CheckoutStrategy = DefaultCheckoutStrategy
 	}

cmd/server_test.go

@@ -51,6 +51,8 @@ func (s *ServerStarterMock) Start() error {
 // Adding a new flag? Add it to this slice for testing in alphabetical
 // order.
 var testFlags = map[string]interface{}{
+	ADHonorRetryAfter:          true,
+	ADTimeoutSeconds:           12,
 	ADTokenFlag:                "ad-token",
 	ADUserFlag:                 "ad-user",
 	ADWebhookPasswordFlag:      "ad-wh-pass",

server/events/vcs/azuredevops_client.go

@@ -6,6 +6,8 @@ import (
 	"net/http"
 	"net/url"
 	"path/filepath"
+	"regexp"
+	"strconv"
 	"strings"
 	"time"
 
@@ -22,14 +24,53 @@ type AzureDevopsClient struct {
 	UserName string
 }
 
+type ThrottleRoundtripper struct {
+	// If azure devops is throttling, this is how long to wait before the next request
+	RetryAfter         time.Duration
+	ParentRoundtripper http.RoundTripper
+}
+
+func NewThrottleRoundtripper(rt http.RoundTripper) http.RoundTripper {
+	return &ThrottleRoundtripper{
+		ParentRoundtripper: rt,
+	}
+}
+
+// RoundTrip follows http.Transport interface and will automatically sleep for retryAfter seconds
+// This allows us to honor the Retry-After header from Azdo
+func (rt *ThrottleRoundtripper) RoundTrip(req *http.Request) (resp *http.Response, err error) {
+	time.Sleep(rt.RetryAfter)
+	rt.RetryAfter = 0
+
+	resp, err = rt.ParentRoundtripper.RoundTrip(req)
+	if err != nil {
+		return
+	}
+
+	retryAfterVal := resp.Header.Get("retry-after")
+	if retryAfterVal != "" {
+		retryAfterSec, err := strconv.Atoi(retryAfterVal)
+		if err == nil {
+			rt.RetryAfter = time.Duration(retryAfterSec) * time.Second
+		}
+	}
+
+	return
+}
+
 // NewAzureDevopsClient returns a valid Azure DevOps client.
-func NewAzureDevopsClient(hostname string, userName string, token string) (*AzureDevopsClient, error) {
+func NewAzureDevopsClient(hostname string, userName string, token string, timeout time.Duration, honorRetryAfterHeader bool) (*AzureDevopsClient, error) {
 	tp := azuredevops.BasicAuthTransport{
 		Username: "",
 		Password: strings.TrimSpace(token),
 	}
 	httpClient := tp.Client()
-	httpClient.Timeout = time.Second * 10
+	httpClient.Timeout = timeout
+
+	if honorRetryAfterHeader {
+		httpClient.Transport = NewThrottleRoundtripper(httpClient.Transport)
+	}
+
 	var adClient, err = azuredevops.NewClient(httpClient)
 	if err != nil {
 		return nil, err
@@ -68,6 +109,14 @@ func (g *AzureDevopsClient) GetModifiedFiles(repo models.Repo, pull models.PullR
 	sourceRefName := strings.Replace(pullRequest.GetSourceRefName(), "refs/heads/", "", 1)
 
 	r, resp, err := g.Client.Git.GetDiffs(g.ctx, owner, project, repoName, targetRefName, sourceRefName)
+	if err != nil {
+		// There is a bug in the underlying azdo library which does not unmarshal requests correctly
+		// If this PR gets merged, then this bug will go away: https://github.com/mcdafydd/go-azuredevops/pull/12
+		matched, err := regexp.MatchString(`cannot unmarshal number [A-Za-z]+ into Go struct field GitCommitDiffs.changeCounts of type azuredevops.VersionControlChangeType`, err.Error())
+		if !matched {
+			return nil, errors.Wrapf(err, "Error getting diff %s to %s", sourceRefName, targetRefName)
+		}
+	}
 	if resp.StatusCode != http.StatusOK {
 		return nil, errors.Wrapf(err, "http response code %d getting diff %s to %s", resp.StatusCode, sourceRefName, targetRefName)
 	}

server/events/vcs/azuredevops_client_test.go

@@ -9,6 +9,7 @@ import (
 	"net/url"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/mcdafydd/go-azuredevops/azuredevops"
 	"github.com/runatlantis/atlantis/server/events/models"
@@ -19,28 +20,39 @@ import (
 
 func TestAzureDevopsClient_MergePull(t *testing.T) {
 	cases := []struct {
-		description string
-		response    string
-		code        int
-		expErr      string
+		description  string
+		response     string
+		code         int
+		expErr       string
+		throttleTime time.Duration
 	}{
 		{
 			"success",
 			adMergeSuccess,
 			200,
 			"",
+			0 * time.Second,
 		},
 		{
 			"405",
 			`{"message":"405 Method Not Allowed"}`,
 			405,
 			"405 {message: 405 Method Not Allowed}",
+			0 * time.Second,
 		},
 		{
 			"406",
 			`{"message":"406 Branch cannot be merged"}`,
 			406,
 			"406 {message: 406 Branch cannot be merged}",
+			0 * time.Second,
+		},
+		{
+			"success",
+			adMergeSuccess,
+			200,
+			"",
+			1 * time.Second,
 		},
 	}
 
@@ -81,8 +93,21 @@ func TestAzureDevopsClient_MergePull(t *testing.T) {
 
 	for _, c := range cases {
 		t.Run(c.description, func(t *testing.T) {
+			// Use lastquery to keep track of the last request that hit our server
+			var lastQuery time.Time = time.Time{}
+
 			testServer := httptest.NewTLSServer(
 				http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					currentTime := time.Now()
+					if !currentTime.After(lastQuery.Add(c.throttleTime)) {
+						t.Errorf("Did not respect retry-after header from azdo")
+						http.Error(w, "throttle not honored", http.StatusBadRequest)
+						return
+					}
+					lastQuery = currentTime
+					// We always set the retry-after
+					w.Header().Set("retry-after", fmt.Sprint(int64(c.throttleTime/time.Second)))
+
 					switch r.RequestURI {
 					// The first request should hit this URL.
 					case "/owner/project/_apis/git/repositories/repo/pullrequests/22?api-version=5.1-preview.1":
@@ -99,7 +124,7 @@ func TestAzureDevopsClient_MergePull(t *testing.T) {
 
 			testServerURL, err := url.Parse(testServer.URL)
 			Ok(t, err)
-			client, err := vcs.NewAzureDevopsClient(testServerURL.Host, "user", "token")
+			client, err := vcs.NewAzureDevopsClient(testServerURL.Host, "user", "token", 5*time.Second, true)
 			client.Client.VsaexBaseURL = *testServerURL
 			Ok(t, err)
 			defer disableSSLVerification()()
@@ -213,7 +238,7 @@ func TestAzureDevopsClient_UpdateStatus(t *testing.T) {
 
 			testServerURL, err := url.Parse(testServer.URL)
 			Ok(t, err)
-			client, err := vcs.NewAzureDevopsClient(testServerURL.Host, "user", "token")
+			client, err := vcs.NewAzureDevopsClient(testServerURL.Host, "user", "token", 1*time.Second, true)
 			Ok(t, err)
 			defer disableSSLVerification()()
 
@@ -237,6 +262,9 @@ func TestAzureDevopsClient_UpdateStatus(t *testing.T) {
 // and concat results.
 func TestAzureDevopsClient_GetModifiedFiles(t *testing.T) {
 	itemRespTemplate := `{
+		"changeCounts": {
+			"Add": 2
+		},
 		"changes": [
 	{
 		"item": {
@@ -276,7 +304,7 @@ func TestAzureDevopsClient_GetModifiedFiles(t *testing.T) {
 
 	testServerURL, err := url.Parse(testServer.URL)
 	Ok(t, err)
-	client, err := vcs.NewAzureDevopsClient(testServerURL.Host, "user", "token")
+	client, err := vcs.NewAzureDevopsClient(testServerURL.Host, "user", "token", 1*time.Second, true)
 	Ok(t, err)
 	defer disableSSLVerification()()
 
@@ -396,7 +424,7 @@ func TestAzureDevopsClient_PullIsMergeable(t *testing.T) {
 			testServerURL, err := url.Parse(testServer.URL)
 			Ok(t, err)
 
-			client, err := vcs.NewAzureDevopsClient(testServerURL.Host, "user", "token")
+			client, err := vcs.NewAzureDevopsClient(testServerURL.Host, "user", "token", 1*time.Second, true)
 			Ok(t, err)
 
 			defer disableSSLVerification()()
@@ -490,7 +518,7 @@ func TestAzureDevopsClient_PullIsApproved(t *testing.T) {
 			testServerURL, err := url.Parse(testServer.URL)
 			Ok(t, err)
 
-			client, err := vcs.NewAzureDevopsClient(testServerURL.Host, "user", "token")
+			client, err := vcs.NewAzureDevopsClient(testServerURL.Host, "user", "token", 1*time.Second, true)
 			Ok(t, err)
 
 			defer disableSSLVerification()()
@@ -535,7 +563,7 @@ func TestAzureDevopsClient_GetPullRequest(t *testing.T) {
 			}))
 		testServerURL, err := url.Parse(testServer.URL)
 		Ok(t, err)
-		client, err := vcs.NewAzureDevopsClient(testServerURL.Host, "user", "token")
+		client, err := vcs.NewAzureDevopsClient(testServerURL.Host, "user", "token", 1*time.Second, true)
 		Ok(t, err)
 		defer disableSSLVerification()()
 
@@ -555,7 +583,7 @@ func TestAzureDevopsClient_GetPullRequest(t *testing.T) {
 }
 
 func TestAzureDevopsClient_MarkdownPullLink(t *testing.T) {
-	client, err := vcs.NewAzureDevopsClient("hostname", "user", "token")
+	client, err := vcs.NewAzureDevopsClient("hostname", "user", "token", 1*time.Second, true)
 	Ok(t, err)
 	pull := models.PullRequest{Num: 1}
 	s, _ := client.MarkdownPullLink(pull)

server/server.go

@@ -203,9 +203,13 @@ func NewServer(userConfig UserConfig, config Config) (*Server, error) {
 		}
 	}
 	if userConfig.AzureDevopsUser != "" {
+		azureDevopsTimeout := 10 * time.Second
+		if userConfig.AzureDevopsTimeoutSeconds != 0 {
+			azureDevopsTimeout = time.Duration(userConfig.AzureDevopsTimeoutSeconds) * time.Second
+		}
 		supportedVCSHosts = append(supportedVCSHosts, models.AzureDevops)
 		var err error
-		azuredevopsClient, err = vcs.NewAzureDevopsClient("dev.azure.com", userConfig.AzureDevopsUser, userConfig.AzureDevopsToken)
+		azuredevopsClient, err = vcs.NewAzureDevopsClient("dev.azure.com", userConfig.AzureDevopsUser, userConfig.AzureDevopsToken, azureDevopsTimeout, userConfig.AzureDevopsHonorRetryAfter)
 		if err != nil {
 			return nil, err
 		}

server/user_config.go

@@ -13,6 +13,8 @@ type UserConfig struct {
 	AtlantisURL                string `mapstructure:"atlantis-url"`
 	Automerge                  bool   `mapstructure:"automerge"`
 	AutoplanFileList           string `mapstructure:"autoplan-file-list"`
+	AzureDevopsHonorRetryAfter bool   `mapstructure:"azuredevops-honor-retry-after"`
+	AzureDevopsTimeoutSeconds  int    `mapstructure:"azuredevops-timeout-seconds"`
 	AzureDevopsToken           string `mapstructure:"azuredevops-token"`
 	AzureDevopsUser            string `mapstructure:"azuredevops-user"`
 	AzureDevopsWebhookPassword string `mapstructure:"azuredevops-webhook-password"`