Address #2539 by caching auth token

BASE/Test/Microsoft.ApplicationInsights.Test/Microsoft.ApplicationInsights.Tests/Extensibility/Implementation/Authentication/CountingMockCredential.cs

@@ -0,0 +1,27 @@
+#if !NET452 && !NET46
+namespace Microsoft.ApplicationInsights.TestFramework.Extensibility.Implementation.Authentication
+{
+    using System.Threading;
+
+    using Azure.Core;
+
+    /// <summary>
+    /// A <see cref="MockCredential"/> that counts the number of calls to <see cref="GetToken(Azure.Core.TokenRequestContext, System.Threading.CancellationToken)"/>.
+    /// </summary>
+    public class CountingMockCredential : MockCredential
+    {
+        private int getTokenCallCount;
+
+        /// <summary>
+        /// Gets or sets the call count.
+        /// </summary>
+        public int GetTokenCallCount { get => getTokenCallCount; private set => getTokenCallCount = value; }
+
+        public override AccessToken GetToken(TokenRequestContext requestContext, CancellationToken cancellationToken)
+        {
+            Interlocked.Increment(ref getTokenCallCount);
+            return base.GetToken(requestContext, cancellationToken);
+        }
+    }
+}
+#endif

BASE/Test/Microsoft.ApplicationInsights.Test/Microsoft.ApplicationInsights.Tests/Extensibility/Implementation/Authentication/TransmissionCredentialEnvelopeTests.cs

@@ -88,6 +88,43 @@ public async Task VerifyTransmissionSendAsync_WithCredential_SetsAuthHeader()
                 var result = await transmission.SendAsync();
             }
         }
+
+        [TestMethod]
+        public async Task VerifyTransmissionSendAsync_WithCachedCredential_ReusesCachedToken()
+        {
+            var mockCredential = new CountingMockCredential();
+            var credentialEnvelope = new CachedReflectionCredentialEnvelope(mockCredential);
+            var authToken = credentialEnvelope.GetToken();
+
+            Assert.AreEqual(1, mockCredential.GetTokenCallCount);
+
+            var handler = new HandlerForFakeHttpClient
+            {
+                InnerHandler = new HttpClientHandler(),
+                OnSendAsync = (req, cancellationToken) =>
+                {
+                    // VALIDATE
+                    Assert.AreEqual(AuthConstants.AuthorizationTokenPrefix.Trim(), req.Headers.Authorization.Scheme);
+                    Assert.AreEqual(authToken.Token, req.Headers.Authorization.Parameter);
+                    Assert.AreEqual(1, mockCredential.GetTokenCallCount, "Expected the token to be cached, thereby avoiding a second call into the token credential.");
+                    
+                    return Task.FromResult<HttpResponseMessage>(new HttpResponseMessage());
+                }
+            };
+
+            using (var fakeHttpClient = new HttpClient(handler))
+            {
+                var expectedContentType = "content/type";
+                var expectedContentEncoding = "contentEncoding";
+                var items = new List<ITelemetry> { new EventTelemetry() };
+
+                // Instantiate Transmission with the mock HttpClient
+                var transmission = new Transmission(testUri, new byte[] { 1, 2, 3, 4, 5 }, fakeHttpClient, expectedContentType, expectedContentEncoding);
+                transmission.CredentialEnvelope = credentialEnvelope;
+
+                var result = await transmission.SendAsync();
+            }
+        }
     }
 }
 #endif

BASE/src/Microsoft.ApplicationInsights/Extensibility/Implementation/Authentication/CachedReflectionCredentialEnvelope.cs

