Azure · Pavel Krymets (pakrym) · Oct 22, 2019 · Oct 22, 2019 · Oct 22, 2019 · Oct 22, 2019
diff --git a/sdk/core/Azure.Core/src/TokenCredential.cs b/sdk/core/Azure.Core/src/TokenCredential.cs
@@ -17,7 +17,7 @@ public abstract class TokenCredential
         /// <param name="requestContext">The <see cref="TokenRequestContext"/> with authentication information.</param>
         /// <param name="cancellationToken">The <see cref="CancellationToken"/> to use.</param>
         /// <returns>A valid <see cref="AccessToken"/>.</returns>
-        public abstract Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken);
+        public abstract ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken);
 
         /// <summary>
         /// Gets an <see cref="AccessToken"/> for the specified set of scopes.

diff --git a/sdk/core/Azure.Core/tests/TestFramework/DiagnosticScopeValidatingInterceptor.cs b/sdk/core/Azure.Core/tests/TestFramework/DiagnosticScopeValidatingInterceptor.cs
@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
+using System.Reflection;
 using System.Threading.Tasks;
 using Azure.Core.Tests;
 using Castle.DynamicProxy;
@@ -36,10 +37,26 @@ public void Intercept(IInvocation invocation)
                     return;
                 }
 
-                var result = (Task)invocation.ReturnValue;
                 try
                 {
-                    result.GetAwaiter().GetResult();
+                    object returnValue = invocation.ReturnValue;
+                    if (returnValue is Task t)
+                    {
+                        t.GetAwaiter().GetResult();
+                    }
+                    else
+                    {
+                        // Await ValueTask
+                        Type returnType = returnValue.GetType();
+                        MethodInfo getAwaiterMethod = returnType.GetMethod("GetAwaiter", BindingFlags.Instance | BindingFlags.Public);
+                        MethodInfo getResultMethod = getAwaiterMethod.ReturnType.GetMethod("GetResult", BindingFlags.Instance | BindingFlags.Public);
+
+                        getResultMethod.Invoke(
+                            getAwaiterMethod.Invoke(returnValue, Array.Empty<object>()),
+                            Array.Empty<object>());
+
+                    }
+
                     expectedEvents.Add(expectedEventPrefix + ".Stop");
                 }
                 catch (Exception ex)

diff --git a/sdk/core/Azure.Core/tests/TestFramework/TestRecording.cs b/sdk/core/Azure.Core/tests/TestFramework/TestRecording.cs
@@ -261,9 +261,9 @@ public void DisableIdReuse()
 
         private class TestCredential : TokenCredential
         {
-            public override Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken)
+            public override ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken)
             {
-                return Task.FromResult(GetToken(requestContext, cancellationToken));
+                return new ValueTask<AccessToken>(GetToken(requestContext, cancellationToken));
             }
 
             public override AccessToken GetToken(TokenRequestContext requestContext, CancellationToken cancellationToken)

diff --git a/sdk/core/Azure.Core/tests/TestFramework/UseSyncMethodsInterceptor.cs b/sdk/core/Azure.Core/tests/TestFramework/UseSyncMethodsInterceptor.cs
@@ -87,7 +87,7 @@ public void Intercept(IInvocation invocation)
                 }
                 else
                 {
-                    invocation.ReturnValue = _taskFromResultMethod.MakeGenericMethod(returnType).Invoke(null, new[] { result });
+                    SetAsyncResult(invocation, returnType, result);
                 }
             }
             catch (TargetInvocationException exception)
