Cache FindContainer (#3304)

Closes #3156.

From now on, we will store additional information about containers in a table (`ContainerInformation`). This is populated when a container is created, or lazily when looking for a container we have not yet stored in the table.

Reads from this table are also cached in memory for a short period of time (currently 10 minutes).

The table is currently used to speed up lookups for containers across storage accounts; we round-robin creation of containers across multiple storage accounts to spread load, so when we read from a container we need to find what storage account it lives in.

The `ContainerInformation` table could also be used in future to store information about access control (see, e.g. #1419).

Author: Porges

src/ApiService/ApiService/Functions/Containers.cs

@@ -41,7 +41,7 @@ private async Async.Task<HttpResponseData> Get(HttpRequestData req) {
                     req,
                     Error.Create(
                         ErrorCode.INVALID_REQUEST,
-                        "invalid container"),
+                        $"invalid container '{get.Name}'"),
                     context: get.Name.String);
             }
 
@@ -75,13 +75,7 @@ private async Async.Task<HttpResponseData> Delete(HttpRequestData req) {
 
         var delete = request.OkV;
         _logger.LogInformation("deleting {ContainerName}", delete.Name);
-        var container = await _context.Containers.FindContainer(delete.Name, StorageType.Corpus);
-
-        var deleted = false;
-        if (container is not null) {
-            deleted = await container.DeleteIfExistsAsync();
-        }
-
+        var deleted = await _context.Containers.DeleteContainerIfExists(delete.Name, StorageType.Corpus);
         return await RequestHandling.Ok(req, deleted);
     }
 

src/ApiService/ApiService/OneFuzzTypes/Model.cs

@@ -173,7 +173,7 @@ public static Error Create(ErrorCode code, params string[] errors)
         => new(code, errors.ToList());
 
     public sealed override string ToString() {
-        var errorsString = Errors != null ? string.Concat("; ", Errors) : string.Empty;
+        var errorsString = Errors != null ? string.Join("; ", Errors) : string.Empty;
         return $"Error {{ Code = {Code}, Errors = {errorsString} }}";
     }
 };
@@ -767,6 +767,11 @@ public record ClientCredentials
     string ClientSecret
 );
 
