From 8256d9781e5fbb09ae18da3bd0b79d0a1a34ed45 Mon Sep 17 00:00:00 2001
From: Craig Miskell <cmiskell@gitlab.com>
Date: Fri, 8 Oct 2021 14:19:47 +1300
Subject: [PATCH] Respect HTTP 429 and Retry-After header

When a 429 is received, sleep for the length of time indicated by Retry-After, then try again.

Try up to 3 times before giving up, although the count can be configured with :ratelimit_retries if you want more or less

If no Retry-After is sent by the server, defaults to 2 seconds (something is better than nothing). That should only happen with quite old and unsupported versions of GitLab
---
 lib/gitlab/request.rb       | 13 ++++++-
 spec/gitlab/request_spec.rb | 73 +++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/lib/gitlab/request.rb b/lib/gitlab/request.rb
index 59b74c32a..ff6958e35 100644
--- a/lib/gitlab/request.rb
+++ b/lib/gitlab/request.rb
@@ -48,7 +48,18 @@ def self.decode(response)
           params[:headers].merge!(authorization_header)
         end
 
-        validate self.class.send(method, endpoint + path, params)
+        retries_left = params[:ratelimit_retries] || 3
+        begin
+          response = self.class.send(method, endpoint + path, params)
+          validate response
+        rescue Gitlab::Error::TooManyRequests => e
+          retries_left -= 1
+          raise e if retries_left.zero?
+
+          wait_time = response.headers['Retry-After'] || 2
+          sleep(wait_time.to_i)
+          retry
+        end
       end
     end
 
diff --git a/spec/gitlab/request_spec.rb b/spec/gitlab/request_spec.rb
index 0c2fabb77..4f196c891 100644
--- a/spec/gitlab/request_spec.rb
+++ b/spec/gitlab/request_spec.rb
@@ -114,4 +114,77 @@
       expect(@request.send(:authorization_header)).to eq('Authorization' => 'Bearer 3225e2804d31fea13fc41fc83bffef00cfaedc463118646b154acc6f94747603')
     end
   end
+
+  describe 'ratelimiting' do
+    before do
+      @request.private_token = 'token'
+      @request.endpoint = 'https://example.com/api/v4'
+      @rpath = "#{@request.endpoint}/version"
+
+      allow(@request).to receive(:httparty)
+    end
+
+    it 'tries 3 times when ratelimited by default' do
+      stub_request(:get, @rpath)
+        .to_return(
+          status: 429,
+          headers: { 'Retry-After' => 1 }
+        )
+
+      expect do
+        @request.get('/version')
+      end.to raise_error(Gitlab::Error::TooManyRequests)
+
+      expect(a_request(:get, @rpath).with(headers: {
+        'PRIVATE_TOKEN' => 'token'
+      }.merge(described_class.headers))).to have_been_made.times(3)
+    end
+
+    it 'tries 4 times when ratelimited with option' do
+      stub_request(:get, @rpath)
+        .to_return(
+          status: 429,
+          headers: { 'Retry-After' => 1 }
+        )
+      expect do
+        @request.get('/version', { ratelimit_retries: 4 })
+      end.to raise_error(Gitlab::Error::TooManyRequests)
+
+      expect(a_request(:get, @rpath).with(headers: {
+        'PRIVATE_TOKEN' => 'token'
+      }.merge(described_class.headers))).to have_been_made.times(4)
+    end
+
+    it 'handles one retry then success' do
+      stub_request(:get, @rpath)
+        .to_return(
+          status: 429,
+          headers: { 'Retry-After' => 1 }
+        ).times(1).then
+        .to_return(
+          status: 200
+        ).times(1)
+
+      @request.get('/version')
+
+      expect(a_request(:get, @rpath).with(headers: {
+        'PRIVATE_TOKEN' => 'token'
+      }.merge(described_class.headers))).to have_been_made.times(2)
+    end
+
+    it 'survives a 429 with no Retry-After header' do
+      stub_request(:get, @rpath)
+        .to_return(
+          status: 429
+        )
+
+      expect do
+        @request.get('/version')
+      end.to raise_error(Gitlab::Error::TooManyRequests)
+
+      expect(a_request(:get, @rpath).with(headers: {
+        'PRIVATE_TOKEN' => 'token'
+      }.merge(described_class.headers))).to have_been_made.times(3)
+    end
+  end
 end