@@ -98,9 +98,51 @@ public void Intercept(IInvocation invocation)
                 }
                 else
                 {
-                    invocation.ReturnValue = _taskFromExceptionMethod.MakeGenericMethod(methodInfo.ReturnType).Invoke(null, new[] { exception.InnerException });
+                    SetAsyncException(invocation, returnType, exception.InnerException);
+                }
+            }
+        }
+
+        private void SetAsyncResult(IInvocation invocation, Type returnType, object result)
+        {
+            Type methodReturnType = invocation.Method.ReturnType;
+            if (methodReturnType.IsGenericType)
+            {
+                if (methodReturnType.GetGenericTypeDefinition() == typeof(Task<>))
+                {
+                    invocation.ReturnValue = _taskFromResultMethod.MakeGenericMethod(returnType).Invoke(null, new[] { result });
+                    return;
+                }
+                if (methodReturnType.GetGenericTypeDefinition() == typeof(ValueTask<>))
+                {
+                    invocation.ReturnValue = Activator.CreateInstance(typeof(ValueTask<>).MakeGenericType(returnType), result);
+                    return;
                 }
             }
+
+            throw new NotSupportedException();
+        }
+
+        private void SetAsyncException(IInvocation invocation, Type returnType, Exception result)
+        {
+            Type methodReturnType = invocation.Method.ReturnType;
+            if (methodReturnType.IsGenericType)
+            {
+                if (methodReturnType.GetGenericTypeDefinition() == typeof(Task<>))
+                {
+                    invocation.ReturnValue = _taskFromExceptionMethod.MakeGenericMethod(returnType).Invoke(null, new[] { result });
+                    return;
+                }
+
+                if (methodReturnType.GetGenericTypeDefinition() == typeof(ValueTask<>))
+                {
+                    var task = _taskFromExceptionMethod.MakeGenericMethod(returnType).Invoke(null, new[] { result });
+                    invocation.ReturnValue = Activator.CreateInstance(typeof(ValueTask<>).MakeGenericType(returnType), task);
+                    return;
+                }
+            }
+
+            throw new NotSupportedException();
         }
 
         private static MethodInfo GetMethod(IInvocation invocation, string nonAsyncMethodName, Type[] types) =>

diff --git a/sdk/eventhub/Azure.Messaging.EventHubs/src/Authorization/EventHubSharedKeyCredential.cs b/sdk/eventhub/Azure.Messaging.EventHubs/src/Authorization/EventHubSharedKeyCredential.cs
@@ -70,7 +70,7 @@ public EventHubSharedKeyCredential(string sharedAccessKeyName,
         ///
         /// <returns>The token representing the shared access signature for this credential.</returns>
         ///
-        public override Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken) => throw new InvalidOperationException(Resources.SharedKeyCredentialCannotGenerateTokens);
+        public override ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken) => throw new InvalidOperationException(Resources.SharedKeyCredentialCannotGenerateTokens);
 
         /// <summary>
         /// Coverts to shared access signature credential.

diff --git a/sdk/eventhub/Azure.Messaging.EventHubs/src/Authorization/EventHubTokenCredential.cs b/sdk/eventhub/Azure.Messaging.EventHubs/src/Authorization/EventHubTokenCredential.cs
@@ -81,7 +81,7 @@ public override AccessToken GetToken(TokenRequestContext requestContext,
         ///
         /// <returns>The token representing the shared access signature for this credential.</returns>
         ///
-        public override Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext,
+        public override ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext,
                                                         CancellationToken cancellationToken) => Credential.GetTokenAsync(requestContext, cancellationToken);
     }
 }
diff --git a/sdk/eventhub/Azure.Messaging.EventHubs/src/Authorization/SharedAccessSignatureCredential.cs b/sdk/eventhub/Azure.Messaging.EventHubs/src/Authorization/SharedAccessSignatureCredential.cs
@@ -58,7 +58,7 @@ public override AccessToken GetToken(TokenRequestContext requestContext,
         ///
         /// <returns>The token representing the shared access signature for this credential.</returns>
         ///
-        public override Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext,
-                                                        CancellationToken cancellationToken) => Task.FromResult(new AccessToken(SharedAccessSignature.Value, SharedAccessSignature.SignatureExpiration));
+        public override ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext,
+                                                        CancellationToken cancellationToken) => new ValueTask<AccessToken>(new AccessToken(SharedAccessSignature.Value, SharedAccessSignature.SignatureExpiration));
     }
 }
