AAD: new AuthenticationTransmissionPolicy for retry scenarios.

BASE/Test/ServerTelemetryChannel.Test/TelemetryChannel.Tests/Implementation/TransmissionPolicy/AuthenticationTransmissionPolicyTests.cs

@@ -0,0 +1,118 @@
+namespace Microsoft.ApplicationInsights.WindowsServer.TelemetryChannel.Implementation.TransmissionPolicy
+{
+    using System;
+    using System.Diagnostics.Tracing;
+    using System.Globalization;
+    using System.Linq;
+    using System.Threading;
+
+    using Microsoft.ApplicationInsights.Channel.Implementation;
+    using Microsoft.ApplicationInsights.Extensibility.Implementation;
+    using Microsoft.ApplicationInsights.TestFramework;
+    using Microsoft.ApplicationInsights.WindowsServer.TelemetryChannel.Helpers;
+    using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+    [TestClass]
+    [TestCategory("TransmissionPolicy")]
+    public class AuthenticationTransmissionPolicyTests
+    {
+        [TestMethod]
+        public void Verify400TriggersThrottling() => this.EvaluateIfStatusCodeTriggersThrottling(ResponseStatusCodes.BadRequest);
+
+        [TestMethod]
+        public void Verify401TriggersThrottling() => this.EvaluateIfStatusCodeTriggersThrottling(ResponseStatusCodes.Unauthorized);
+
+        [TestMethod]
+        public void Verify403TriggersThrottling() => this.EvaluateIfStatusCodeTriggersThrottling(ResponseStatusCodes.Forbidden);
+
+        [TestMethod]
+        public void Verify200DoesNotTriggerThrottling() => this.EvaluateIfStatusCodeIgnored(ResponseStatusCodes.Success);
+
+        [TestMethod]
+        public void VerifyOtherDoesNotTriggerThrottling() => this.EvaluateIfStatusCodeIgnored(000);
+
+        private void EvaluateIfStatusCodeTriggersThrottling(int statusCode)
+        {
+            var retryAfterSeconds = BackoffLogicManager.SlotDelayInSeconds;
+            var waitForTheFirstApplyAsync = TimeSpan.FromMilliseconds(100);
+            var waitForTheSecondApplyAsync = TimeSpan.FromMilliseconds(retryAfterSeconds * 1000 + 500);
+
+            // SETUP
+            var transmitter = new StubTransmitterEvalOnApply();
+
+            var policy = new AuthenticationTransmissionPolicy()
+            {
+                Enabled = true,
+            };
+            policy.Initialize(transmitter);
+
+            // ACT
+            transmitter.InvokeTransmissionSentEvent(statusCode);
+
+            // ASSERT: First Handle will trigger Throttle and delay.
+            Assert.IsTrue(transmitter.IsApplyInvoked(waitForTheFirstApplyAsync));
+
+            Assert.AreEqual(0, policy.MaxSenderCapacity);
+            Assert.IsNull(policy.MaxBufferCapacity);
+            Assert.IsNull(policy.MaxStorageCapacity);
+
+            // ASSERT: Throttle expires and policy will be reset.
+            Assert.IsTrue(transmitter.IsApplyInvoked(waitForTheSecondApplyAsync));
+
+            Assert.IsNull(policy.MaxSenderCapacity);
+            Assert.IsNull(policy.MaxBufferCapacity);
+            Assert.IsNull(policy.MaxStorageCapacity);
+        }
+
+        private void EvaluateIfStatusCodeIgnored(int statusCode)
+        {
+            var waitForTheFirstApplyAsync = TimeSpan.FromMilliseconds(100);
+
+            // SETUP
+            var transmitter = new StubTransmitterEvalOnApply();
+
+            var policy = new AuthenticationTransmissionPolicy()
+            {
+                Enabled = true,
+            };
+            policy.Initialize(transmitter);
+
+            // ACT
+            transmitter.InvokeTransmissionSentEvent(statusCode);
+
+            // ASSERT: The Apply event handler should not be called.
+            Assert.IsFalse(transmitter.IsApplyInvoked(waitForTheFirstApplyAsync));
+
+            // ASSERT: Capacities should have default values.
+            Assert.IsNull(policy.MaxSenderCapacity);
+            Assert.IsNull(policy.MaxBufferCapacity);
+            Assert.IsNull(policy.MaxStorageCapacity);
+        }
+
+        private class StubTransmitterEvalOnApply : StubTransmitter
+        {
+            private AutoResetEvent autoResetEvent;
+
+            public StubTransmitterEvalOnApply()
+            {
+                this.autoResetEvent = new AutoResetEvent(false);
+                this.OnApplyPolicies = () => this.autoResetEvent.Set();
+            }
+
+            public void InvokeTransmissionSentEvent(int responseStatusCode)
+            {
+                this.OnTransmissionSent(new TransmissionProcessedEventArgs(
+                    transmission: new StubTransmission(),
+                    exception: null,
+                    response: new HttpWebResponseWrapper()
+                    {
+                        StatusCode = responseStatusCode,
+                        StatusDescription = null,
+                    }
+                ));
+            }
+
+            public bool IsApplyInvoked(TimeSpan timeout) => this.autoResetEvent.WaitOne(timeout);
+        }
+    }
+}

