Skip to content

Commit

Permalink
caches token (#452)
Browse files Browse the repository at this point in the history
  • Loading branch information
@giventocode
giventocode committed Oct 10, 2023
1 parent cd78c00 commit 7774107
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 11 deletions.
28 changes: 27 additions & 1 deletion src/Tes.Runner.Test/Storage/TerraUrlTransformationStrategyTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ public class TerraUrlTransformationStrategyTests
private Mock<TerraWsmApiClient> mockTerraWsmApiClient = null!;
private RuntimeOptions runtimeOptions = null!;
private SasTokenApiParameters capturedSasTokenApiParameters = null!;
private const int SasExpirationInSeconds = 5;

[TestInitialize]
public void SetUp()
Expand All @@ -35,7 +36,7 @@ public void SetUp()
mockTerraWsmApiClient = new Mock<TerraWsmApiClient>();
capturedSasTokenApiParameters = new SasTokenApiParameters("", 0, "", "");
SetupWsmClientWithAssumingSuccess();
transformationStrategy = new TerraUrlTransformationStrategy(runtimeOptions.Terra, mockTerraWsmApiClient.Object);
transformationStrategy = new TerraUrlTransformationStrategy(runtimeOptions.Terra, mockTerraWsmApiClient.Object, SasExpirationInSeconds);
}

private void SetupWsmClientWithAssumingSuccess()
Expand Down Expand Up @@ -106,5 +107,30 @@ public async Task TransformUrlWithStrategyAsync_NonTerraStorageAccount_SourceUrl

Assert.AreEqual(new Uri(sourceUrl).ToString(), sasUrl.ToString());
}

[TestMethod]
public async Task TransformUrlWithStrategyAsync_RequestsSasTokenMoreThanOnce_SasTokenIsCached()
{
var sourceUrl = $"{stubTerraBlobUrl}/blob";
await transformationStrategy.TransformUrlWithStrategyAsync(sourceUrl, BlobSasPermissions.Read);
await transformationStrategy.TransformUrlWithStrategyAsync(sourceUrl, BlobSasPermissions.Read);
mockTerraWsmApiClient.Verify(w => w.GetSasTokenAsync(It.IsAny<Guid>(),
It.IsAny<Guid>(), It.IsAny<SasTokenApiParameters>(), It.IsAny<CancellationToken>()), Times.Once);
}