diff --git a/sdk/eventhub/Azure.Messaging.EventHubs/tests/Amqp/AmqpEventHubClientTests.cs b/sdk/eventhub/Azure.Messaging.EventHubs/tests/Amqp/AmqpEventHubClientTests.cs
@@ -232,7 +232,7 @@ public void GetPropertiesAsyncCreatesTheRequest()
 
             mockCredential
                 .Setup(credential => credential.GetTokenAsync(It.IsAny<TokenRequestContext>(), It.Is<CancellationToken>(value => value == cancellationSource.Token)))
-                .Returns(Task.FromResult(new AccessToken(tokenValue, DateTimeOffset.MaxValue)))
+                .Returns(new ValueTask<AccessToken>(new AccessToken(tokenValue, DateTimeOffset.MaxValue)))
                 .Verifiable();
 
             mockConverter
@@ -273,7 +273,7 @@ public void GetPropertiesAsyncRespectsTheRetryPolicy(RetryOptions retryOptions)
 
             mockCredential
                 .Setup(credential => credential.GetTokenAsync(It.IsAny<TokenRequestContext>(), It.Is<CancellationToken>(value => value == cancellationSource.Token)))
-                .Returns(Task.FromResult(new AccessToken(tokenValue, DateTimeOffset.MaxValue)));
+                .Returns(new ValueTask<AccessToken>(new AccessToken(tokenValue, DateTimeOffset.MaxValue)));
 
             mockConverter
                 .Setup(converter => converter.CreateEventHubPropertiesRequest(It.Is<string>(value => value == eventHubName), It.Is<string>(value => value == tokenValue)))
@@ -354,7 +354,7 @@ public void GetPartitionPropertiesAsyncCreatesTheRequest()
 
             mockCredential
                 .Setup(credential => credential.GetTokenAsync(It.IsAny<TokenRequestContext>(), It.Is<CancellationToken>(value => value == cancellationSource.Token)))
-                .Returns(Task.FromResult(new AccessToken(tokenValue, DateTimeOffset.MaxValue)))
+                .Returns(new ValueTask<AccessToken>(new AccessToken(tokenValue, DateTimeOffset.MaxValue)))
                 .Verifiable();
 
             mockConverter