BASE/Test/ServerTelemetryChannel.Test/TelemetryChannel.Tests/Implementation/TransmissionPolicy/ThrottlingTransmissionPolicyTest.cs

@@ -22,67 +22,31 @@ public class HandleTransmissionSentEvent : ErrorHandlingTransmissionPolicyTest
             [TestMethod]
             public void AssertTooManyRequestsStopsSending()
             {
-                this.PositiveTest(ResponseStatusCodes.ResponseCodeTooManyRequests, 0, null, null, false);
+                this.EvaluateIfStatusCodeTriggersThrottling(ResponseStatusCodes.ResponseCodeTooManyRequests, 0, null, null, false);
             }
 
             [TestMethod]
             public void AssertTooManyRequestsStopsSendingWithFlushAsyncTask()
             {
-                this.PositiveTest(ResponseStatusCodes.ResponseCodeTooManyRequests, 0, 0, null, true);
+                this.EvaluateIfStatusCodeTriggersThrottling(ResponseStatusCodes.ResponseCodeTooManyRequests, 0, 0, null, true);
             }
 
             [TestMethod]
             public void AssertTooManyRequestsOverExtendedTimeStopsSendingAndCleansCache()
             {
-                this.PositiveTest(ResponseStatusCodes.ResponseCodeTooManyRequestsOverExtendedTime, 0, 0, 0, false);
+                this.EvaluateIfStatusCodeTriggersThrottling(ResponseStatusCodes.ResponseCodeTooManyRequestsOverExtendedTime, 0, 0, 0, false);
             }
 
             [TestMethod]
             public void AssertPaymentRequiredDoesntChangeCapacity()
             {
-                var transmitter = new StubTransmitter();
-                transmitter.OnApplyPolicies = () =>
-                {
-                    throw new Exception("Apply shouldn't be called because unsupported response code was passed");
-                };
-
-                var policy = new ThrottlingTransmissionPolicy();
-                policy.Initialize(transmitter);
-
-                transmitter.OnTransmissionSent(
-                    new TransmissionProcessedEventArgs(
-                        new StubTransmission(), null, new HttpWebResponseWrapper()
-                        {
-                            StatusCode = ResponseCodePaymentRequired
-                        }));
-
-                Assert.IsNull(policy.MaxSenderCapacity);
-                Assert.IsNull(policy.MaxBufferCapacity);
-                Assert.IsNull(policy.MaxStorageCapacity);
+                this.EvaluateIfStatusCodeIgnored(ResponseCodePaymentRequired);
             }
 
             [TestMethod]
             public void AssertUnsupportedResponseCodeDoesnotChangeCapacity()
             {
-                var transmitter = new StubTransmitter();
-                transmitter.OnApplyPolicies = () =>
-                {
-                    throw new Exception("Apply shouldn't be called because unsupported response code was passed");
-                };
-
-                var policy = new ThrottlingTransmissionPolicy();
-                policy.Initialize(transmitter);
-
-                transmitter.OnTransmissionSent(
-                    new TransmissionProcessedEventArgs(
-                        new StubTransmission(), null, new HttpWebResponseWrapper()
-                        {
-                            StatusCode = ResponseCodeUnsupported
-                        }));
-
-                Assert.IsNull(policy.MaxSenderCapacity);
-                Assert.IsNull(policy.MaxBufferCapacity);
-                Assert.IsNull(policy.MaxStorageCapacity);
+                this.EvaluateIfStatusCodeIgnored(ResponseCodeUnsupported);
             }
 
             [TestMethod]
@@ -112,49 +76,86 @@ public void CannotParseRetryAfterWritesToEventSource()
                 }
             }
 
