Add retry support (#1502)

* Base design * add_test * adjust test * handle nil * Fix tests * add testing sleep time without actually sleeping * Add documentation * change code generator * pre-commit fixes * change how configs are made * pre-commit fixes * update readme doc * pre-commit fixes * adds validation for backoff_base * pre-commit fixes * change generator * pre-commit fixes * add validation back * pre-commit fixes --------- Co-authored-by: Thomas Hervé <[email protected]> Co-authored-by: ci.datadog-api-spec <[email protected]>
DataDog · Aug 23, 2023 · 0af80ad · 0af80ad
1 parent 803beb2
commit 0af80ad
Show file tree

Hide file tree

Showing 7 changed files with 288 additions and 106 deletions.
diff --git a/.generator/src/generator/templates/api_client.j2 b/.generator/src/generator/templates/api_client.j2
@@ -44,71 +44,99 @@ module {{ module_name }}
     #   the data deserialized from response body (could be nil), response status code and response headers.
     def call_api(http_method, path, opts = {})
       request = build_request(http_method, path, opts)
-      if opts[:stream_body]
-        tempfile = nil
-        encoding = nil
-
-        response = request.perform do | chunk |
-          unless tempfile
-            content_disposition = chunk.http_response.header['Content-Disposition']
-            if content_disposition && content_disposition =~ /filename=/i
-              filename = content_disposition[/filename=['"]?([^'"\s]+)['"]?/, 1]
-              prefix = sanitize_filename(filename)
-            else
-              prefix = 'download-'
-            end
-            prefix = prefix + '-' unless prefix.end_with?('-')
-            unless encoding
-              encoding = chunk.encoding
+      attempt = 0
+      loop do
+        if opts[:stream_body]
+          tempfile = nil
+          encoding = nil
+
+          response = request.perform do | chunk |
+            unless tempfile
+              content_disposition = chunk.http_response.header['Content-Disposition']
+              if content_disposition && content_disposition =~ /filename=/i
+                filename = content_disposition[/filename=['"]?([^'"\s]+)['"]?/, 1]
+                prefix = sanitize_filename(filename)
+              else
+                prefix = 'download-'
+              end
+              prefix = prefix + '-' unless prefix.end_with?('-')
+              unless encoding
+                encoding = chunk.encoding
+              end
+              tempfile = Tempfile.open(prefix, @config.temp_folder_path, encoding: encoding)
+              @tempfile = tempfile
             end
-            tempfile = Tempfile.open(prefix, @config.temp_folder_path, encoding: encoding)
-            @tempfile = tempfile
+            chunk.force_encoding(encoding)
+            tempfile.write(chunk)
+          end
+          if tempfile
+            tempfile.close
+            @config.logger.info "Temp file written to #{tempfile.path}, please copy the file to a proper folder "\
+                              "with e.g. `FileUtils.cp(tempfile.path, '/new/file/path')` otherwise the temp file "\
+                              "will be deleted automatically with GC. It's also recommended to delete the temp file "\
+                              "explicitly with `tempfile.delete`"
           end
-          chunk.force_encoding(encoding)
-          tempfile.write(chunk)
+        else
+          response = request.perform
         end
-        if tempfile
-          tempfile.close
-          @config.logger.info "Temp file written to #{tempfile.path}, please copy the file to a proper folder "\
-                            "with e.g. `FileUtils.cp(tempfile.path, '/new/file/path')` otherwise the temp file "\
-                            "will be deleted automatically with GC. It's also recommended to delete the temp file "\
-                            "explicitly with `tempfile.delete`"
+
+        if @config.debugging
+          @config.logger.debug "HTTP response body ~BEGIN~\n#{response.body}\n~END~\n"
         end
-      else
-        response = request.perform
-      end
 
-      if @config.debugging
-        @config.logger.debug "HTTP response body ~BEGIN~\n#{response.body}\n~END~\n"
-      end
+        unless response.success?
+          if response.request_timeout?
+            fail APIError.new('Connection timed out')
+          elsif response.code == 0
+            # Errors from libcurl will be made visible here
+            fail APIError.new(:code => 0,
+                              :message => response.return_message)
+          else
+            body = response.body
+            if response.headers['Content-Encoding'].eql?('gzip') && !(body.nil? || body.empty?) then
+              gzip = Zlib::Inflate.new(Zlib::MAX_WBITS + 16)
+              body = gzip.inflate(body)
+              gzip.close
+            end
+            if should_retry(attempt, @config.max_retries, response.code, @config.enable_retry)
+              sleep calculate_retry_interval(response, @config.backoff_base, @config.backoff_multiplier, attempt, @config.timeout)
+              attempt = attempt + 1
+              next
+            else
+              fail APIError.new(:code => response.code,
+                :response_headers => response.headers,
+                :response_body => body),
+              response.message
+            end
+          end
+        end
 
-      unless response.success?
-        if response.request_timeout?
-          fail APIError.new('Connection timed out')
-        elsif response.code == 0
-          # Errors from libcurl will be made visible here
-          fail APIError.new(:code => 0,
-                            :message => response.return_message)
+        if opts[:return_type]
+          data = deserialize(opts[:api_version], response, opts[:return_type])
         else
-          body = response.body
-          if response.headers['Content-Encoding'].eql?('gzip') && !(body.nil? || body.empty?) then
-            gzip = Zlib::Inflate.new(Zlib::MAX_WBITS + 16)
-            body = gzip.inflate(body)
-            gzip.close
-          end
-          fail APIError.new(:code => response.code,
-                            :response_headers => response.headers,
-                            :response_body => body),
-               response.message
+          data = nil
         end
+        return data, response.code, response.headers
       end
+    end
+
+    # Check if an http request should be retried
+    def should_retry(attempt, max_retries, http_code, enable_retry)
+      (http_code == 429 || http_code >= 500) && max_retries > attempt && enable_retry
+    end
 
-      if opts[:return_type]
-        data = deserialize(opts[:api_version], response, opts[:return_type])
+    # Calculate the sleep interval between 2 retry attempts
+    def calculate_retry_interval(response, backoff_base, backoff_multiplier, attempt, timeout)
+      reset_header = response.headers['X-Ratelimit-Reset']
+      if  !reset_header.nil? && !reset_header.empty?
+        sleep_time = reset_header.to_i
       else
-        data = nil
+        sleep_time = (backoff_multiplier**attempt) * backoff_base
+        if timeout && timeout > 0
+          sleep_time = [timeout, sleep_time].min
+        end
       end
-      return data, response.code, response.headers
+      sleep_time
     end
 
     # Build the HTTP request

diff --git a/.generator/src/generator/templates/configuration.j2 b/.generator/src/generator/templates/configuration.j2
@@ -138,6 +138,16 @@ module {{ module_name }}
     # Password for proxy server authentication
     attr_accessor :http_proxypass
 
+    # Enable retry when rate limited
+    attr_accessor :enable_retry
+
+    # Retry backoff calculation parameters
+    attr_accessor :backoff_base
+    attr_accessor :backoff_multiplier
+
+    # Maximum number of retry attempts allowed
+    attr_accessor :max_retries
+
     def initialize
       {%- set default_server = openapi.servers[0]|format_server %}
       @scheme = '{{ default_server.scheme }}'
@@ -149,6 +159,10 @@ module {{ module_name }}
       @server_operation_variables = {}
       @api_key = {}
       @api_key_prefix = {}
+      @enable_retry = false
+      @backoff_base = 2
+      @backoff_multiplier = 2
+      @max_retries = 3
       @timeout = nil
       @client_side_validation = true
       @verify_ssl = true
@@ -188,6 +202,13 @@ module {{ module_name }}
       @@default ||= Configuration.new
     end
 
+    def backoff_base=(value)
+      if value < 2
+        raise ArgumentError, 'backoff_base cannot be smaller than 2'
+      end
+      @backoff_base = value
+    end
+
     def configure
       yield(self) if block_given?
     end

diff --git a/README.md b/README.md
@@ -165,6 +165,29 @@ api_instance.list_incidents_with_pagination() do |incident|
 end
 ```
 
+### Retry
+
+To enable the client to retry when rate limited (status 429) or status 500 and above:
+
+```ruby
+config = DatadogAPIClient::Configuration.new
+config.enable_retry = true
+client = DatadogAPIClient::APIClient.new(config)
+```
+
+The interval between 2 retry attempts will be the value of the `x-ratelimit-reset` response header when available.
+If not, it will be :
+
+```ruby
+(config.backoffMultiplier ** current_retry_count) * config.backoffBase
+```
+
+The maximum number of retry attempts is `3` by default and can be modified with
+
+```ruby
+config.maxRetries
+```
+
 ## Documentation
 
 If you are interested in general documentation for all public Datadog API endpoints, checkout the [general documentation site][api docs].

diff --git a/lib/datadog_api_client/api_client.rb b/lib/datadog_api_client/api_client.rb
@@ -55,71 +55,99 @@ def self.default
     #   the data deserialized from response body (could be nil), response status code and response headers.
     def call_api(http_method, path, opts = {})
       request = build_request(http_method, path, opts)
-      if opts[:stream_body]
-        tempfile = nil
-        encoding = nil
-
-        response = request.perform do | chunk |
-          unless tempfile
-            content_disposition = chunk.http_response.header['Content-Disposition']
-            if content_disposition && content_disposition =~ /filename=/i
-              filename = content_disposition[/filename=['"]?([^'"\s]+)['"]?/, 1]
-              prefix = sanitize_filename(filename)
-            else
-              prefix = 'download-'
-            end
-            prefix = prefix + '-' unless prefix.end_with?('-')
-            unless encoding
-              encoding = chunk.encoding
+      attempt = 0
+      loop do
+        if opts[:stream_body]
+          tempfile = nil
+          encoding = nil
+
+          response = request.perform do | chunk |
+            unless tempfile
+              content_disposition = chunk.http_response.header['Content-Disposition']
+              if content_disposition && content_disposition =~ /filename=/i
+                filename = content_disposition[/filename=['"]?([^'"\s]+)['"]?/, 1]
+                prefix = sanitize_filename(filename)
+              else
+                prefix = 'download-'
+              end
+              prefix = prefix + '-' unless prefix.end_with?('-')
+              unless encoding
+                encoding = chunk.encoding
+              end
+              tempfile = Tempfile.open(prefix, @config.temp_folder_path, encoding: encoding)
+              @tempfile = tempfile
             end
-            tempfile = Tempfile.open(prefix, @config.temp_folder_path, encoding: encoding)
-            @tempfile = tempfile
+            chunk.force_encoding(encoding)
+            tempfile.write(chunk)
+          end
+          if tempfile
+            tempfile.close
+            @config.logger.info "Temp file written to #{tempfile.path}, please copy the file to a proper folder "\
+                              "with e.g. `FileUtils.cp(tempfile.path, '/new/file/path')` otherwise the temp file "\
+                              "will be deleted automatically with GC. It's also recommended to delete the temp file "\
+                              "explicitly with `tempfile.delete`"
           end
-          chunk.force_encoding(encoding)
-          tempfile.write(chunk)
+        else
+          response = request.perform
         end
-        if tempfile
-          tempfile.close
-          @config.logger.info "Temp file written to #{tempfile.path}, please copy the file to a proper folder "\
-                            "with e.g. `FileUtils.cp(tempfile.path, '/new/file/path')` otherwise the temp file "\
-                            "will be deleted automatically with GC. It's also recommended to delete the temp file "\
-                            "explicitly with `tempfile.delete`"
+
+        if @config.debugging
+          @config.logger.debug "HTTP response body ~BEGIN~\n#{response.body}\n~END~\n"
         end
-      else
-        response = request.perform
-      end
 
-      if @config.debugging
-        @config.logger.debug "HTTP response body ~BEGIN~\n#{response.body}\n~END~\n"
-      end
+        unless response.success?
+          if response.request_timeout?
+            fail APIError.new('Connection timed out')
+          elsif response.code == 0
+            # Errors from libcurl will be made visible here
+            fail APIError.new(:code => 0,
+                              :message => response.return_message)
+          else
+            body = response.body
+            if response.headers['Content-Encoding'].eql?('gzip') && !(body.nil? || body.empty?) then
+              gzip = Zlib::Inflate.new(Zlib::MAX_WBITS + 16)
+              body = gzip.inflate(body)
+              gzip.close
+            end
+            if should_retry(attempt, @config.max_retries, response.code, @config.enable_retry)
+              sleep calculate_retry_interval(response, @config.backoff_base, @config.backoff_multiplier, attempt, @config.timeout)
+              attempt = attempt + 1
+              next
+            else
+              fail APIError.new(:code => response.code,
+                :response_headers => response.headers,
+                :response_body => body),
+              response.message
+            end
+          end
+        end
 
-      unless response.success?
-        if response.request_timeout?
-          fail APIError.new('Connection timed out')
-        elsif response.code == 0
-          # Errors from libcurl will be made visible here
-          fail APIError.new(:code => 0,
-                            :message => response.return_message)
+        if opts[:return_type]
+          data = deserialize(opts[:api_version], response, opts[:return_type])
         else
-          body = response.body
-          if response.headers['Content-Encoding'].eql?('gzip') && !(body.nil? || body.empty?) then
-            gzip = Zlib::Inflate.new(Zlib::MAX_WBITS + 16)
-            body = gzip.inflate(body)
-            gzip.close
-          end
-          fail APIError.new(:code => response.code,
-                            :response_headers => response.headers,
-                            :response_body => body),
-               response.message
+          data = nil
         end
+        return data, response.code, response.headers
       end
+    end
+
+    # Check if an http request should be retried
+    def should_retry(attempt, max_retries, http_code, enable_retry)
+      (http_code == 429 || http_code >= 500) && max_retries > attempt && enable_retry
+    end
 
-      if opts[:return_type]
-        data = deserialize(opts[:api_version], response, opts[:return_type])
+    # Calculate the sleep interval between 2 retry attempts
+    def calculate_retry_interval(response, backoff_base, backoff_multiplier, attempt, timeout)
+      reset_header = response.headers['X-Ratelimit-Reset']
+      if  !reset_header.nil? && !reset_header.empty?
+        sleep_time = reset_header.to_i
       else
-        data = nil
+        sleep_time = (backoff_multiplier**attempt) * backoff_base
+        if timeout && timeout > 0
+          sleep_time = [timeout, sleep_time].min
+        end
       end
-      return data, response.code, response.headers
+      sleep_time
     end
 
     # Build the HTTP request