@@ -396,7 +396,7 @@ public void GetPartitionPropertiesAsyncRespectsTheRetryPolicy(RetryOptions retry
 
             mockCredential
                 .Setup(credential => credential.GetTokenAsync(It.IsAny<TokenRequestContext>(), It.Is<CancellationToken>(value => value == cancellationSource.Token)))
-                .Returns(Task.FromResult(new AccessToken(tokenValue, DateTimeOffset.MaxValue)));
+                .Returns(new ValueTask<AccessToken>(new AccessToken(tokenValue, DateTimeOffset.MaxValue)));
 
             mockConverter
                 .Setup(converter => converter.CreatePartitionPropertiesRequest(It.Is<string>(value => value == eventHubName), It.Is<string>(value => value == partitionId), It.Is<string>(value => value == tokenValue)))

diff --git a/sdk/eventhub/Azure.Messaging.EventHubs/tests/Amqp/AmqpEventHubConsumerTests.cs b/sdk/eventhub/Azure.Messaging.EventHubs/tests/Amqp/AmqpEventHubConsumerTests.cs
@@ -271,7 +271,7 @@ public void ReceiveAsyncRespectsTheRetryPolicy(RetryOptions retryOptions)
 
             mockCredential
                 .Setup(credential => credential.GetTokenAsync(It.IsAny<TokenRequestContext>(), It.Is<CancellationToken>(value => value == cancellationSource.Token)))
-                .Returns(Task.FromResult(new AccessToken(tokenValue, DateTimeOffset.MaxValue)));
+                .Returns(new ValueTask<AccessToken>(new AccessToken(tokenValue, DateTimeOffset.MaxValue)));
 
             mockScope
                .Setup(scope => scope.OpenConsumerLinkAsync(

diff --git a/sdk/eventhub/Azure.Messaging.EventHubs/tests/Amqp/CbsTokenProviderTests.cs b/sdk/eventhub/Azure.Messaging.EventHubs/tests/Amqp/CbsTokenProviderTests.cs
@@ -47,7 +47,7 @@ public async Task GetTokenAsyncPassesAlongTheClaims()
 
             mockCredential
                 .Setup(credential => credential.GetTokenAsync(It.Is<TokenRequestContext>(value => value.Scopes == requiredClaims), It.IsAny<CancellationToken>()))
-                .Returns(Task.FromResult(new AccessToken("blah", DateTimeOffset.Parse("2015-10-27T00:00:00Z"))))
+                .Returns(new ValueTask<AccessToken>(new AccessToken("blah", DateTimeOffset.Parse("2015-10-27T00:00:00Z"))))
                 .Verifiable();
 
             await provider.GetTokenAsync(new Uri("http://www.here.com"), "nobody", requiredClaims);
@@ -70,7 +70,7 @@ public async Task GetTokenAsyncPopulatesUsingTheCredentialAccessToken()
 
             mockCredential
                 .Setup(credential => credential.GetTokenAsync(It.IsAny<TokenRequestContext>(), It.IsAny<CancellationToken>()))
-                .Returns(Task.FromResult(new AccessToken(tokenValue, expires)));
+                .Returns(new ValueTask<AccessToken>(new AccessToken(tokenValue, expires)));
 
             CbsToken cbsToken = await provider.GetTokenAsync(new Uri("http://www.here.com"), "nobody", new string[0]);
 
@@ -114,7 +114,7 @@ public async Task GetTokenAsyncSetsTheCorrectTypeForOtherTokens()
 
             mockCredential
                 .Setup(credential => credential.GetTokenAsync(It.IsAny<TokenRequestContext>(), It.IsAny<CancellationToken>()))
-                .Returns(Task.FromResult(new AccessToken(tokenValue, expires)));
+                .Returns(new ValueTask<AccessToken>(new AccessToken(tokenValue, expires)));
 
             CbsToken cbsToken = await provider.GetTokenAsync(new Uri("http://www.here.com"), "nobody", new string[0]);
 

diff --git a/sdk/identity/Azure.Identity/src/AuthorizationCodeCredential.cs b/sdk/identity/Azure.Identity/src/AuthorizationCodeCredential.cs
@@ -87,7 +87,7 @@ public override AccessToken GetToken(TokenRequestContext requestContext, Cancell
         /// <param name="requestContext">The details of the authentication request.</param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
         /// <returns>An <see cref="AccessToken"/> which can be used to authenticate service client calls.</returns>
-        public override async Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
+        public override async ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
         {
             using DiagnosticScope scope = _clientDiagnostics.CreateScope("Azure.Identity.AuthorizationCodeCredential.GetToken");
 

diff --git a/sdk/identity/Azure.Identity/src/ChainedTokenCredential.cs b/sdk/identity/Azure.Identity/src/ChainedTokenCredential.cs
@@ -71,7 +71,7 @@ public override AccessToken GetToken(TokenRequestContext requestContext, Cancell
         /// <param name="requestContext">The details of the authentication request.</param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
         /// <returns>The first non default <see cref="AccessToken"/> returned by the specified sources.  If all credentials in the chain return default a default <see cref="AccessToken"/> is returned.</returns>
-        public override async Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
+        public override async ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
         {
             AccessToken token = new AccessToken();
 

diff --git a/sdk/identity/Azure.Identity/src/ClientCertificateCredential.cs b/sdk/identity/Azure.Identity/src/ClientCertificateCredential.cs
@@ -80,7 +80,7 @@ public override AccessToken GetToken(TokenRequestContext requestContext, Cancell
         /// <param name="requestContext">The details of the authentication request.</param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
         /// <returns>An <see cref="AccessToken"/> which can be used to authenticate service client calls.</returns>
-        public override async Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
+        public override async ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
         {
             return await _client.AuthenticateAsync(TenantId, ClientId, ClientCertificate, requestContext.Scopes, cancellationToken).ConfigureAwait(false);
         }

diff --git a/sdk/identity/Azure.Identity/src/ClientSecretCredential.cs b/sdk/identity/Azure.Identity/src/ClientSecretCredential.cs
@@ -65,7 +65,7 @@ public ClientSecretCredential(string tenantId, string clientId, string clientSec
         /// <param name="requestContext">The details of the authentication request.</param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
         /// <returns>An <see cref="AccessToken"/> which can be used to authenticate service client calls.</returns>
-        public override async Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
+        public override async ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
         {
             return await _client.AuthenticateAsync(TenantId, ClientId, ClientSecret, requestContext.Scopes, cancellationToken).ConfigureAwait(false);
         }

diff --git a/sdk/identity/Azure.Identity/src/DefaultAzureCredential.cs b/sdk/identity/Azure.Identity/src/DefaultAzureCredential.cs
@@ -85,7 +85,7 @@ public override AccessToken GetToken(TokenRequestContext requestContext, Cancell
                 throw new AuthenticationFailedException(CredentialNotFoundMessage);
             }
 
-            public override Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken)
+            public override ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken)
             {
                 throw new AuthenticationFailedException(CredentialNotFoundMessage);
             }

diff --git a/sdk/identity/Azure.Identity/src/DeviceCodeCredential.cs b/sdk/identity/Azure.Identity/src/DeviceCodeCredential.cs
@@ -124,7 +124,7 @@ public override AccessToken GetToken(TokenRequestContext requestContext, Cancell
         /// <param name="requestContext">The details of the authentication request.</param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
         /// <returns>An <see cref="AccessToken"/> which can be used to authenticate service client calls.</returns>
-        public override async Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
+        public override async ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
         {
             using DiagnosticScope scope = _clientDiagnostics.CreateScope("Azure.Identity.DeviceCodeCredential.GetToken");
 

diff --git a/sdk/identity/Azure.Identity/src/EnvironmentCredential.cs b/sdk/identity/Azure.Identity/src/EnvironmentCredential.cs
@@ -83,7 +83,7 @@ public override AccessToken GetToken(TokenRequestContext requestContext, Cancell
         /// <param name="requestContext">The details of the authentication request.</param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
         /// <returns>An <see cref="AccessToken"/> which can be used to authenticate service client calls, or a default <see cref="AccessToken"/>.</returns>
-        public override async Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
+        public override async ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
         {
             return (_credential != null) ? await _credential.GetTokenAsync(requestContext, cancellationToken).ConfigureAwait(false) : default;
         }

diff --git a/sdk/identity/Azure.Identity/src/InteractiveBrowserCredential.cs b/sdk/identity/Azure.Identity/src/InteractiveBrowserCredential.cs
@@ -82,7 +82,7 @@ public override AccessToken GetToken(TokenRequestContext requestContext, Cancell
         /// <param name="requestContext">The details of the authentication request.</param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
         /// <returns>An <see cref="AccessToken"/> which can be used to authenticate service client calls.</returns>
-        public override async Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
+        public override async ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
         {
             if (_account != null)
             {

diff --git a/sdk/identity/Azure.Identity/src/ManagedIdentityCredential.cs b/sdk/identity/Azure.Identity/src/ManagedIdentityCredential.cs
@@ -42,7 +42,7 @@ public ManagedIdentityCredential(string clientId = null, TokenCredentialOptions
         /// <param name="requestContext">The details of the authentication request.</param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
         /// <returns>An <see cref="AccessToken"/> which can be used to authenticate service client calls, or a default <see cref="AccessToken"/> if no managed identity is available.</returns>
-        public override async Task<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
+        public override async ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken = default)
         {
             return await _client.AuthenticateAsync(requestContext.Scopes, _clientId, cancellationToken).ConfigureAwait(false);
         }