-            private void PositiveTest(int responseCode, int? expectedSenderCapacity, int? expectedBufferCapacity, int? expectedStorageCapacity, bool hasFlushTask)
+            private void EvaluateIfStatusCodeTriggersThrottling(int responseCode, int? expectedSenderCapacity, int? expectedBufferCapacity, int? expectedStorageCapacity, bool hasFlushTask)
             {
                 const int RetryAfterSeconds = 2;
-                string retryAfter = DateTime.Now.ToUniversalTime().AddSeconds(RetryAfterSeconds).ToString("R", CultureInfo.InvariantCulture);
-                const int WaitForTheFirstApplyAsync = 100;
-                int waitForTheSecondApplyAsync = (RetryAfterSeconds * 1000) /*to milliseconds*/ +
-                    500 /**magic number to wait for other code before/after 
-                         * timer which calls 2nd ApplyAsync
-                         **/;
-
-                var policyApplied = new AutoResetEvent(false);
-                var transmitter = new StubTransmitter
-                {
-                    OnApplyPolicies = () => policyApplied.Set(),
-                };
+                var waitForTheFirstApplyAsync = TimeSpan.FromMilliseconds(100);
+                var waitForTheSecondApplyAsync = TimeSpan.FromMilliseconds(RetryAfterSeconds * 1000 + 500);
+
+                // SETUP
+                var transmitter = new StubTransmitterEvalOnApply();
 
                 var policy = new ThrottlingTransmissionPolicy();
                 policy.Initialize(transmitter);
 
-                transmitter.OnTransmissionSent(
-                    new TransmissionProcessedEventArgs(
-                        transmission: new StubTransmission() { IsFlushAsyncInProgress = hasFlushTask }, 
-                        exception: null, 
-                        response: new HttpWebResponseWrapper()
-                        {
-                            StatusCode = responseCode,
-                            StatusDescription = null,
-                            RetryAfterHeader = retryAfter
-                        }));
-
-                Assert.IsTrue(policyApplied.WaitOne(WaitForTheFirstApplyAsync));
+                transmitter.InvokeTransmissionSentEvent(responseCode, TimeSpan.FromSeconds(RetryAfterSeconds), hasFlushTask);
+
+                // ASSERT: First Handle will trigger Throttle and delay.
+                Assert.IsTrue(transmitter.IsApplyInvoked(waitForTheFirstApplyAsync));
 
                 Assert.AreEqual(expectedSenderCapacity, policy.MaxSenderCapacity);
                 Assert.AreEqual(expectedBufferCapacity, policy.MaxBufferCapacity);
                 Assert.AreEqual(expectedStorageCapacity, policy.MaxStorageCapacity);
 
-                Assert.IsTrue(policyApplied.WaitOne(waitForTheSecondApplyAsync));
+                // ASSERT: Throttle expires and policy will be reset.
+                Assert.IsTrue(transmitter.IsApplyInvoked(waitForTheSecondApplyAsync));
 
-                // Check that it resets after retry-after interval
                 Assert.IsNull(policy.MaxSenderCapacity);
                 Assert.IsNull(policy.MaxBufferCapacity);
                 Assert.IsNull(policy.MaxStorageCapacity);
             }
+
+            private void EvaluateIfStatusCodeIgnored(int statusCode)
+            {
+                var waitForTheFirstApplyAsync = TimeSpan.FromMilliseconds(100);
+
+                // SETUP
+                var transmitter = new StubTransmitterEvalOnApply();
+
+                var policy = new AuthenticationTransmissionPolicy()
+                {
+                    Enabled = true,
+                };
+                policy.Initialize(transmitter);
+
+                // ACT
+                transmitter.InvokeTransmissionSentEvent(statusCode, default, false);
+
+                // ASSERT: The Apply event handler should not be called.
+                Assert.IsFalse(transmitter.IsApplyInvoked(waitForTheFirstApplyAsync));
+
+                // ASSERT: Capacities should have default values.
+                Assert.IsNull(policy.MaxSenderCapacity);
+                Assert.IsNull(policy.MaxBufferCapacity);
+                Assert.IsNull(policy.MaxStorageCapacity);
+            }
+
+            private class StubTransmitterEvalOnApply : StubTransmitter
+            {
+                private AutoResetEvent autoResetEvent;
+
+                public StubTransmitterEvalOnApply()
+                {
+                    this.autoResetEvent = new AutoResetEvent(false);
+                    this.OnApplyPolicies = () => this.autoResetEvent.Set();
+                }
+
+                public void InvokeTransmissionSentEvent(int responseStatusCode, TimeSpan retryAfter, bool isFlushAsyncInProgress)
+                {
+                    this.OnTransmissionSent(new TransmissionProcessedEventArgs(
+                        transmission: new StubTransmission() { IsFlushAsyncInProgress = isFlushAsyncInProgress },
+                        exception: null,
+                        response: new HttpWebResponseWrapper()
+                        {
+                            StatusCode = responseStatusCode,
+                            StatusDescription = null,
+                            RetryAfterHeader = DateTime.Now.ToUniversalTime().Add(retryAfter).ToString("R", CultureInfo.InvariantCulture),
+                        }
+                    ));
+                }
+
+                public bool IsApplyInvoked(TimeSpan timeout) => this.autoResetEvent.WaitOne(timeout);
+            }
         }
     }
 }

BASE/src/ServerTelemetryChannel/Implementation/ResponseStatusCodes.cs

@@ -11,8 +11,9 @@ internal class ResponseStatusCodes
         public const int InternalServerError = 500;
         public const int BadGateway = 502;
         public const int ServiceUnavailable = 503;
-        public const int GatewayTimeout = 504; 
-        public const int Unauthorized = 401; // AAD
-        public const int Forbidden = 403; // AAD
+        public const int GatewayTimeout = 504;
+        public const int BadRequest = 400; // AAD: AI resource was configured for AAD, but SDK is using older api. (v2 and v2.1).
+        public const int Unauthorized = 401; // AAD: token is either absent, invalid, or expired.
+        public const int Forbidden = 403; // AAD: Provided credentials do not grant access to ingest telemetry.
     }
 }