From fbab0bae094d3675c2225920d994e413e8d1159a Mon Sep 17 00:00:00 2001 From: Itay Date: Mon, 16 Sep 2024 08:07:23 +0000 Subject: [PATCH 1/3] Add rate limit handling for GitHub client --- go.mod | 1 + go.sum | 2 + server/events/vcs/github_client.go | 12 ++++-- server/events/vcs/github_client_test.go | 55 +++++++++++++++++++++++++ 4 files changed, 67 insertions(+), 3 deletions(-) diff --git a/go.mod b/go.mod index 3e5f0f3f2a..8b4ea324a7 100644 --- a/go.mod +++ b/go.mod @@ -83,6 +83,7 @@ require ( github.com/go-fed/httpsig v1.1.0 // indirect github.com/go-playground/locales v0.14.1 // indirect github.com/go-playground/universal-translator v0.18.1 // indirect + github.com/gofri/go-github-ratelimit v1.1.0 // indirect github.com/golang-jwt/jwt/v4 v4.5.0 // indirect github.com/golang/protobuf v1.5.3 // indirect github.com/google/go-cmp v0.6.0 // indirect diff --git a/go.sum b/go.sum index c679f45898..d8ad927d8e 100644 --- a/go.sum +++ b/go.sum @@ -156,6 +156,8 @@ github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4 github.com/go-test/deep v1.0.4/go.mod h1:wGDj63lr65AM2AQyKZd/NYHGb0R+1RLqB8NKt3aSFNA= github.com/go-test/deep v1.1.1 h1:0r/53hagsehfO4bzD2Pgr/+RgHqhmf+k1Bpse2cTu1U= github.com/go-test/deep v1.1.1/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE= +github.com/gofri/go-github-ratelimit v1.1.0 h1:ijQ2bcv5pjZXNil5FiwglCg8wc9s8EgjTmNkqjw8nuk= +github.com/gofri/go-github-ratelimit v1.1.0/go.mod h1:OnCi5gV+hAG/LMR7llGhU7yHt44se9sYgKPnafoL7RY= github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ= github.com/golang-jwt/jwt/v4 v4.5.0 h1:7cYmW1XlMY7h7ii7UhUyChSgS5wUJEnm9uZVTGqOWzg= github.com/golang-jwt/jwt/v4 v4.5.0/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0= diff --git a/server/events/vcs/github_client.go b/server/events/vcs/github_client.go index 040a85958b..2fd7576963 100644 --- a/server/events/vcs/github_client.go +++ b/server/events/vcs/github_client.go @@ -22,6 +22,7 @@ import ( "strings" "time" + "github.com/gofri/go-github-ratelimit/github_ratelimit" "github.com/google/go-github/v59/github" "github.com/pkg/errors" "github.com/runatlantis/atlantis/server/events/command" @@ -121,15 +122,20 @@ func NewGithubClient(hostname string, credentials GithubCredentials, config Gith return nil, errors.Wrap(err, "error initializing github authentication transport") } + transportWithRateLimit, err := github_ratelimit.NewRateLimitWaiterClient(transport.Transport, github_ratelimit.WithTotalSleepLimit(time.Minute, nil)) + if err != nil { + return nil, errors.Wrap(err, "error initializing github rate limit transport") + } + var graphqlURL string var client *github.Client if hostname == "github.com" { - client = github.NewClient(transport) + client = github.NewClient(transportWithRateLimit) graphqlURL = "https://api.github.com/graphql" } else { apiURL := resolveGithubAPIURL(hostname) // TODO: Deprecated: Use NewClient(httpClient).WithEnterpriseURLs(baseURL, uploadURL) instead - client, err = github.NewEnterpriseClient(apiURL.String(), apiURL.String(), transport) //nolint:staticcheck + client, err = github.NewEnterpriseClient(apiURL.String(), apiURL.String(), transportWithRateLimit) //nolint:staticcheck if err != nil { return nil, err } @@ -137,7 +143,7 @@ func NewGithubClient(hostname string, credentials GithubCredentials, config Gith } // Use the client from shurcooL's githubv4 library for queries. - v4Client := githubv4.NewEnterpriseClient(graphqlURL, transport) + v4Client := githubv4.NewEnterpriseClient(graphqlURL, transportWithRateLimit) user, err := credentials.GetUser() logger.Debug("GH User: %s", user) diff --git a/server/events/vcs/github_client_test.go b/server/events/vcs/github_client_test.go index ef97c64245..f7769836e5 100644 --- a/server/events/vcs/github_client_test.go +++ b/server/events/vcs/github_client_test.go @@ -11,6 +11,7 @@ import ( "os" "strings" "testing" + "time" "github.com/runatlantis/atlantis/server/events/command" "github.com/runatlantis/atlantis/server/events/models" @@ -1628,3 +1629,57 @@ func TestGithubClient_GetPullLabels_EmptyResponse(t *testing.T) { Ok(t, err) Equals(t, 0, len(labels)) } + +func TestGithubClient_SecondaryRateLimitHandling_CreateComment(t *testing.T) { + logger := logging.NewNoopLogger(t) + calls := 0 + maxCalls := 2 + + testServer := httptest.NewTLSServer( + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost || r.URL.Path != "/api/v3/repos/owner/repo/issues/1/comments" { + t.Errorf("Unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + return + } + + if calls < maxCalls { + // Secondary rate limiting, x-ratelimit-remaining must be > 0 + w.Header().Set("x-ratelimit-remaining", "1") + w.Header().Set("x-ratelimit-reset", fmt.Sprintf("%d", time.Now().Unix()+1)) + w.WriteHeader(http.StatusForbidden) + json.NewEncoder(w).Encode(map[string]string{"message": "You have exceeded a secondary rate limit"}) + } else { + w.WriteHeader(http.StatusCreated) + json.NewEncoder(w).Encode(map[string]interface{}{ + "id": 1, + "body": "Test comment", + }) + } + calls++ + }), + ) + + testServerURL, err := url.Parse(testServer.URL) + Ok(t, err) + + client, err := vcs.NewGithubClient(testServerURL.Host, &vcs.GithubUserCredentials{"user", "pass"}, vcs.GithubConfig{}, 0, logger) + Ok(t, err) + defer disableSSLVerification()() + + // Simulate creating a comment + repo := models.Repo{ + FullName: "owner/repo", + Owner: "owner", + Name: "repo", + } + pullNum := 1 + comment := "Test comment" + + err = client.CreateComment(logger, repo, pullNum, comment, "") + Ok(t, err) + + // Verify that the number of calls is greater than maxCalls, indicating that retries occurred + Assert(t, calls > maxCalls, "Expected more than %d calls due to rate limiting, but got %d", maxCalls, calls) + +} From f064f7bda21987843c99b3f3a73d1866369173e8 Mon Sep 17 00:00:00 2001 From: Itay Date: Mon, 16 Sep 2024 15:56:52 +0000 Subject: [PATCH 2/3] Add go-github-ratelimit package and update response handling in GitHub client test --- go.mod | 2 +- server/events/vcs/github_client_test.go | 7 ++----- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/go.mod b/go.mod index 8b4ea324a7..e36d4f66ad 100644 --- a/go.mod +++ b/go.mod @@ -12,6 +12,7 @@ require ( github.com/go-ozzo/ozzo-validation v3.6.0+incompatible github.com/go-playground/validator/v10 v10.22.0 github.com/go-test/deep v1.1.1 + github.com/gofri/go-github-ratelimit v1.1.0 github.com/golang-jwt/jwt/v5 v5.2.1 github.com/google/go-github/v59 v59.0.0 github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 @@ -83,7 +84,6 @@ require ( github.com/go-fed/httpsig v1.1.0 // indirect github.com/go-playground/locales v0.14.1 // indirect github.com/go-playground/universal-translator v0.18.1 // indirect - github.com/gofri/go-github-ratelimit v1.1.0 // indirect github.com/golang-jwt/jwt/v4 v4.5.0 // indirect github.com/golang/protobuf v1.5.3 // indirect github.com/google/go-cmp v0.6.0 // indirect diff --git a/server/events/vcs/github_client_test.go b/server/events/vcs/github_client_test.go index f7769836e5..e2ab00dd88 100644 --- a/server/events/vcs/github_client_test.go +++ b/server/events/vcs/github_client_test.go @@ -1648,13 +1648,10 @@ func TestGithubClient_SecondaryRateLimitHandling_CreateComment(t *testing.T) { w.Header().Set("x-ratelimit-remaining", "1") w.Header().Set("x-ratelimit-reset", fmt.Sprintf("%d", time.Now().Unix()+1)) w.WriteHeader(http.StatusForbidden) - json.NewEncoder(w).Encode(map[string]string{"message": "You have exceeded a secondary rate limit"}) + w.Write([]byte(`{"message": "You have exceeded a secondary rate limit"}`)) } else { w.WriteHeader(http.StatusCreated) - json.NewEncoder(w).Encode(map[string]interface{}{ - "id": 1, - "body": "Test comment", - }) + w.Write([]byte(`{"id": 1, "body": "Test comment"}`)) } calls++ }), From e409712a824baed7102866d6b4fc7c4f298bac4b Mon Sep 17 00:00:00 2001 From: Itay Date: Sun, 22 Sep 2024 07:41:14 +0000 Subject: [PATCH 3/3] fix lint check --- server/events/vcs/github_client_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/server/events/vcs/github_client_test.go b/server/events/vcs/github_client_test.go index e2ab00dd88..ab8481dfba 100644 --- a/server/events/vcs/github_client_test.go +++ b/server/events/vcs/github_client_test.go @@ -1648,10 +1648,10 @@ func TestGithubClient_SecondaryRateLimitHandling_CreateComment(t *testing.T) { w.Header().Set("x-ratelimit-remaining", "1") w.Header().Set("x-ratelimit-reset", fmt.Sprintf("%d", time.Now().Unix()+1)) w.WriteHeader(http.StatusForbidden) - w.Write([]byte(`{"message": "You have exceeded a secondary rate limit"}`)) + w.Write([]byte(`{"message": "You have exceeded a secondary rate limit"}`)) // nolint: errcheck } else { w.WriteHeader(http.StatusCreated) - w.Write([]byte(`{"id": 1, "body": "Test comment"}`)) + w.Write([]byte(`{"id": 1, "body": "Test comment"}`)) // nolint: errcheck } calls++ }),