Skip to content
This repository has been archived by the owner on Nov 1, 2023. It is now read-only.

Commit

Permalink
Add caching to Storage.cs (#2102)
Browse files Browse the repository at this point in the history 
* Add caching to Storage.cs

* make Storage disposable

* make Storage disposable
  • Loading branch information
@chkeita
chkeita authored Jun 29, 2022
1 parent 7e0565d commit 63bff8b
Show file tree
Hide file tree
Showing 3 changed files with 61 additions and 47 deletions.
4 changes: 2 additions & 2 deletions src/ApiService/ApiService/Program.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,6 @@ public static void Main() {
.AddScoped<ITaskOperations, TaskOperations>()
.AddScoped<ITaskEventOperations, TaskEventOperations>()
.AddScoped<IQueue, Queue>()
.AddScoped<IStorage, Storage>()
.AddScoped<IProxyOperations, ProxyOperations>()
.AddScoped<IProxyForwardOperations, ProxyForwardOperations>()
.AddScoped<IConfigOperations, ConfigOperations>()
Expand All @@ -112,10 +111,11 @@ public static void Main() {
.AddScoped<IRequestHandling, RequestHandling>()
.AddScoped<IOnefuzzContext, OnefuzzContext>()
.AddScoped<IEndpointAuthorization, EndpointAuthorization>()
.AddScoped<INodeMessageOperations, NodeMessageOperations>()

.AddSingleton<ICreds, Creds>()
.AddSingleton<IServiceConfig, ServiceConfiguration>()
.AddSingleton<INodeMessageOperations, NodeMessageOperations>()
.AddSingleton<IStorage, Storage>()
.AddHttpClient();
}
)
Expand Down
2 changes: 1 addition & 1 deletion src/ApiService/ApiService/onefuzzlib/Containers.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ public Containers(ILogTracer log, IStorage storage, ICreds creds, IServiceConfig

var containerName = _config.OneFuzzStoragePrefix + container.ContainerName;

var containers = _storage.GetAccounts(storageType)
var containers = _storage.GetAccounts(storageType).AsEnumerable()
.Reverse()
.Select(async account => (await GetBlobService(account))?.GetBlobContainerClient(containerName));

Expand Down
102 changes: 58 additions & 44 deletions src/ApiService/ApiService/onefuzzlib/Storage.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
using Azure.Core;
using Azure.ResourceManager;
using Azure.ResourceManager.Storage;
using Microsoft.Extensions.Caching.Memory;

namespace Microsoft.OneFuzz.Service;

Expand All @@ -27,17 +28,21 @@ public interface IStorage {
public IEnumerable<string> GetAccounts(StorageType storageType);
}

public class Storage : IStorage {
private ICreds _creds;
private ArmClient _armClient;
private ILogTracer _log;
private IServiceConfig _config;
public sealed class Storage : IStorage, IDisposable {
private readonly ICreds _creds;
private readonly ArmClient _armClient;
private readonly ILogTracer _log;
private readonly IServiceConfig _config;
private readonly MemoryCache _cache;

public Storage(ICreds creds, ILogTracer log, IServiceConfig config) {
_creds = creds;
_armClient = creds.ArmClient;
_log = log;
_config = config;
_cache = new MemoryCache(new MemoryCacheOptions() {

});
}

public string GetFuncStorage() {
Expand All @@ -54,40 +59,41 @@ public ArmClient GetMgmtClient() {
return _armClient;
}

// TODO: @cached
public IEnumerable<string> CorpusAccounts() {
var skip = GetFuncStorage();
var results = new List<string> { GetFuzzStorage() };
return _cache.GetOrCreate<List<string>>("CorpusAccounts", cacheEntry => {
var skip = GetFuncStorage();
var results = new List<string> { GetFuzzStorage() };

var client = GetMgmtClient();
var group = _creds.GetResourceGroupResourceIdentifier();
var client = GetMgmtClient();
var group = _creds.GetResourceGroupResourceIdentifier();

const string storageTypeTagKey = "storage_type";
const string storageTypeTagKey = "storage_type";

var resourceGroup = client.GetResourceGroupResource(group);
foreach (var account in resourceGroup.GetStorageAccounts()) {
if (account.Id == skip) {
continue;
}
var resourceGroup = client.GetResourceGroupResource(group);
foreach (var account in resourceGroup.GetStorageAccounts()) {
if (account.Id == skip) {
continue;
}

if (results.Contains(account.Id!)) {
continue;
}
if (results.Contains(account.Id!)) {
continue;
}

if (string.IsNullOrEmpty(account.Data.PrimaryEndpoints.Blob)) {
continue;
}
if (string.IsNullOrEmpty(account.Data.PrimaryEndpoints.Blob)) {
continue;
}

if (!account.Data.Tags.ContainsKey(storageTypeTagKey)
|| account.Data.Tags[storageTypeTagKey] != "corpus") {
continue;
}
if (!account.Data.Tags.ContainsKey(storageTypeTagKey)
|| account.Data.Tags[storageTypeTagKey] != "corpus") {
continue;
}

results.Add(account.Id!);
}
results.Add(account.Id!);
}

_log.Info($"corpus accounts: {JsonSerializer.Serialize(results)}");
return results;
_log.Info($"corpus accounts: {JsonSerializer.Serialize(results)}");
return results;
});
}

public string GetPrimaryAccount(StorageType storageType) {
Expand All @@ -99,22 +105,26 @@ public string GetPrimaryAccount(StorageType storageType) {
};
}

public async Async.Task<(string, string)> GetStorageAccountNameAndKey(string accountId) {
var resourceId = new ResourceIdentifier(accountId);
var armClient = GetMgmtClient();
var storageAccount = armClient.GetStorageAccountResource(resourceId);
var keys = await storageAccount.GetKeysAsync();
var key = keys.Value.Keys.FirstOrDefault() ?? throw new Exception("no keys found");
return (resourceId.Name, key.Value);
public Async.Task<(string, string)> GetStorageAccountNameAndKey(string accountId) {
return _cache.GetOrCreateAsync<(string, string)>($"GetStorageAccountNameAndKey-{accountId}", async cacheEntry => {
var resourceId = new ResourceIdentifier(accountId);
var armClient = GetMgmtClient();
var storageAccount = armClient.GetStorageAccountResource(resourceId);
var keys = await storageAccount.GetKeysAsync();
var key = keys.Value.Keys.FirstOrDefault() ?? throw new Exception("no keys found");
return (resourceId.Name, key.Value);
});
}

public async Async.Task<string?> GetStorageAccountNameKeyByName(string accountName) {
var armClient = GetMgmtClient();
var resourceGroup = _creds.GetResourceGroupResourceIdentifier();
var storageAccount = await armClient.GetResourceGroupResource(resourceGroup).GetStorageAccountAsync(accountName);
var keys = await storageAccount.Value.GetKeysAsync();
var key = keys.Value.Keys.FirstOrDefault();
return key?.Value;
public Async.Task<string?> GetStorageAccountNameKeyByName(string accountName) {
return _cache.GetOrCreateAsync<string?>($"GetStorageAccountNameKeyByName-{accountName}", async cacheEntry => {
var armClient = GetMgmtClient();
var resourceGroup = _creds.GetResourceGroupResourceIdentifier();
var storageAccount = await armClient.GetResourceGroupResource(resourceGroup).GetStorageAccountAsync(accountName);
var keys = await storageAccount.Value.GetKeysAsync();
var key = keys.Value.Keys.FirstOrDefault();
return key?.Value;
});
}

public string ChooseAccounts(StorageType storageType) {
Expand Down Expand Up @@ -158,4 +168,8 @@ public Uri GetQueueEndpoint(string accountId)

public Uri GetBlobEndpoint(string accountId)
=> new($"https://{accountId}.blob.core.windows.net/");

public void Dispose() {
_cache.Dispose();
}
}

0 comments on commit 63bff8b

Please sign in to comment.