ExternalVendorCode/onefuzz
1
0
Fork 0
mirror of https://github.com/microsoft/onefuzz.git synced 2025-06-13 02:28:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add caching to Storage.cs (#2102)

Browse Source 
* Add caching to Storage.cs

* make Storage disposable

* make Storage disposable
This commit is contained in:
Cheick Keita
2022-06-29 16:40:31 -07:00
committed by GitHub
parent 7e0565dac1
commit 63bff8bedc
3 changed files with 64 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/ApiService/ApiService/Program.cs
View File

@ -85,7 +85,6 @@ public class Program {
.AddScoped<ITaskOperations, TaskOperations>() .AddScoped<ITaskOperations, TaskOperations>()
.AddScoped<ITaskEventOperations, TaskEventOperations>() .AddScoped<ITaskEventOperations, TaskEventOperations>()
.AddScoped<IQueue, Queue>() .AddScoped<IQueue, Queue>()
.AddScoped<IStorage, Storage>()
.AddScoped<IProxyOperations, ProxyOperations>() .AddScoped<IProxyOperations, ProxyOperations>()
.AddScoped<IProxyForwardOperations, ProxyForwardOperations>() .AddScoped<IProxyForwardOperations, ProxyForwardOperations>()
.AddScoped<IConfigOperations, ConfigOperations>() .AddScoped<IConfigOperations, ConfigOperations>()
@ -112,10 +111,11 @@ public class Program {
.AddScoped<IRequestHandling, RequestHandling>() .AddScoped<IRequestHandling, RequestHandling>()
.AddScoped<IOnefuzzContext, OnefuzzContext>() .AddScoped<IOnefuzzContext, OnefuzzContext>()
.AddScoped<IEndpointAuthorization, EndpointAuthorization>() .AddScoped<IEndpointAuthorization, EndpointAuthorization>()
.AddScoped<INodeMessageOperations, NodeMessageOperations>()
.AddSingleton<ICreds, Creds>() .AddSingleton<ICreds, Creds>()
.AddSingleton<IServiceConfig, ServiceConfiguration>() .AddSingleton<IServiceConfig, ServiceConfiguration>()
.AddSingleton<INodeMessageOperations, NodeMessageOperations>() .AddSingleton<IStorage, Storage>()
.AddHttpClient(); .AddHttpClient();
} }
) )

2
src/ApiService/ApiService/onefuzzlib/Containers.cs
View File

