updated to take cached task

Azure · Apr 9, 2018 · d609fc8 · d609fc8
1 parent a315382
commit d609fc8
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 25 deletions.
diff --git a/src/Microsoft.Azure.WebJobs.Host/Executors/DefaultStorageAccountProvider.cs b/src/Microsoft.Azure.WebJobs.Host/Executors/DefaultStorageAccountProvider.cs
@@ -16,7 +16,7 @@ internal class DefaultStorageAccountProvider : IStorageAccountProvider
         private readonly IStorageCredentialsValidator _storageCredentialsValidator;
         private readonly IStorageAccountParser _storageAccountParser;
         private readonly IServiceProvider _services;
-        private readonly ConcurrentDictionary<string, IStorageAccount> _accounts = new ConcurrentDictionary<string, IStorageAccount>();
+        private readonly ConcurrentDictionary<string, Task<IStorageAccount>> _accounts = new ConcurrentDictionary<string, Task<IStorageAccount>>();
 
         private IStorageAccount _dashboardAccount;
         private bool _dashboardAccountSet;
@@ -188,17 +188,8 @@ private async Task<IStorageAccount> CreateAndValidateAccountAsync(string connect
             return account;
         }
 
-        public async Task<IStorageAccount> TryGetAccountAsync(string connectionStringName, CancellationToken cancellationToken)
-        {
-            IStorageAccount account;
-            if (!_accounts.TryGetValue(connectionStringName, out account))
-            {
-                // in rare cases createAndValidateAccountAsync could be called multiple times for the same account
-                account = await CreateAndValidateAccountAsync(connectionStringName, cancellationToken);
-                _accounts.AddOrUpdate(connectionStringName, (cs) => account, (cs, a) => account);
-            }
-            return account;
-        }
+        public Task<IStorageAccount> TryGetAccountAsync(string connectionStringName, CancellationToken cancellationToken) =>
+            _accounts.GetOrAdd(connectionStringName, s => CreateAndValidateAccountAsync(s, cancellationToken));
 
         private IStorageAccount ParseAccount(string connectionStringName)
         {

diff --git a/src/Microsoft.Azure.WebJobs.Host/Queues/Bindings/QueueBindingProvider.cs b/src/Microsoft.Azure.WebJobs.Host/Queues/Bindings/QueueBindingProvider.cs
@@ -108,9 +108,8 @@ private ParameterDescriptor ToWriteParameterDescriptorForCollector(QueueAttribut
 
             private ParameterDescriptor ToParameterDescriptorForCollector(QueueAttribute attr, ParameterInfo parameter, INameResolver nameResolver, FileAccess access)
             {
-                Task<IStorageAccount> t = Task.Run(() =>
-                    _accountProvider.GetStorageAccountAsync(attr, CancellationToken.None, nameResolver));
-                IStorageAccount account = t.GetAwaiter().GetResult();
+                // Avoid using the sync over async pattern (Async().GetAwaiter().GetResult()) whenever possible
+                IStorageAccount account = _accountProvider.GetStorageAccountAsync(attr, CancellationToken.None, nameResolver).GetAwaiter().GetResult();
 
                 string accountName = account.Credentials.AccountName;
 
@@ -184,7 +183,19 @@ private IAsyncCollector<IStorageQueueMessage> BuildFromQueueAttribute(QueueAttri
 
             internal IStorageQueue GetQueue(QueueAttribute attrResolved)
             {
-                var account = Task.Run(() => _accountProvider.GetStorageAccountAsync(attrResolved, CancellationToken.None)).GetAwaiter().GetResult();
+                // Avoid using the sync over async pattern (Async().GetAwaiter().GetResult()) whenever possible
+                var account = _accountProvider.GetStorageAccountAsync(attrResolved, CancellationToken.None).GetAwaiter().GetResult();
+                return GetQueue(attrResolved, account);
+            }
+
+            internal async Task<IStorageQueue> GetQueueAsync(QueueAttribute attrResolved)
+            {
+                var account = await _accountProvider.GetStorageAccountAsync(attrResolved, CancellationToken.None);
+                return GetQueue(attrResolved, account);
+            }
+
+            internal static IStorageQueue GetQueue(QueueAttribute attrResolved, IStorageAccount account)
+            {
                 var client = account.CreateQueueClient();
 
                 string queueName = attrResolved.QueueName.ToLowerInvariant();
@@ -209,7 +220,7 @@ async Task<IStorageQueue> IAsyncConverter<QueueAttribute, IStorageQueue>.Convert
                 QueueAttribute attrResolved,
                 CancellationToken cancellation)
             {
-                IStorageQueue queue = _bindingProvider.GetQueue(attrResolved);
+                IStorageQueue queue = await _bindingProvider.GetQueueAsync(attrResolved);
                 await queue.CreateIfNotExistsAsync(CancellationToken.None);
                 return queue;
             }
@@ -258,4 +269,4 @@ public QueueAsyncCollector(IStorageQueue queue, IMessageEnqueuedWatcher messageE
             }
         }
     }
-}
+}
diff --git a/src/Microsoft.Azure.WebJobs.Host/Tables/TableExtension.cs b/src/Microsoft.Azure.WebJobs.Host/Tables/TableExtension.cs
@@ -68,16 +68,27 @@ public void Initialize(ExtensionConfigContext context)
         // Get the storage table from the attribute.
         private IStorageTable GetTable(TableAttribute attribute)
         {
-            var account = Task.Run(() => this._accountProvider.GetStorageAccountAsync(attribute, CancellationToken.None)).GetAwaiter().GetResult();
+            // Avoid using the sync over async pattern (Async().GetAwaiter().GetResult()) whenever possible
+            var account = this._accountProvider.GetStorageAccountAsync(attribute, CancellationToken.None).GetAwaiter().GetResult();
+            return GetTable(attribute, account);
+        }
+
+        private async Task<IStorageTable> GetTableAsync(TableAttribute attribute)
+        {
+            var account = await this._accountProvider.GetStorageAccountAsync(attribute, CancellationToken.None);
+            return GetTable(attribute, account);
+        }
+
+        private static IStorageTable GetTable(TableAttribute attribute, IStorageAccount account)
+        {
             var tableClient = account.CreateTableClient();
             return tableClient.GetTableReference(attribute.TableName);
         }
 
         private ParameterDescriptor ToParameterDescriptorForCollector(TableAttribute attribute, ParameterInfo parameter, INameResolver nameResolver)
         {
-            Task<IStorageAccount> t = Task.Run(() =>
-                _accountProvider.GetStorageAccountAsync(attribute, CancellationToken.None, nameResolver));
-            IStorageAccount account = t.GetAwaiter().GetResult();
+            // Avoid using the sync over async pattern (Async().GetAwaiter().GetResult()) whenever possible
+            IStorageAccount account = _accountProvider.GetStorageAccountAsync(attribute, CancellationToken.None, nameResolver).GetAwaiter().GetResult();
             string accountName = account.Credentials.AccountName;
 
             return new TableParameterDescriptor
@@ -224,15 +235,15 @@ IStorageTable IConverter<TableAttribute, IStorageTable>.Convert(TableAttribute a
 
             async Task<CloudTable> IAsyncConverter<TableAttribute, CloudTable>.ConvertAsync(TableAttribute attribute, CancellationToken cancellation)
             {
-                var table = _bindingProvider.GetTable(attribute);
+                var table = await _bindingProvider.GetTableAsync(attribute);
                 await table.CreateIfNotExistsAsync(CancellationToken.None);
 
                 return table.SdkObject;
             }
 
             async Task<JObject> IAsyncConverter<TableAttribute, JObject>.ConvertAsync(TableAttribute attribute, CancellationToken cancellation)
             {
-                var table = _bindingProvider.GetTable(attribute);
+                var table = await _bindingProvider.GetTableAsync(attribute);
 
                 IStorageTableOperation retrieve = table.CreateRetrieveOperation<DynamicTableEntity>(
                   attribute.PartitionKey, attribute.RowKey);
@@ -252,7 +263,7 @@ async Task<JObject> IAsyncConverter<TableAttribute, JObject>.ConvertAsync(TableA
             // Used as an alternative to binding to IQueryable.
             async Task<JArray> IAsyncConverter<TableAttribute, JArray>.ConvertAsync(TableAttribute attribute, CancellationToken cancellation)
             {
-                var table = _bindingProvider.GetTable(attribute).SdkObject;
+                var table = (await _bindingProvider.GetTableAsync(attribute)).SdkObject;
 
                 string finalQuery = attribute.Filter;
                 if (!string.IsNullOrEmpty(attribute.PartitionKey))

diff --git a/test/Microsoft.Azure.WebJobs.Host.EndToEndTests/AsyncCancellationEndToEndTests.cs b/test/Microsoft.Azure.WebJobs.Host.EndToEndTests/AsyncCancellationEndToEndTests.cs
@@ -22,13 +22,17 @@ public class AsyncCancellationEndToEndTests : IDisposable
         private static EventWaitHandle _functionStarted;
         private static EventWaitHandle _functionCompleted;
         private static bool _tokenCancelled;
+        private static SynchronizationContext _oldContext;
 
         private readonly CloudStorageAccount _storageAccount;
         private readonly RandomNameResolver _resolver;
         private readonly JobHostConfiguration _hostConfiguration;
 
         public AsyncCancellationEndToEndTests()
         {
+            // Run tests in multithreaded environment
+            _oldContext = SynchronizationContext.Current;
+            SynchronizationContext.SetSynchronizationContext(null);
             _resolver = new RandomNameResolver();
 
             _hostConfiguration = new JobHostConfiguration()
@@ -58,6 +62,7 @@ public void Dispose()
                     testQueue.Delete();
                 }
             }
+            SynchronizationContext.SetSynchronizationContext(_oldContext);
         }