@@ -0,0 +1,111 @@
+namespace Microsoft.ApplicationInsights.Extensibility.Implementation.Authentication
+{
+    using System;
+    using System.Threading;
+    using System.Threading.Tasks;
+
+    /// <summary>
+    /// This is a version of <see cref="ReflectionCredentialEnvelope"/> that caches and reuses
+    /// the auth token for most of its lifetime, thereby preventing repeated calls to fetch
+    /// new tokens.
+    /// </summary>
+    internal sealed class CachedReflectionCredentialEnvelope : ReflectionCredentialEnvelope
+    {
+        /// <summary>
+        /// The token refresh interval. This is the minimum lifetime required
+        /// for the auth token. A cached token can be re-used if its remaining
+        /// lifetime is at least as long as this interval.
+        /// </summary>
+        private readonly TimeSpan tokenRefreshOffset;
+
+        /// <summary>
+        /// The cached token, if any.
+        /// </summary>
+        private AuthToken? cachedToken;
+
+        /// <summary>
+        /// Create an instance of <see cref="CachedReflectionCredentialEnvelope"/>.
+        /// </summary>
+        /// <param name="tokenCredential">An instance of Azure.Core.TokenCredential.</param>
+        /// <remarks>
+        /// The default 5 minute token refresh interval matches the default
+        /// tokenRefreshOffset in Azure.Core's BearerTokenAuthenticationPolicy. It's
+        /// reasonable to assume that callers will use the auth token to make an
+        /// authenticated call well within 5 minutes of calling GetToken.
+        /// </remarks>
+        public CachedReflectionCredentialEnvelope(object tokenCredential) : this(tokenCredential, TimeSpan.FromMinutes(5))
+        {
+        }
+
+        /// <summary>
+        /// Create an instance of <see cref="CachedReflectionCredentialEnvelope"/>.
+        /// </summary>
+        /// <param name="tokenCredential">An instance of Azure.Core.TokenCredential.</param>
+        /// <param name="tokenRefreshOffset">The remaining lifetime allowed before a cached token must be refreshed.</param>
+        public CachedReflectionCredentialEnvelope(object tokenCredential, TimeSpan tokenRefreshOffset) : base(tokenCredential)
+        {
+            this.tokenRefreshOffset = tokenRefreshOffset;
+        }
+
+        /// <summary>
+        /// Gets an Azure.Core.AccessToken.
+        /// </summary>
+        /// <remarks>
+        /// Whomever uses this MUST verify that it's called within <see cref="SdkInternalOperationsMonitor.Enter"/> otherwise dependency calls will be tracked.
+        /// </remarks>
+        /// <param name="cancellationToken">The System.Threading.CancellationToken to use.</param>
+        /// <returns>A valid Azure.Core.AccessToken.</returns>
+        public override AuthToken GetToken(CancellationToken cancellationToken = default)
+        {
+            if (TryUseCachedToken(out AuthToken authToken))
+            {
+                return authToken;
+            }
+
+            return (this.cachedToken = base.GetToken(cancellationToken)).Value;
+        }
+
+        /// <summary>
+        /// Gets an Azure.Core.AccessToken.
+        /// </summary>
+        /// <remarks>
+        /// Whomever uses this MUST verify that it's called within <see cref="SdkInternalOperationsMonitor.Enter"/> otherwise dependency calls will be tracked.
+        /// </remarks>
+        /// <param name="cancellationToken">The System.Threading.CancellationToken to use.</param>
+        /// <returns>A valid Azure.Core.AccessToken.</returns>
+        public override async Task<AuthToken> GetTokenAsync(CancellationToken cancellationToken = default)
+        {
+            if (TryUseCachedToken(out AuthToken authToken))
+            {
+                return authToken;
+            }
+
+            return (this.cachedToken = await base.GetTokenAsync(cancellationToken)).Value;
+        }
+
+        /// <summary>
+        /// Check whether we can use the cached authentication token because its remaining
+        /// lifetime is longer than the minimum required lifetime (5 minutes).
+        /// </summary>
+        /// <param name="cachedToken">On success, the value of the cached token.</param>
+        /// <returns>True if the cached token can be used. False otherwise.</returns>
+        private bool TryUseCachedToken(out AuthToken cachedToken)
+        {
+            AuthToken? token = this.cachedToken;
+            if (token.HasValue)
+            {
+                TimeSpan timeRemaining = token.Value.ExpiresOn - DateTimeOffset.UtcNow;
+                if (timeRemaining >= tokenRefreshOffset)
+                {
+                    cachedToken = token.Value;
+                    return true;
+                }
+
+                // Cached token must be refreshed.
+            }
+
+            cachedToken = default;
+            return false;
+        }
+    }
+}

BASE/src/Microsoft.ApplicationInsights/Extensibility/TelemetryConfiguration.cs

@@ -426,7 +426,7 @@ public void Dispose()
         /// <exception cref="ArgumentException">An ArgumentException is thrown if the provided object does not inherit Azure.Core.TokenCredential.</exception>
         public void SetAzureTokenCredential(object tokenCredential)
         {
-            this.CredentialEnvelope = new ReflectionCredentialEnvelope(tokenCredential);
+            this.CredentialEnvelope = new CachedReflectionCredentialEnvelope(tokenCredential);
             this.SetTelemetryChannelCredentialEnvelope();
 
             // Update Ingestion Endpoint.

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## VNext
 - [Upgrade System.Diagnostics.PerformanceCounter to version 6.0.0 to address CVE-2021-24112](https://github.com/microsoft/ApplicationInsights-dotnet/pull/2707)
+- [Cache authentication token to improve performance when using AAD authentication](https://github.com/microsoft/ApplicationInsights-dotnet/issues/2539)
 
 ## Version 2.22.0-beta1
 - Update endpoint redirect header name for QuickPulse module to v2