Merged PR 747603: Allow BuildXL to use multiple Blob L3 shards

Author: jbayardo

Public/Src/Cache/VerticalStore/MemoizationStoreAdapter/BlobCacheFactory.cs

@@ -6,15 +6,16 @@
 using System.ComponentModel;
 using System.Diagnostics.CodeAnalysis;
 using System.Diagnostics.ContractsLight;
+using System.Linq;
 using System.Threading.Tasks;
 using BuildXL.Cache.ContentStore.Distributed.Blob;
+using BuildXL.Cache.ContentStore.Interfaces.Auth;
 using BuildXL.Cache.Interfaces;
 using BuildXL.Cache.MemoizationStore.Distributed.Stores;
-using BuildXL.Utilities.Core;
+using BuildXL.Utilities.Collections;
 using BuildXL.Utilities.Configuration;
+using BuildXL.Utilities.Core;
 using AbsolutePath = BuildXL.Cache.ContentStore.Interfaces.FileSystem.AbsolutePath;
-using System.Security.Principal;
-using BuildXL.Cache.ContentStore.Interfaces.Auth;
 
 namespace BuildXL.Cache.MemoizationStoreAdapter
 {
@@ -28,7 +29,7 @@ namespace BuildXL.Cache.MemoizationStoreAdapter
     /// Current limitations while we flesh things out:
     /// 1) APIs around tracking named sessions are not implemented
     /// </remarks>
-    public partial class BlobCacheFactory : ICacheFactory
+    public class BlobCacheFactory : ICacheFactory
     {
         /// <summary>
         /// Inheritable configuration settings for cache factories that wish to configure a connection to a blob cache
@@ -76,6 +77,15 @@ public abstract class BlobCacheConfig
             /// </remarks>
             [DefaultValue(0)]
             public int RetentionPolicyInDays { get; set; }
+
+            /// <nodoc />
+            [DefaultValue("default")]
+            public string Universe { get; set; }
+
+            /// <nodoc />
+            [DefaultValue("default")]
+            public string Namespace { get; set; }
+
         }
 
         /// <summary>
@@ -89,14 +99,6 @@ public sealed class Config : BlobCacheConfig
             [DefaultValue(typeof(CacheId))]
             public CacheId CacheId { get; set; }
 
-            /// <nodoc />
-            [DefaultValue("default")]
-            public string Universe { get; set; }
-
-            /// <nodoc />
-            [DefaultValue("default")]
-            public string Namespace { get; set; }
-
             /// <summary>
             /// Path to the log file for the cache.
             /// </summary>
@@ -157,7 +159,7 @@ public async Task<Possible<ICache, Failure>> InitializeCacheAsync(Config configu
                     logger: new DisposeLogger(() => new EtwFileLog(logPath.Path, configuration.CacheId), configuration.LogFlushIntervalSeconds),
                     statsFile: new AbsolutePath(logPath.Path + ".stats"),
                     precedingStateDegradationFailures: failures);
-                
+
                 var startupResult = await cache.StartupAsync();
                 if (!startupResult.Succeeded)
                 {
@@ -172,14 +174,12 @@ public async Task<Possible<ICache, Failure>> InitializeCacheAsync(Config configu
             }
         }
 
-        private static MemoizationStore.Interfaces.Caches.ICache CreateCache(Config configuration)
+        internal static MemoizationStore.Interfaces.Caches.IFullCache CreateCache(BlobCacheConfig configuration)
         {
-            IAzureStorageCredentials credentials = GetAzureCredentialsFromBlobFactoryConfig(configuration);
-
-            var accountName = BlobCacheStorageAccountName.Parse(credentials.GetAccountName());
+            var credentials = LoadAzureCredentials(configuration);
 
             var factoryConfiguration = new AzureBlobStorageCacheFactory.Configuration(
-                ShardingScheme: new ShardingScheme(ShardingAlgorithm.SingleShard, new List<BlobCacheStorageAccountName> { accountName }),
+                ShardingScheme: new ShardingScheme(ShardingAlgorithm.JumpHash, credentials.Keys.ToList()),
                 Universe: configuration.Universe,
                 Namespace: configuration.Namespace,
                 RetentionPolicyInDays: configuration.RetentionPolicyInDays);
@@ -188,14 +188,22 @@ private static MemoizationStore.Interfaces.Caches.ICache CreateCache(Config conf
         }
 
         /// <nodoc />
-        internal static IAzureStorageCredentials GetAzureCredentialsFromBlobFactoryConfig(BlobCacheConfig configuration)
+        internal static Dictionary<BlobCacheStorageAccountName, IAzureStorageCredentials> LoadAzureCredentials(BlobCacheConfig configuration)
         {
-            IAzureStorageCredentials credentials;
+            var credentials = new Dictionary<BlobCacheStorageAccountName, IAzureStorageCredentials>();
             var connectionString = Environment.GetEnvironmentVariable(configuration.ConnectionStringEnvironmentVariableName);
 
             if (!string.IsNullOrEmpty(connectionString))
             {
-                credentials = new SecretBasedAzureStorageCredentials(connectionString);
+                credentials.AddRange(
+                    connectionString.Split(' ')
+                        .Select(
+                            secret =>
+                            {
+                                var credential = new SecretBasedAzureStorageCredentials(secret.Trim());
+                                var accountName = BlobCacheStorageAccountName.Parse(credential.GetAccountName());
+                                return new KeyValuePair<BlobCacheStorageAccountName, IAzureStorageCredentials>(accountName, credential);
+                            }));
             }
             else if (configuration.ManagedIdentityId is not null && configuration.StorageAccountEndpoint is not null)
             {
@@ -207,7 +215,8 @@ internal static IAzureStorageCredentials GetAzureCredentialsFromBlobFactoryConfi
                     throw new InvalidOperationException($"'{configuration.StorageAccountEndpoint}' does not represent a valid URI.");
                 }
 
-                credentials = new ManagedIdentityAzureStorageCredentials(configuration.ManagedIdentityId, uri);
+                var credential = new ManagedIdentityAzureStorageCredentials(configuration.ManagedIdentityId, uri);
+                credentials.Add(BlobCacheStorageAccountName.Parse(credential.GetAccountName()), credential);
             }
             else
             {

Public/Src/Cache/VerticalStore/MemoizationStoreAdapter/EphemeralCacheFactory.cs

@@ -6,20 +6,18 @@
 using System.ComponentModel;
 using System.Diagnostics.CodeAnalysis;
 using System.Diagnostics.ContractsLight;
+using System.Linq;
 using System.Threading.Tasks;
-using BuildXL.Cache.Interfaces;
-using BuildXL.Cache.MemoizationStore.Distributed.Stores;
-using BuildXL.Utilities.Core;
-using BuildXL.Utilities.Configuration;
-using AbsolutePath = BuildXL.Cache.ContentStore.Interfaces.FileSystem.AbsolutePath;
 using BuildXL.Cache.ContentStore.Distributed;
 using BuildXL.Cache.ContentStore.Distributed.Blob;
 using BuildXL.Cache.ContentStore.Grpc;
-using BuildXL.Cache.ContentStore.Interfaces.Auth;
+using BuildXL.Cache.ContentStore.Interfaces.Logging;
 using BuildXL.Cache.ContentStore.Interfaces.Tracing;
-using BuildXL.Cache.ContentStore.Logging;
 using BuildXL.Cache.ContentStore.Tracing.Internal;
-using BuildXL.Cache.ContentStore.Interfaces.Logging;
+using BuildXL.Cache.Interfaces;
+using BuildXL.Utilities.Configuration;
+using BuildXL.Utilities.Core;
+using AbsolutePath = BuildXL.Cache.ContentStore.Interfaces.FileSystem.AbsolutePath;
 
 namespace BuildXL.Cache.MemoizationStoreAdapter;
 
@@ -39,18 +37,6 @@ public sealed class FactoryConfiguration : BlobCacheFactory.BlobCacheConfig
         [DefaultValue(typeof(CacheId))]
         public CacheId CacheId { get; set; }
 
-        /// <nodoc />
-        [DefaultValue("EphemeralCacheConnectionString")]
-        public string ManagementConnectionStringEnvironmentVariableName { get; set; }
-
-        /// <nodoc />
-        [DefaultValue("default")]
-        public string Universe { get; set; }
-
-        /// <nodoc />
-        [DefaultValue("default")]
-        public string Namespace { get; set; }
-
         /// <summary>
         /// Path to the log file for the cache.
         /// </summary>
@@ -160,12 +146,9 @@ public async Task<Possible<ICache, Failure>> InitializeCacheAsync(FactoryConfigu
 
         if (configuration.DatacenterWide)
         {
-            var connectionString = Environment.GetEnvironmentVariable(configuration.ManagementConnectionStringEnvironmentVariableName);
-            if (string.IsNullOrEmpty(connectionString))
-            {
-                throw new InvalidOperationException($"Can't find a connection string in environment variable '{configuration.ManagementConnectionStringEnvironmentVariableName}'.");
-            }
-            var credentials = new SecretBasedAzureStorageCredentials(connectionString);
+            var accounts = BlobCacheFactory.LoadAzureCredentials(configuration);
+            var sorted = ShardingScheme.SortAccounts(accounts.Keys.ToList());
+            var credentials = accounts[sorted.First()];
 
             factoryConfiguration = new ContentStore.Distributed.Ephemeral.EphemeralCacheFactory.DatacenterWideCacheConfiguration()
             {
@@ -188,24 +171,10 @@ public async Task<Possible<ICache, Failure>> InitializeCacheAsync(FactoryConfigu
             };
         }
 
-        var persistentCache = CreateBlobCache(configuration);
+        var persistentCache = BlobCacheFactory.CreateCache(configuration);
         return await Cache.ContentStore.Distributed.Ephemeral.EphemeralCacheFactory.CreateAsync(context, factoryConfiguration, persistentCache);
     }
 
-    private static MemoizationStore.Interfaces.Caches.IFullCache CreateBlobCache(FactoryConfiguration configuration)
-    {
-        var credentials = BlobCacheFactory.GetAzureCredentialsFromBlobFactoryConfig(configuration);
-        var accountName = BlobCacheStorageAccountName.Parse(credentials.GetAccountName());
-
-        var factoryConfiguration = new AzureBlobStorageCacheFactory.Configuration(
-            ShardingScheme: new ShardingScheme(ShardingAlgorithm.SingleShard, new List<BlobCacheStorageAccountName> { accountName }),
-            Universe: configuration.Universe,
-            Namespace: configuration.Namespace,
-            RetentionPolicyInDays: configuration.RetentionPolicyInDays);
-
-        return AzureBlobStorageCacheFactory.Create(factoryConfiguration, new StaticBlobCacheSecretsProvider(credentials));
-    }
-
     /// <inheritdoc />
     public IEnumerable<Failure> ValidateConfiguration(ICacheConfigData cacheData)
     {