[TestMethod]
public async Task TransformUrlWithStrategyAsync_RequestsSasTokenMoreThanOnceAfterExpiration_SasTokenIsRenewed()
{
var sourceUrl = $"{stubTerraBlobUrl}/blob";
var sasUrl1 = await transformationStrategy.TransformUrlWithStrategyAsync(sourceUrl, BlobSasPermissions.Read);
Assert.IsNotNull(sasUrl1);

await Task.Delay(TimeSpan.FromSeconds(SasExpirationInSeconds + 1));

var sasUrl2 = await transformationStrategy.TransformUrlWithStrategyAsync(sourceUrl, BlobSasPermissions.Read);
Assert.IsNotNull(sasUrl2);
mockTerraWsmApiClient.Verify(w => w.GetSasTokenAsync(It.IsAny<Guid>(),
It.IsAny<Guid>(), It.IsAny<SasTokenApiParameters>(), It.IsAny<CancellationToken>()), Times.Exactly(2));
}
}
}
48 changes: 38 additions & 10 deletions src/Tes.Runner/Storage/TerraUrlTransformationStrategy.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
using Azure.Core;
using Azure.Storage.Blobs;
using Azure.Storage.Sas;
using Microsoft.Extensions.Caching.Memory;
using Microsoft.Extensions.Logging;
using Tes.ApiClients;
using Tes.ApiClients.Models.Terra;
Expand All @@ -16,30 +17,39 @@ namespace Tes.Runner.Storage
{
public class TerraUrlTransformationStrategy : IUrlTransformationStrategy
{
public const int TokenExpirationInSeconds = 3600 * 24; //1 day, max time allowed by Terra.
public const int CacheExpirationInSeconds = TokenExpirationInSeconds - 1800; // 30 minutes less than token expiration

private const int MaxNumberOfContainerResources = 10000;
private const int TokenExpirationInSeconds = 3600 * 24 * 7; //7 days, max Azure Batch node runtime.
private const string LzStorageAccountNamePattern = "lz[0-9a-f]*";

private readonly TerraWsmApiClient terraWsmApiClient;
private readonly TerraRuntimeOptions terraRuntimeOptions;
private const string LzStorageAccountNamePattern = "lz[0-9a-f]*";
private readonly ILogger<TerraUrlTransformationStrategy> logger = PipelineLoggerFactory.Create<TerraUrlTransformationStrategy>();
private readonly IMemoryCache memoryCache;
private readonly int cacheExpirationInSeconds;

public TerraUrlTransformationStrategy(TerraRuntimeOptions terraRuntimeOptions, TokenCredential tokenCredential)
public TerraUrlTransformationStrategy(TerraRuntimeOptions terraRuntimeOptions, TokenCredential tokenCredential, int cacheExpirationInSeconds = CacheExpirationInSeconds)
{
ArgumentNullException.ThrowIfNull(terraRuntimeOptions);
ArgumentException.ThrowIfNullOrEmpty(terraRuntimeOptions.WsmApiHost, nameof(terraRuntimeOptions.WsmApiHost));

terraWsmApiClient = TerraWsmApiClient.CreateTerraWsmApiClient(terraRuntimeOptions.WsmApiHost, tokenCredential);
this.terraRuntimeOptions = terraRuntimeOptions;
memoryCache = new MemoryCache(new MemoryCacheOptions());
this.cacheExpirationInSeconds = cacheExpirationInSeconds;
}


public TerraUrlTransformationStrategy(TerraRuntimeOptions terraRuntimeOptions, TerraWsmApiClient terraWsmApiClient)
public TerraUrlTransformationStrategy(TerraRuntimeOptions terraRuntimeOptions, TerraWsmApiClient terraWsmApiClient, int cacheExpirationInSeconds = CacheExpirationInSeconds)
{
ArgumentNullException.ThrowIfNull(terraRuntimeOptions);
ArgumentNullException.ThrowIfNull(terraWsmApiClient);

this.terraWsmApiClient = terraWsmApiClient;
this.terraRuntimeOptions = terraRuntimeOptions;
memoryCache = new MemoryCache(new MemoryCacheOptions());
this.cacheExpirationInSeconds = cacheExpirationInSeconds;
}

public async Task<Uri> TransformUrlWithStrategyAsync(string sourceUrl, BlobSasPermissions blobSasPermissions)
Expand All @@ -65,7 +75,7 @@ public async Task<Uri> TransformUrlWithStrategyAsync(string sourceUrl, BlobSasPe
/// <returns>URL with a SAS token</returns>
private async Task<Uri> GetMappedSasUrlFromWsmAsync(TerraBlobInfo blobInfo, BlobSasPermissions blobSasPermissions)
{
var tokenInfo = await GetWorkspaceBlobSasTokenFromWsmAsync(blobInfo, blobSasPermissions);
var tokenInfo = await GetWorkspaceSasTokenFromWsmAsync(blobInfo, blobSasPermissions);

logger.LogInformation($"Successfully obtained the SAS URL from Terra. WSM resource ID:{blobInfo.WsmContainerResourceId}");

Expand All @@ -82,25 +92,43 @@ private async Task<Uri> GetMappedSasUrlFromWsmAsync(TerraBlobInfo blobInfo, Blob
return uriBuilder.Uri;
}

private async Task<WsmSasTokenApiResponse> GetWorkspaceBlobSasTokenFromWsmAsync(TerraBlobInfo blobInfo, BlobSasPermissions sasBlobPermissions)
private async Task<WsmSasTokenApiResponse> GetWorkspaceSasTokenFromWsmAsync(TerraBlobInfo blobInfo, BlobSasPermissions sasBlobPermissions)
{
var tokenParams = CreateTokenParamsFromOptions(blobInfo.BlobName, sasBlobPermissions);
var tokenParams = CreateTokenParamsFromOptions(sasBlobPermissions);

logger.LogInformation(
$"Getting SAS URL from Terra. WSM workspace ID:{blobInfo.WorkspaceId}");

return await terraWsmApiClient.GetSasTokenAsync(
var cacheKey = $"{blobInfo.WorkspaceId}-{blobInfo.WsmContainerResourceId}-{tokenParams.SasPermission}";

if (memoryCache.TryGetValue(cacheKey, out WsmSasTokenApiResponse? tokenInfo))
{
if (tokenInfo is null)
{
throw new InvalidOperationException("The value retrieved from the cache is null");
}

logger.LogInformation($"SAS URL found in cache. WSM resource ID:{blobInfo.WsmContainerResourceId}");
return tokenInfo;
}

var token = await terraWsmApiClient.GetSasTokenAsync(
blobInfo.WorkspaceId,
blobInfo.WsmContainerResourceId,
tokenParams, CancellationToken.None);

memoryCache.Set(cacheKey, token, TimeSpan.FromSeconds(cacheExpirationInSeconds));

return token;
}

private SasTokenApiParameters CreateTokenParamsFromOptions(string blobName, BlobSasPermissions sasPermissions)
private SasTokenApiParameters CreateTokenParamsFromOptions(BlobSasPermissions sasPermissions)
{
return new(
terraRuntimeOptions.SasAllowedIpRange ?? string.Empty,
TokenExpirationInSeconds,
ToWsmBlobSasPermissions(sasPermissions), blobName);
//setting blob name to empty string to get a SAS token for the container
ToWsmBlobSasPermissions(sasPermissions), SasBlobName: String.Empty);
}

private string ToWsmBlobSasPermissions(BlobSasPermissions blobSasPermissions)
Expand Down

0 comments on commit 7774107

Please sign in to comment.