Refresh AWS credentials asynchronously (#2642)

CONTRIBUTING.md

@@ -51,7 +51,10 @@ Specifically, here are a few things that we would appreciate help on:
 The AWS SDK for Ruby is unit tested using RSpec. You can run the unit tests of the SDK after cloning the repo:
 
     bundle install
-    bundle exec rake test
+    bundle exec rake test:spec
+
+If you want to run `PURE_RUBY` tests, then `export PURE_RUBY=1` into your environment. This will skip installing
+packages like `oj` for instance.
 
 To run integration tests, create a `integration-test-config.json` file at the root of this repository. It should
 contain a `"region"` and credentials. Running rake test when this file is present will enable integration tests.

gems/aws-sdk-cognitoidentity/lib/aws-sdk-cognitoidentity/customizations/cognito_identity_credentials.rb

@@ -85,6 +85,7 @@ def initialize(options = {})
         @identity_id = options.delete(:identity_id)
         @custom_role_arn = options.delete(:custom_role_arn)
         @logins = options.delete(:logins) || {}
+        @async_refresh = false
 
         if !@identity_pool_id && !@identity_id
           raise ArgumentError,
@@ -113,6 +114,8 @@ def identity_id
       private
 
       def refresh
+        @before_refresh.call(self) if @before_refresh
+
         resp = @client.get_credentials_for_identity(
           identity_id: identity_id,
           custom_role_arn: @custom_role_arn

gems/aws-sdk-core/CHANGELOG.md

@@ -1,6 +1,8 @@
 Unreleased Changes
 ------------------
 
+* Feature - Asynchronously refresh AWS credentials (#2641).
+
 * Issue - Add x-amz-region-set to list of headers deleted for re-sign.
 
 3.129.1 (2022-03-10)

gems/aws-sdk-core/lib/aws-sdk-core/assume_role_credentials.rb

@@ -54,6 +54,7 @@ def initialize(options = {})
         end
       end
       @client = client_opts[:client] || STS::Client.new(client_opts)
+      @async_refresh = true
       super
     end
 

gems/aws-sdk-core/lib/aws-sdk-core/assume_role_web_identity_credentials.rb

@@ -17,7 +17,7 @@ module Aws
   #       ...
   #     )
   #     For full list of parameters accepted
-  #     @see Aws::STS::Client#assume_role_with_web_identity 
+  #     @see Aws::STS::Client#assume_role_with_web_identity
   #
   #
   # If you omit `:client` option, a new {STS::Client} object will be
@@ -48,6 +48,7 @@ def initialize(options = {})
       client_opts = {}
       @assume_role_web_identity_params = {}
       @token_file = options.delete(:web_identity_token_file)
+      @async_refresh = true
       options.each_pair do |key, value|
         if self.class.assume_role_web_identity_options.include?(key)
           @assume_role_web_identity_params[key] = value

gems/aws-sdk-core/lib/aws-sdk-core/ecs_credentials.rb

@@ -62,6 +62,7 @@ def initialize options = {}
       @http_read_timeout = options[:http_read_timeout] || 5
       @http_debug_output = options[:http_debug_output]
       @backoff = backoff(options[:backoff])
+      @async_refresh = false
       super
     end
 

gems/aws-sdk-core/lib/aws-sdk-core/instance_profile_credentials.rb

@@ -79,6 +79,7 @@ def initialize(options = {})
       @token_ttl = options[:token_ttl] || 21_600
       @token = nil
       @no_refresh_until = nil
+      @async_refresh = false
       super
     end
 

gems/aws-sdk-core/lib/aws-sdk-core/process_credentials.rb

@@ -27,6 +27,7 @@ class ProcessCredentials
     def initialize(process)
       @process = process
       @credentials = credentials_from_process(@process)
+      @async_refresh = false
 
       super
     end
@@ -73,9 +74,9 @@ def refresh
       @credentials = credentials_from_process(@process)
     end
 
-    def near_expiration?
+    def near_expiration?(expiration_length)
       # are we within 5 minutes of expiration?
-      @expiration && (Time.now.to_i + 5 * 60) > @expiration.to_i
+      @expiration && (Time.now.to_i + expiration_length) > @expiration.to_i
     end
   end
 end

gems/aws-sdk-core/lib/aws-sdk-core/refreshing_credentials.rb

@@ -17,6 +17,9 @@ module Aws
   # @api private
   module RefreshingCredentials
 
+    SYNC_EXPIRATION_LENGTH = 300 # 5 minutes
+    ASYNC_EXPIRATION_LENGTH = 600 # 10 minutes
+
     def initialize(options = {})
       @mutex = Mutex.new
       @before_refresh = options.delete(:before_refresh) if Hash === options
@@ -27,13 +30,13 @@ def initialize(options = {})
 
     # @return [Credentials]
     def credentials
-      refresh_if_near_expiration
+      refresh_if_near_expiration!
       @credentials
     end
 
     # @return [Time,nil]
     def expiration
-      refresh_if_near_expiration
+      refresh_if_near_expiration!
       @expiration
     end
 
@@ -49,24 +52,39 @@ def refresh!
 
     private
 
-    # Refreshes instance metadata credentials if they are within
-    # 5 minutes of expiration.
-    def refresh_if_near_expiration
-      if near_expiration?
+    # Refreshes credentials asynchronously and synchronously.
+    # If we are near to expiration, block while getting new credentials.
+    # Otherwise, if we're approaching expiration, use the existing credentials
+    # but attempt a refresh in the background.
+    def refresh_if_near_expiration!
+      # Note: This check is an optimization. Rather than acquire the mutex on every #refresh_if_near_expiration
+      # call, we check before doing so, and then we check within the mutex to avoid a race condition.
+      # See issue: https://github.com/aws/aws-sdk-ruby/issues/2641 for more info.
+      if near_expiration?(SYNC_EXPIRATION_LENGTH)
         @mutex.synchronize do
-          if near_expiration?
+          if near_expiration?(SYNC_EXPIRATION_LENGTH)
             @before_refresh.call(self) if @before_refresh
-
             refresh
           end
         end
+      elsif @async_refresh && near_expiration?(ASYNC_EXPIRATION_LENGTH)
+        unless @mutex.locked?
+          Thread.new do
+            @mutex.synchronize do
+              if near_expiration?(ASYNC_EXPIRATION_LENGTH)
+                @before_refresh.call(self) if @before_refresh
+                refresh
+              end
+            end
+          end
+        end
       end
     end
 
-    def near_expiration?
+    def near_expiration?(expiration_length)
       if @expiration
-        # are we within 5 minutes of expiration?
-        (Time.now.to_i + 5 * 60) > @expiration.to_i
+        # Are we within expiration?
+        (Time.now.to_i + expiration_length) > @expiration.to_i
       else
         true
       end

gems/aws-sdk-core/lib/aws-sdk-core/sso_credentials.rb

@@ -86,6 +86,7 @@ def initialize(options = {})
       options[:region] = @sso_region
       options[:credentials] = nil
       @client = options[:client] || Aws::SSO::Client.new(options)
+      @async_refresh = true
       super
     end
 

gems/aws-sdk-core/spec/aws/assume_role_credentials_spec.rb

@@ -91,6 +91,17 @@ module Aws
       expect(c.expiration).to eq(in_one_hour)
     end
 
+    it 'refreshes asynchronously' do
+      # expiration 6 minutes out, within the async exp time window
+      allow(credentials).to receive(:expiration).and_return(Time.now + (6*60))
+      expect(client).to receive(:assume_role).at_least(2).times
+      expect(Thread).to receive(:new).and_yield
+      c = AssumeRoleCredentials.new(
+        role_arn: 'arn',
+        role_session_name: 'session')
+      c.credentials
+    end
+
     it 'refreshes credentials automatically when they are near expiration' do
       allow(credentials).to receive(:expiration).and_return(Time.now)
       expect(client).to receive(:assume_role).exactly(4).times

gems/aws-sdk-core/spec/aws/assume_role_web_identity_credentials_spec.rb

@@ -151,6 +151,20 @@ module Aws
       expect(c.expiration).to eq(in_one_hour)
     end
 
+    it 'refreshes asynchronously' do
+      # expiration 6 minutes out, within the async exp time window
+      allow(credentials).to receive(:expiration).and_return(Time.now + (6*60))
+      expect(client).to receive(:assume_role_with_web_identity).exactly(2).times
+      expect(File).to receive(:read).with(token_file_path).exactly(2).times
+      expect(Thread).to receive(:new).and_yield
+
+      c = AssumeRoleWebIdentityCredentials.new(
+        role_arn: 'arn',
+        web_identity_token_file: token_file_path,
+        role_session_name: 'session')
+      c.credentials
+    end
+
     it 'auto refreshes credentials when near expiration' do
       allow(credentials).to receive(:expiration).and_return(Time.now)
       expect(client).to receive(:assume_role_with_web_identity).exactly(4).times
@@ -164,6 +178,5 @@ module Aws
       c.credentials
       c.credentials
     end
-
   end
 end