@ -85,7 +85,7 @@ public class Containers : IContainers {
var containerName = _config.OneFuzzStoragePrefix + container.ContainerName; var containerName = _config.OneFuzzStoragePrefix + container.ContainerName;
var containers = _storage.GetAccounts(storageType) var containers = _storage.GetAccounts(storageType).AsEnumerable()
.Reverse() .Reverse()
.Select(async account => (await GetBlobService(account))?.GetBlobContainerClient(containerName)); .Select(async account => (await GetBlobService(account))?.GetBlobContainerClient(containerName));

108
src/ApiService/ApiService/onefuzzlib/Storage.cs
View File

@ -2,6 +2,7 @@
using Azure.Core; using Azure.Core;
using Azure.ResourceManager; using Azure.ResourceManager;
using Azure.ResourceManager.Storage; using Azure.ResourceManager.Storage;
using Microsoft.Extensions.Caching.Memory;
namespace Microsoft.OneFuzz.Service; namespace Microsoft.OneFuzz.Service;
@ -27,17 +28,21 @@ public interface IStorage {
public IEnumerable<string> GetAccounts(StorageType storageType); public IEnumerable<string> GetAccounts(StorageType storageType);
} }
public class Storage : IStorage { public sealed class Storage : IStorage, IDisposable {
private ICreds _creds; private readonly ICreds _creds;
private ArmClient _armClient; private readonly ArmClient _armClient;
private ILogTracer _log; private readonly ILogTracer _log;
private IServiceConfig _config; private readonly IServiceConfig _config;
private readonly MemoryCache _cache;
public Storage(ICreds creds, ILogTracer log, IServiceConfig config) { public Storage(ICreds creds, ILogTracer log, IServiceConfig config) {
_creds = creds; _creds = creds;
_armClient = creds.ArmClient; _armClient = creds.ArmClient;
_log = log; _log = log;
_config = config; _config = config;
_cache = new MemoryCache(new MemoryCacheOptions() {
});
} }
public string GetFuncStorage() { public string GetFuncStorage() {
@ -54,40 +59,41 @@ public class Storage : IStorage {
return _armClient; return _armClient;
} }
// TODO: @cached
public IEnumerable<string> CorpusAccounts() { public IEnumerable<string> CorpusAccounts() {
var skip = GetFuncStorage(); return _cache.GetOrCreate<List<string>>("CorpusAccounts", cacheEntry => {
var results = new List<string> { GetFuzzStorage() }; var skip = GetFuncStorage();
var results = new List<string> { GetFuzzStorage() };
var client = GetMgmtClient(); var client = GetMgmtClient();
var group = _creds.GetResourceGroupResourceIdentifier(); var group = _creds.GetResourceGroupResourceIdentifier();
const string storageTypeTagKey = "storage_type"; const string storageTypeTagKey = "storage_type";
var resourceGroup = client.GetResourceGroupResource(group); var resourceGroup = client.GetResourceGroupResource(group);
foreach (var account in resourceGroup.GetStorageAccounts()) { foreach (var account in resourceGroup.GetStorageAccounts()) {
if (account.Id == skip) { if (account.Id == skip) {
continue; continue;
}
if (results.Contains(account.Id!)) {
continue;
}
if (string.IsNullOrEmpty(account.Data.PrimaryEndpoints.Blob)) {
continue;
}
if (!account.Data.Tags.ContainsKey(storageTypeTagKey)
|| account.Data.Tags[storageTypeTagKey] != "corpus") {
continue;
}
results.Add(account.Id!);
} }
if (results.Contains(account.Id!)) { _log.Info($"corpus accounts: {JsonSerializer.Serialize(results)}");
continue; return results;
} });
if (string.IsNullOrEmpty(account.Data.PrimaryEndpoints.Blob)) {
continue;
}
if (!account.Data.Tags.ContainsKey(storageTypeTagKey)
|| account.Data.Tags[storageTypeTagKey] != "corpus") {
continue;
}
results.Add(account.Id!);
}
_log.Info($"corpus accounts: {JsonSerializer.Serialize(results)}");
return results;
} }
public string GetPrimaryAccount(StorageType storageType) { public string GetPrimaryAccount(StorageType storageType) {
@ -99,22 +105,26 @@ public class Storage : IStorage {
}; };
} }
public async Async.Task<(string, string)> GetStorageAccountNameAndKey(string accountId) { public Async.Task<(string, string)> GetStorageAccountNameAndKey(string accountId) {
var resourceId = new ResourceIdentifier(accountId); return _cache.GetOrCreateAsync<(string, string)>($"GetStorageAccountNameAndKey-{accountId}", async cacheEntry => {
var armClient = GetMgmtClient(); var resourceId = new ResourceIdentifier(accountId);
var storageAccount = armClient.GetStorageAccountResource(resourceId); var armClient = GetMgmtClient();
var keys = await storageAccount.GetKeysAsync(); var storageAccount = armClient.GetStorageAccountResource(resourceId);
var key = keys.Value.Keys.FirstOrDefault() ?? throw new Exception("no keys found"); var keys = await storageAccount.GetKeysAsync();
return (resourceId.Name, key.Value); 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) { public Async.Task<string?> GetStorageAccountNameKeyByName(string accountName) {
var armClient = GetMgmtClient(); return _cache.GetOrCreateAsync<string?>($"GetStorageAccountNameKeyByName-{accountName}", async cacheEntry => {
var resourceGroup = _creds.GetResourceGroupResourceIdentifier(); var armClient = GetMgmtClient();
var storageAccount = await armClient.GetResourceGroupResource(resourceGroup).GetStorageAccountAsync(accountName); var resourceGroup = _creds.GetResourceGroupResourceIdentifier();
var keys = await storageAccount.Value.GetKeysAsync(); var storageAccount = await armClient.GetResourceGroupResource(resourceGroup).GetStorageAccountAsync(accountName);
var key = keys.Value.Keys.FirstOrDefault(); var keys = await storageAccount.Value.GetKeysAsync();
return key?.Value; var key = keys.Value.Keys.FirstOrDefault();
return key?.Value;
});
} }
public string ChooseAccounts(StorageType storageType) { public string ChooseAccounts(StorageType storageType) {
@ -158,4 +168,8 @@ public class Storage : IStorage {
public Uri GetBlobEndpoint(string accountId) public Uri GetBlobEndpoint(string accountId)
=> new($"https://{accountId}.blob.core.windows.net/"); => new($"https://{accountId}.blob.core.windows.net/");
public void Dispose() {
_cache.Dispose();
}
} }