+public record ContainerInformation(
+    [PartitionKey] StorageType Type,
+    [RowKey] Container Name,
+    string ResourceId // full ARM resource ID for the container
+) : EntityBase;
 
 public record AgentConfig(
     ClientCredentials? ClientCredentials,

src/ApiService/ApiService/OneFuzzTypes/ReturnTypes.cs

@@ -2,6 +2,16 @@
 
 namespace Microsoft.OneFuzz.Service {
 
+    public static class Result {
+        public readonly record struct OkPlaceholder();
+        public readonly record struct OkPlaceholder<T>(T Value);
+        public readonly record struct ErrPlaceholder<T>(T Value);
+
+        public static OkPlaceholder Ok() => new();
+        public static OkPlaceholder<T> Ok<T>(T value) => new(value);
+        public static ErrPlaceholder<T> Error<T>(T value) => new(value);
+    }
+
     public readonly struct ResultVoid<T_Error> {
         public static ResultVoid<T_Error> Ok() => new();
         public static ResultVoid<T_Error> Error(T_Error err) => new(err);
@@ -13,15 +23,16 @@ public readonly struct ResultVoid<T_Error> {
         public bool IsOk { get; }
 
         public T_Error? ErrorV { get; }
-    }
 
+        public static implicit operator ResultVoid<T_Error>(Result.OkPlaceholder _ok) => Ok();
+        public static implicit operator ResultVoid<T_Error>(Result.ErrPlaceholder<T_Error> err) => Error(err.Value);
+    }
 
     public readonly struct Result<T_Ok, T_Error> {
         public static Result<T_Ok, T_Error> Ok(T_Ok ok) => new(ok);
         public static Result<T_Ok, T_Error> Error(T_Error err) => new(err);
 
         private Result(T_Ok ok) => (OkV, ErrorV, IsOk) = (ok, default, true);
-
         private Result(T_Error error) => (ErrorV, OkV, IsOk) = (error, default, false);
 
         [MemberNotNullWhen(returnValue: true, member: nameof(OkV))]
@@ -31,6 +42,9 @@ public readonly struct Result<T_Ok, T_Error> {
         public T_Error? ErrorV { get; }
 
         public T_Ok? OkV { get; }
+
+        public static implicit operator Result<T_Ok, T_Error>(Result.OkPlaceholder<T_Ok> ok) => Ok(ok.Value);
+        public static implicit operator Result<T_Ok, T_Error>(Result.ErrPlaceholder<T_Error> err) => Error(err.Value);
     }
 
     public static class OneFuzzResult {
@@ -61,6 +75,8 @@ public readonly struct OneFuzzResult<T_Ok> {
 
         // Allow simple conversion of Errors to Results.
         public static implicit operator OneFuzzResult<T_Ok>(Error err) => new(err);
+        public static implicit operator OneFuzzResult<T_Ok>(Result.OkPlaceholder<T_Ok> ok) => Ok(ok.Value);
+        public static implicit operator OneFuzzResult<T_Ok>(Result.ErrPlaceholder<Error> err) => Error(err.Value);
     }
 
     public readonly struct OneFuzzResultVoid {
@@ -83,5 +99,7 @@ public readonly struct OneFuzzResultVoid {
 
         // Allow simple conversion of Errors to Results.
         public static implicit operator OneFuzzResultVoid(Error err) => new(err);
+        public static implicit operator OneFuzzResultVoid(Result.OkPlaceholder _ok) => Ok;
+        public static implicit operator OneFuzzResultVoid(Result.ErrPlaceholder<Error> err) => Error(err.Value);
     }
 }

src/ApiService/ApiService/onefuzzlib/Containers.cs

@@ -2,12 +2,15 @@
 using System.IO.Compression;
 using System.Threading;
 using System.Threading.Tasks;
+using ApiService.OneFuzzLib.Orm;
 using Azure;
 using Azure.Core;
+using Azure.ResourceManager.Storage;
 using Azure.Storage.Blobs;
 using Azure.Storage.Blobs.Models;
 using Azure.Storage.Blobs.Specialized;
 using Azure.Storage.Sas;
+using Microsoft.Extensions.Caching.Memory;
 using Microsoft.Extensions.Logging;
 namespace Microsoft.OneFuzz.Service;
 
@@ -20,6 +23,7 @@ public interface IContainers {
     public Async.Task<BlobContainerClient?> GetOrCreateContainerClient(Container container, StorageType storageType, IDictionary<string, string>? metadata);
 
     public Async.Task<BlobContainerClient?> FindContainer(Container container, StorageType storageType);
+    public Async.Task<bool> DeleteContainerIfExists(Container container, StorageType storageType);
 
     public Async.Task<Uri?> GetFileSasUrl(Container container, string name, StorageType storageType, BlobSasPermissions permissions, TimeSpan? duration = null);
     public Async.Task SaveBlob(Container container, string name, string data, StorageType storageType, DateOnly? expiresOn = null);
@@ -40,19 +44,26 @@ public interface IContainers {
     public Async.Task DeleteAllExpiredBlobs();
 }
 
-public class Containers : IContainers {
+public class Containers : Orm<ContainerInformation>, IContainers {
     private readonly ILogger _log;
     private readonly IStorage _storage;
     private readonly IServiceConfig _config;
-    private readonly IOnefuzzContext _context;
+    private readonly IMemoryCache _cache;
 
     static readonly TimeSpan CONTAINER_SAS_DEFAULT_DURATION = TimeSpan.FromDays(30);
+    static readonly TimeSpan CONTAINER_INFO_EXPIRATION_TIME = TimeSpan.FromMinutes(10);
+
+    public Containers(
+        ILogger<Containers> log,
+        IStorage storage,
+        IServiceConfig config,
+        IOnefuzzContext context,
+        IMemoryCache cache) : base(log, context) {
 
-    public Containers(ILogger<Containers> log, IStorage storage, IServiceConfig config, IOnefuzzContext context) {
         _log = log;
         _storage = storage;
         _config = config;
-        _context = context;
+        _cache = cache;
 
         _getInstanceId = new Lazy<Async.Task<Guid>>(async () => {
             var (data, tags) = await GetBlob(WellKnownContainers.BaseConfig, "instance_id", StorageType.Config);
@@ -74,7 +85,6 @@ public Containers(ILogger<Containers> log, IStorage storage, IServiceConfig conf
 
     public async Async.Task<(BinaryData? data, IDictionary<string, string>? tags)> GetBlob(Container container, string name, StorageType storageType) {
         var client = await FindContainer(container, storageType);
-
         if (client == null) {
             return (null, null);
         }
@@ -125,32 +135,105 @@ private static readonly BlobContainerSasPermissions _containerCreatePermissions
             return null;
         }
 
-        return cc;
-    }
+        // record the fact that we created a new container
+        var resourceId = BlobContainerResource.CreateResourceIdentifier(
+            account.SubscriptionId,
+            account.ResourceGroupName,
+            account.Name,
+            containerName);
 
+        _ = await SetContainerInformation(container, storageType, resourceId);
 
-    public async Async.Task<BlobContainerClient?> FindContainer(Container container, StorageType storageType) {
-        // # check secondary accounts first by searching in reverse.
-        // #
-        // # By implementation, the primary account is specified first, followed by
-        // # any secondary accounts.
-        // #
-        // # Secondary accounts, if they exist, are preferred for containers and have
-        // # increased IOP rates, this should be a slight optimization
+        return cc;
+    }
 
+    private async Task<ResourceIdentifier?> FindContainerInAccounts(Container container, StorageType storageType) {
         var containerName = _config.OneFuzzStoragePrefix + container;
-
+        // Check secondary accounts first by searching in reverse.
+        // 
+        // By implementation, the primary account is specified first, followed by
+        // any secondary accounts.
+        // 
+        // Secondary accounts, if they exist, are preferred for containers and have
+        // increased IOP rates, this should be a slight optimization
         foreach (var account in _storage.GetAccounts(storageType).Reverse()) {
             var accountClient = await _storage.GetBlobServiceClientForAccount(account);
             var containerClient = accountClient.GetBlobContainerClient(containerName);
             if (await containerClient.ExistsAsync()) {
-                return containerClient;
+                return BlobContainerResource.CreateResourceIdentifier(
+                    account.SubscriptionId,
+                    account.ResourceGroupName,
+                    account.Name,
+                    containerName);
             }
         }
 
         return null;
     }
 
+    private sealed record ContainerKey(StorageType storageType, Container container);
+    private async Task<ContainerInformation> SetContainerInformation(Container container, StorageType storageType, ResourceIdentifier resourceId) {
+        var containerInfo = new ContainerInformation(storageType, container, resourceId.ToString());
+        _ = await Replace(containerInfo);
+        _ = _cache.Set(new ContainerKey(storageType, container), containerInfo, CONTAINER_INFO_EXPIRATION_TIME);
+        return containerInfo;
+    }
+
+    private async Task<bool> DeleteContainerInformation(Container container, StorageType storageType) {
+        var result = await DeleteIfExists(storageType.ToString(), container.ToString());
+        _cache.Remove(new ContainerKey(storageType, container));
+        return result.IsOk && result.OkV;
+    }
+
+    private async Task<ContainerInformation?> LoadContainerInformation(Container container, StorageType storageType) {
+        // first, try cache
+        var info = _cache.Get<ContainerInformation>(new ContainerKey(storageType, container));
+        if (info is not null) {
+            return info;
+        }
+
+        // next try the table
+        var result = await QueryAsync(Query.SingleEntity(storageType.ToString(), container.ToString())).FirstOrDefaultAsync();
+        if (result is not null) {
+            _ = _cache.Set(new ContainerKey(storageType, container), result, CONTAINER_INFO_EXPIRATION_TIME);
+            return result;
+        }
+
+        // we don't have metadata in the table about this account yet, find it:
+        var resourceId = await FindContainerInAccounts(container, storageType);
+        if (resourceId is null) {
+            // never negatively-cache container info, so containers created by other instances
+            // can be found instantly
+            return null;
+        }
+
+        // we found the container, insert it into the table (and cache) so we find it next time:
+        return await SetContainerInformation(container, storageType, resourceId);
+    }
+
+    public async Async.Task<BlobContainerClient?> FindContainer(Container container, StorageType storageType) {
+        var containerInfo = await LoadContainerInformation(container, storageType);
+        if (containerInfo is null) {
+            return null;
+        }
+
+        return await _storage.GetBlobContainerClientForContainerResource(new ResourceIdentifier(containerInfo.ResourceId));
+    }
+
+    public async Async.Task<bool> DeleteContainerIfExists(Container container, StorageType storageType) {
+        var client = await FindContainer(container, storageType);
+        if (client is null) {
+            // container doesn't exist
+            return false;
+        }
+
+        var (_, result) = await (
+            DeleteContainerInformation(container, storageType),
+            client.DeleteIfExistsAsync());
+
+        return result;
+    }
+
     public async Async.Task<Uri?> GetFileSasUrl(Container container, string name, StorageType storageType, BlobSasPermissions permissions, TimeSpan? duration = null) {
         var client = await FindContainer(container, storageType);
         if (client is null) {

src/ApiService/ApiService/onefuzzlib/Storage.cs

@@ -1,4 +1,5 @@
-using System.Text.Json;
+using System.Diagnostics;
+using System.Text.Json;
 using System.Threading.Tasks;
 using Azure.Core;
 using Azure.Data.Tables;
@@ -47,6 +48,21 @@ public Task<BlobServiceClient> GetBlobServiceClientForAccount(ResourceIdentifier
 
     public Task<BlobServiceClient> GetBlobServiceClientForAccountName(string accountName);
 
+    public async Task<BlobContainerClient> GetBlobContainerClientForContainerResource(ResourceIdentifier containerResourceId) {
+        if (containerResourceId.ResourceType != BlobContainerResource.ResourceType) {
+            throw new ArgumentException("must be a container resource ID", nameof(containerResourceId));
+        }
+
+        // parent of container resource = blob services resource
+        // parent of blob services resource = storage account resource
+        var accountId = containerResourceId.Parent?.Parent ?? throw new ArgumentException("resource ID must have parent", nameof(containerResourceId));
+
+        Debug.Assert(accountId.ResourceType == StorageAccountResource.ResourceType);
+
+        var accountClient = await GetBlobServiceClientForAccount(accountId);
+        return accountClient.GetBlobContainerClient(containerResourceId.Name);
+    }
+
     public Uri GenerateBlobContainerSasUri(
         BlobContainerSasPermissions permissions,
         BlobContainerClient containerClient,

src/ApiService/ApiService/onefuzzlib/orm/EntityConverter.cs

@@ -257,6 +257,8 @@ public async Async.Task<TableEntity> ToTableEntity<T>(T typedEntity) where T : E
                 return int.Parse(stringValue);
             } else if (ef.type == typeof(long)) {
                 return long.Parse(stringValue);
+            } else if (ef.type.IsEnum) {
+                return Enum.Parse(ef.type, stringValue);
             } else if (ef.type.IsClass) {
                 try {
                     if (ef.type.GetMethod("Parse", BindingFlags.Static | BindingFlags.Public) is MethodInfo mi) {
@@ -275,7 +277,6 @@ public async Async.Task<TableEntity> ToTableEntity<T>(T typedEntity) where T : E
         var fieldName = ef.columnName;
         var obj = entity[fieldName];
         if (obj == null) {
-
             if (ef.parameterInfo.HasDefaultValue) {
                 return ef.parameterInfo.DefaultValue;
             }