Skip to content

Commit

Permalink
Perform background refresh of credentials during preempt expiry time …
Browse files Browse the repository at this point in the history 
…period
  • Loading branch information
@kevinstapleton
kevinstapleton committed Nov 1, 2024
1 parent ad13fd4 commit c221644
Show file tree
Hide file tree
Showing 3 changed files with 329 additions and 31 deletions.
126 changes: 96 additions & 30 deletions sdk/src/Core/Amazon.Runtime/Credentials/RefreshingAWSCredentials.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,11 +36,7 @@ public abstract class RefreshingAWSCredentials : AWSCredentials, IDisposable
/// </summary>
public class CredentialsRefreshState
{
public ImmutableCredentials Credentials
{
get;
set;
}
public ImmutableCredentials Credentials { get; set; }
public DateTime Expiration { get; set; }

public CredentialsRefreshState()
Expand All @@ -61,6 +57,16 @@ internal bool IsExpiredWithin(TimeSpan preemptExpiryTime)
var exp = Expiration.ToUniversalTime();
return (now > exp - preemptExpiryTime);
}

internal TimeSpan GetTimeToLive(TimeSpan preemptExpiryTime)
{
#pragma warning disable CS0612,CS0618 // Type or member is obsolete
var now = AWSSDKUtils.CorrectedUtcNow;
#pragma warning restore CS0612,CS0618 // Type or member is obsolete
var exp = Expiration.ToUniversalTime();

return exp - now + preemptExpiryTime;
}
}

/// <summary>
Expand Down Expand Up @@ -120,46 +126,106 @@ public TimeSpan PreemptExpiryTime
/// <returns></returns>
public override ImmutableCredentials GetCredentials()
{
_updateGeneratedCredentialsSemaphore.Wait();
try
// We save the currentState as it might be modified or cleared.
var tempState = currentState;

var ttl = tempState?.GetTimeToLive(PreemptExpiryTime);

if (ttl > TimeSpan.Zero)
{
// We save the currentState as it might be modified or cleared.
var tempState = currentState;
// If credentials are expired or we don't have any state yet, update
if (ShouldUpdateState(tempState, PreemptExpiryTime))
if (ttl < PreemptExpiryTime)
{
tempState = GenerateNewCredentials();
UpdateToGeneratedCredentials(tempState, PreemptExpiryTime);
currentState = tempState;
// background refresh (fire & forget)
if (_updateGeneratedCredentialsSemaphore.Wait(0))
{
_ = System.Threading.Tasks.Task.Run(GenerateCredentialsAndUpdateState);
}
}
return tempState.Credentials.Copy();
}
finally
else
{
_updateGeneratedCredentialsSemaphore.Release();
// If credentials are expired, update
_updateGeneratedCredentialsSemaphore.Wait();
tempState = GenerateCredentialsAndUpdateState();
}

return tempState.Credentials.Copy();

CredentialsRefreshState GenerateCredentialsAndUpdateState()
{
System.Diagnostics.Debug.Assert(_updateGeneratedCredentialsSemaphore.CurrentCount == 0);

try
{
var tempState = currentState;
// double-check that the credentials still need updating
// as it's possible that multiple requests were queued acquiring the semaphore
if (ShouldUpdateState(tempState, PreemptExpiryTime))
{
tempState = GenerateNewCredentials();
UpdateToGeneratedCredentials(tempState, PreemptExpiryTime);
currentState = tempState;
}

return tempState;
}
finally
{
_updateGeneratedCredentialsSemaphore.Release();
}
}
}

#if AWS_ASYNC_API
public override async System.Threading.Tasks.Task<ImmutableCredentials> GetCredentialsAsync()
{
await _updateGeneratedCredentialsSemaphore.WaitAsync().ConfigureAwait(false);
try
// We save the currentState as it might be modified or cleared.
var tempState = currentState;

var ttl = tempState?.GetTimeToLive(PreemptExpiryTime);

if (ttl > TimeSpan.Zero)
{
// We save the currentState as it might be modified or cleared.
var tempState = currentState;
// If credentials are expired, update
if (ShouldUpdateState(tempState, PreemptExpiryTime))
if (ttl < PreemptExpiryTime)
{
tempState = await GenerateNewCredentialsAsync().ConfigureAwait(false);
UpdateToGeneratedCredentials(tempState, PreemptExpiryTime);
currentState = tempState;
// background refresh (fire & forget)
if (_updateGeneratedCredentialsSemaphore.Wait(0))
{
_ = GenerateCredentialsAndUpdateStateAsync();
}
}
return tempState.Credentials.Copy();
}
finally
else
{
// If credentials are expired, update
await _updateGeneratedCredentialsSemaphore.WaitAsync().ConfigureAwait(false);
tempState = await GenerateCredentialsAndUpdateStateAsync().ConfigureAwait(false);
}

return tempState.Credentials.Copy();

async System.Threading.Tasks.Task<CredentialsRefreshState> GenerateCredentialsAndUpdateStateAsync()
{
_updateGeneratedCredentialsSemaphore.Release();
System.Diagnostics.Debug.Assert(_updateGeneratedCredentialsSemaphore.CurrentCount == 0);

try
{
var tempState = currentState;
// double-check that the credentials still need updating
// as it's possible that multiple requests were queued acquiring the semaphore
if (ShouldUpdateState(tempState, PreemptExpiryTime))
{
tempState = await GenerateNewCredentialsAsync().ConfigureAwait(false);
UpdateToGeneratedCredentials(tempState, PreemptExpiryTime);
currentState = tempState;
}

return tempState;
}
finally
{
_updateGeneratedCredentialsSemaphore.Release();
}
}
}
#endif
Expand Down Expand Up @@ -262,7 +328,7 @@ protected virtual CredentialsRefreshState GenerateNewCredentials()
throw new NotImplementedException();
}
#if AWS_ASYNC_API
/// <summary>
/// <summary>
/// When overridden in a derived class, generates new credentials and new expiration date.
///
/// Called on first credentials request and when expiration date is in the past.
Expand Down
2 changes: 1 addition & 1 deletion sdk/test/NetStandard/UnitTests/AWSSDK.UnitTests.Custom.NetStandard.csproj
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ This project file should not be used as part of a release pipeline.
<PackageReference Include="System.Text.RegularExpressions" Version="4.3.1" />
<PackageReference Include="xunit" Version="2.4.2" />
<PackageReference Include="Moq" Version="4.18.4" />
<PackageReference Include="xunit.runner.visualstudio" Version="2.4.5" >
<PackageReference Include="xunit.runner.visualstudio" Version="2.4.5">
<PrivateAssets>all</PrivateAssets>
<IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
</PackageReference>
Expand Down
Loading

0 comments on commit c221644

Please sign in to comment.