fix: retry when encounting the GitHub rate limit

internal/scm/github/github.go

@@ -290,26 +290,27 @@ func (g *Github) CreatePullRequest(ctx context.Context, repo scm.Repository, prR
 
 func (g *Github) createPullRequest(ctx context.Context, repo repository, prRepo repository, newPR scm.NewPullRequest) (*github.PullRequest, error) {
 	head := fmt.Sprintf("%s:%s", prRepo.ownerName, newPR.Head)
-	pr, _, err := g.ghClient.PullRequests.Create(ctx, repo.ownerName, repo.name, &github.NewPullRequest{
-		Title: &newPR.Title,
-		Body:  &newPR.Body,
-		Head:  &head,
-		Base:  &newPR.Base,
-		Draft: &newPR.Draft,
-	})
-	if err != nil {
-		return nil, err
-	}
 
-	return pr, nil
+	pr, _, err := retry(ctx, func() (*github.PullRequest, *github.Response, error) {
+		return g.ghClient.PullRequests.Create(ctx, repo.ownerName, repo.name, &github.NewPullRequest{
+			Title: &newPR.Title,
+			Body:  &newPR.Body,
+			Head:  &head,
+			Base:  &newPR.Base,
+			Draft: &newPR.Draft,
+		})
+	})
+	return pr, err
 }
 
 func (g *Github) addReviewers(ctx context.Context, repo repository, newPR scm.NewPullRequest, createdPR *github.PullRequest) error {
 	if len(newPR.Reviewers) == 0 {
 		return nil
 	}
-	_, _, err := g.ghClient.PullRequests.RequestReviewers(ctx, repo.ownerName, repo.name, createdPR.GetNumber(), github.ReviewersRequest{
-		Reviewers: newPR.Reviewers,
+	_, _, err := retry(ctx, func() (*github.PullRequest, *github.Response, error) {
+		return g.ghClient.PullRequests.RequestReviewers(ctx, repo.ownerName, repo.name, createdPR.GetNumber(), github.ReviewersRequest{
+			Reviewers: newPR.Reviewers,
+		})
 	})
 	return err
 }
@@ -318,7 +319,9 @@ func (g *Github) addAssignees(ctx context.Context, repo repository, newPR scm.Ne
 	if len(newPR.Assignees) == 0 {
 		return nil
 	}
-	_, _, err := g.ghClient.Issues.AddAssignees(ctx, repo.ownerName, repo.name, createdPR.GetNumber(), newPR.Assignees)
+	_, _, err := retry(ctx, func() (*github.Issue, *github.Response, error) {
+		return g.ghClient.Issues.AddAssignees(ctx, repo.ownerName, repo.name, createdPR.GetNumber(), newPR.Assignees)
+	})
 	return err
 }
 
@@ -504,14 +507,18 @@ func (g *Github) MergePullRequest(ctx context.Context, pullReq scm.PullRequest)
 		return errors.New("none of the configured merge types was permitted")
 	}
 
-	_, _, err = g.ghClient.PullRequests.Merge(ctx, pr.ownerName, pr.repoName, pr.number, "", &github.PullRequestOptions{
-		MergeMethod: mergeTypeGhName[mergeTypes[0]],
+	_, _, err = retry(ctx, func() (*github.PullRequestMergeResult, *github.Response, error) {
+		return g.ghClient.PullRequests.Merge(ctx, pr.ownerName, pr.repoName, pr.number, "", &github.PullRequestOptions{
+			MergeMethod: mergeTypeGhName[mergeTypes[0]],
+		})
 	})
 	if err != nil {
 		return err
 	}
 
-	_, err = g.ghClient.Git.DeleteRef(ctx, pr.prOwnerName, pr.prRepoName, fmt.Sprintf("heads/%s", pr.branchName))
+	_, err = retryWithoutReturn(ctx, func() (*github.Response, error) {
+		return g.ghClient.Git.DeleteRef(ctx, pr.prOwnerName, pr.prRepoName, fmt.Sprintf("heads/%s", pr.branchName))
+	})
 
 	// Ignore errors about the reference not existing since it may be the case that GitHub has already deleted the branch
 	if err != nil && !strings.Contains(err.Error(), "Reference does not exist") {
@@ -546,9 +553,12 @@ func (g *Github) ForkRepository(ctx context.Context, repo scm.Repository, newOwn
 	g.modLock()
 	defer g.modUnlock()
 
-	createdRepo, _, err := g.ghClient.Repositories.CreateFork(ctx, r.ownerName, r.name, &github.RepositoryCreateForkOptions{
-		Organization: newOwner,
+	createdRepo, _, err := retry(ctx, func() (*github.Repository, *github.Response, error) {
+		return g.ghClient.Repositories.CreateFork(ctx, r.ownerName, r.name, &github.RepositoryCreateForkOptions{
+			Organization: newOwner,
+		})
 	})
+
 	if err != nil {
 		if _, isAccepted := err.(*github.AcceptedError); !isAccepted {
 			return nil, err

internal/scm/github/retry.go

@@ -0,0 +1,81 @@
+package github
+
+import (
+	"context"
+	"math"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/google/go-github/v47/github"
+	"github.com/pkg/errors"
+	log "github.com/sirupsen/logrus"
+)
+
+const retryHeader = "Retry-After"
+
+var sleep = func(ctx context.Context, d time.Duration) error {
+	log.Infof("Hit rate limit, sleeping for %s", d)
+	select {
+	case <-ctx.Done():
+		return errors.New("aborted while waiting for rate-limit")
+	case <-time.After(d):
+		return nil
+	}
+}
+
+// retry runs a GitHub API request and retries it if a temporary error occurred
+func retry[K any](ctx context.Context, fn func() (K, *github.Response, error)) (K, *github.Response, error) {
+	var val K
+	resp, err := retryWithoutReturn(ctx, func() (*github.Response, error) {
+		var resp *github.Response
+		var err error
+		val, resp, err = fn()
+		return resp, err
+	})
+	return val, resp, err
+}
+
+// retryWithoutReturn runs a GitHub API request with no return value and retries it if a temporary error occurred
+func retryWithoutReturn(ctx context.Context, fn func() (*github.Response, error)) (*github.Response, error) {
+	tries := 0
+
+	for {
+		tries++
+
+		githubResp, err := fn()
+		if err == nil { // NB!
+			return githubResp, nil
+		}
+
+		// Get the number of retry seconds (if any)
+		retryAfter := 0
+		if githubResp != nil && githubResp.Header != nil {
+			retryAfterStr := githubResp.Header.Get(retryHeader)
+			if retryAfterStr != "" {
+				var err error
+				if retryAfter, err = strconv.Atoi(retryAfterStr); err != nil {
+					return githubResp, errors.WithMessage(err, "could not convert Retry-After header")
+				}
+			}
+		}
+
+		switch {
+		// If GitHub has specified how long we should wait, use that information
+		case retryAfter != 0:
+			err := sleep(ctx, time.Duration(retryAfter)*time.Second)
+			if err != nil {
+				return githubResp, err
+			}
+		// If secondary rate limit error, use an exponential back-off to determine the wait
+		case strings.Contains(err.Error(), "secondary rate limit"):
+			err := sleep(ctx, time.Duration(math.Pow(float64(tries), 3))*10*time.Second)
+			if err != nil {
+				return githubResp, err
+			}
+		// If any other error, return the error
+		default:
+			return githubResp, err
+		}
+	}
+}

internal/scm/github/retry_test.go

@@ -0,0 +1,199 @@
+package github
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/google/go-github/v47/github"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_retryWithoutReturn(t *testing.T) {
+	tests := []struct {
+		name      string
+		responses []response
+		wantErr   bool
+		sleep     time.Duration
+	}{
+		{
+			name: "one fail",
+			responses: []response{
+				secondaryRateLimitError,
+				okResponse,
+			},
+			sleep: 10 * time.Second,
+		},
+		{
+			name: "two fails",
+			responses: []response{
+				secondaryRateLimitError,
+				secondaryRateLimitError,
+				okResponse,
+			},
+			sleep: 1*time.Minute + 30*time.Second,
+		},
+		{
+			name: "three fails",
+			responses: []response{
+				secondaryRateLimitError,
+				secondaryRateLimitError,
+				secondaryRateLimitError,
+				okResponse,
+			},
+			sleep: 6*time.Minute + 0*time.Second,
+		},
+		{
+			name: "four fails",
+			responses: []response{
+				secondaryRateLimitError,
+				secondaryRateLimitError,
+				secondaryRateLimitError,
+				secondaryRateLimitError,
+				okResponse,
+			},
+			sleep: 16*time.Minute + 40*time.Second,
+		},
+		{
+			name: "a real fail",
+			responses: []response{
+				realError,
+			},
+			sleep:   0,
+			wantErr: true,
+		},
+		{
+			name: "two temporary fails and a real one",
+			responses: []response{
+				secondaryRateLimitError,
+				secondaryRateLimitError,
+				realError,
+			},
+			sleep:   1*time.Minute + 30*time.Second,
+			wantErr: true,
+		},
+		{
+			name: "primary rate limit",
+			responses: []response{
+				primaryRateLimitError,
+				okResponse,
+			},
+			sleep:   1*time.Minute + 37*time.Second,
+			wantErr: false,
+		},
+		{
+			name: "two primary rate limit errors in a row",
+			responses: []response{
+				primaryRateLimitError,
+				primaryRateLimitError,
+				okResponse,
+			},
+			sleep:   3*time.Minute + 14*time.Second,
+			wantErr: false,
+		},
+		{
+			name: "secondary rate limit followed by primary rate limit",
+			responses: []response{
+				secondaryRateLimitError,
+				primaryRateLimitError,
+				okResponse,
+			},
+			sleep:   1*time.Minute + 47*time.Second,
+			wantErr: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var slept time.Duration
+			sleep = func(ctx context.Context, d time.Duration) error {
+				slept += d
+				return nil
+			}
+
+			call := 0
+			fn := func() (*github.Response, error) {
+				resp := tt.responses[call]
+				call++
+				return resp.response, resp.err
+			}
+
+			if _, err := retryWithoutReturn(context.Background(), fn); (err != nil) != tt.wantErr {
+				t.Errorf("retry() error = %v, wantErr %v", err, tt.wantErr)
+			}
+
+			assert.Equal(t, tt.sleep, slept)
+		})
+	}
+}
+
+func Test_retry(t *testing.T) {
+	var slept time.Duration
+	sleep = func(ctx context.Context, d time.Duration) error {
+		slept += d
+		return nil
+	}
+
+	call := 0
+	fn := func() (*github.PullRequest, *github.Response, error) {
+		call++
+		if call == 4 {
+			return &github.PullRequest{
+				ID: &[]int64{100}[0],
+			}, okResponse.response, nil
+		}
+		return nil, secondaryRateLimitError.response, secondaryRateLimitError.err
+	}
+
+	pr, resp, err := retry(context.Background(), fn)
+	assert.Equal(t, int64(100), *pr.ID)
+	assert.Equal(t, 200, resp.StatusCode)
+	assert.NoError(t, err)
+}
+
+type response struct {
+	response *github.Response
+	err      error
+}
+
+var okResponse = response{
+	err: nil,
+	response: &github.Response{
+		Response: &http.Response{
+			StatusCode: 200,
+		},
+	},
+}
+var realError = createResponse(http.StatusNotFound, errors.New("something went wrong"))
+var secondaryErrorMsg = "You have exceeded a secondary rate limit and have been temporarily blocked from content creation. Please retry your request again later."
+var secondaryRateLimitError = createResponse(http.StatusForbidden, errors.New(secondaryErrorMsg))
+var primaryRateLimitError = response{
+	response: &github.Response{
+		Response: &http.Response{
+			StatusCode: http.StatusForbidden,
+			Header: http.Header{
+				"Retry-After": []string{"97"},
+			},
+		},
+	},
+	err: errors.New("rate limited"),
+}
+
+func createResponse(statusCode int, err error) response {
+	return response{
+		response: createGithubResponse(statusCode, err.Error()),
+		err:      err,
+	}
+}
+
+func createGithubResponse(statusCode int, errorMsg string) *github.Response {
+	return &github.Response{
+		Response: &http.Response{
+			StatusCode: statusCode,
+			Body:       io.NopCloser(strings.NewReader(errorMsg)),
+		},
+	}
+}