Add CmdletContext

src/Authentication.Abstractions/AuthenticationTelemetry.cs

@@ -1,48 +1,19 @@
-// ----------------------------------------------------------------------------------
-//
-// Copyright Microsoft Corporation
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-// http://www.apache.org/licenses/LICENSE-2.0
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-// ----------------------------------------------------------------------------------
-
-using System.Collections.Generic;
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Interfaces;
 
 namespace Microsoft.Azure.Commands.Common.Authentication.Abstractions
 {
-    /// <summary>
-    /// A model class for a list of authenction telemetry records.
-    /// </summary>
-    public class AuthenticationTelemetry
+    public class AuthenticationTelemetry: IAzureTelemetry <AuthTelemetryRecord> 
     {
-        /// <summary>
-        /// The first record of authentication telemetry data, usually describes the main method of this authentication process.
-        /// </summary>
-        public IAuthTelemetryRecord Head { get; } = null;
+        //public delegate string RequestIdAccquirer();
 
-        /// <summary>
-        /// The left part of authentication telemetry records.
-        /// </summary>
-        public IList<IAuthTelemetryRecord> Tail { get; } = new List<IAuthTelemetryRecord>();
+        //public static RequestIdAccquirer GetRequestId;
 
-        public AuthenticationTelemetry(IEnumerable<IAuthTelemetryRecord> records)
-        {
-            var enumerator = records.GetEnumerator();
-            if (enumerator.MoveNext())
-            {
-                Head = enumerator.Current;
-            }
+        //private static string requestId;
 
-            while (enumerator.MoveNext())
-            {
-                Tail.Add(enumerator.Current);
-            }
+        public AuthenticationTelemetryData GetTelemetryRecord(ICmdletContext cmdletContext)
+        {
+            var records = PopTelemetryRecord(cmdletContext);
+            return records == null ? null : new AuthenticationTelemetryData(records);
         }
     }
 }

src/Authentication.Abstractions/AuthenticationTelemetryData.cs

@@ -0,0 +1,48 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using System.Collections.Generic;
+
+namespace Microsoft.Azure.Commands.Common.Authentication.Abstractions
+{
+    /// <summary>
+    /// A model class for a list of authenction telemetry records.
+    /// </summary>
+    public class AuthenticationTelemetryData
+    {
+        /// <summary>
+        /// The first record of authentication telemetry data, usually describes the main method of this authentication process.
+        /// </summary>
+        public IAuthTelemetryRecord Head { get; } = null;
+
+        /// <summary>
+        /// The left part of authentication telemetry records.
+        /// </summary>
+        public IList<IAuthTelemetryRecord> Tail { get; } = new List<IAuthTelemetryRecord>();
+
+        public AuthenticationTelemetryData(IEnumerable<IAuthTelemetryRecord> records)
+        {
+            var enumerator = records.GetEnumerator();
+            if (enumerator.MoveNext())
+            {
+                Head = enumerator.Current;
+            }
+
+            while (enumerator.MoveNext())
+            {
+                Tail.Add(enumerator.Current);
+            }
+        }
+    }
+}

src/Authentication.Abstractions/AzureCmdletContext.cs

@@ -0,0 +1,30 @@
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Interfaces;
+
+using System;
+using System.Collections.Generic;
+using System.Text;
+
+namespace Microsoft.Azure.Commands.Common.Authentication.Abstractions
+{
+    public class AzureCmdletContext : ICmdletContext
+    {
+        private string cmdletId;
+
+        public const ICmdletContext CmdletNone = null;
+
+        public AzureCmdletContext(string id)
+        {
+            cmdletId = id;
+        }
+        public string CmdletId
+        { 
+            get => cmdletId;
+            set => cmdletId = value;
+        }
+
+        public bool IsValid
+        { 
+            get => string.IsNullOrEmpty(cmdletId);
+        }
+    }
+}

src/Authentication.Abstractions/Interfaces/IAuthenticationFactory.cs

@@ -15,6 +15,7 @@
 using Microsoft.Rest;
 using System;
 using System.Security;
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Interfaces;
 
 namespace Microsoft.Azure.Commands.Common.Authentication.Abstractions
 {
@@ -43,6 +44,7 @@ IAccessToken Authenticate(
             string promptBehavior,
             Action<string> promptAction,
             IAzureTokenCache tokenCache,
+            ICmdletContext cmdletContext,
             string resourceId = AzureEnvironment.Endpoint.ActiveDirectoryServiceEndpointResourceId);
 
         /// <summary>
@@ -63,22 +65,25 @@ IAccessToken Authenticate(
             SecureString password,
             string promptBehavior,
             Action<string> promptAction,
+            ICmdletContext cmdletContext,
             string resourceId = AzureEnvironment.Endpoint.ActiveDirectoryServiceEndpointResourceId);
 
         /// <summary>
         /// Get AutoRest credentials for the given context
         /// </summary>
         /// <param name="context">The target azure context</param>
+        /// <param name="cmdletContext">The caller cmdlet context</param>
         /// <returns>AutoRest client credentials targeting the given context</returns>
-        ServiceClientCredentials GetServiceClientCredentials(IAzureContext context);
+        ServiceClientCredentials GetServiceClientCredentials(IAzureContext context, ICmdletContext cmdletContext);
 
         /// <summary>
         /// Get AutoRest credebntials using the given context and named endpoint
         /// </summary>
         /// <param name="context">The context to use for authentication</param>
+        /// <param name="cmdletContext">The caller cmdlet context</param>
         /// <param name="targetEndpoint">The named endpoint the AutoRest client will target</param>
         /// <returns>AutoRest client crentials targeting the given context and endpoint</returns>
-        ServiceClientCredentials GetServiceClientCredentials(IAzureContext context, string targetEndpoint);
+        ServiceClientCredentials GetServiceClientCredentials(IAzureContext context, string targetEndpoint, ICmdletContext cmdletContext);
 
         /// <summary>
         /// Get service client credentials with initial token and delegate for renewing
@@ -94,10 +99,5 @@ IAccessToken Authenticate(
         /// <param name="account">The account to remove credentials for</param>
         /// <param name="tokenCache">The TokenCache to remove credentials from</param>
         void RemoveUser(IAzureAccount account, IAzureTokenCache tokenCache);
-
-        /// <summary>
-        /// Get the information to be recorded in Telemetry
-        /// </summary>
-        AuthenticationTelemetry GetDataForTelemetry(string requestId);
     }
 }

src/Authentication.Abstractions/Interfaces/IAzureTelemetry.cs

@@ -0,0 +1,34 @@
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+
+namespace Microsoft.Azure.Commands.Common.Authentication.Abstractions.Interfaces
+{
+    public abstract class IAzureTelemetry <T>
+    {
+        private ConcurrentDictionary<string, IList<T>> telemetryDataAccquirer = new ConcurrentDictionary<string, IList<T>>();
+
+        public bool PushTelemetryRecord(ICmdletContext cmdletContext, T record)
+        {
+            if (cmdletContext != null && cmdletContext.IsValid && record != null)
+            {
+                if (!telemetryDataAccquirer.ContainsKey(cmdletContext.CmdletId))
+                {
+                    telemetryDataAccquirer[cmdletContext.CmdletId] = new List<T>();
+                }
+                telemetryDataAccquirer[cmdletContext.CmdletId].Add(record);
+                return true;
+            }
+            return false;
+        }
+
+        public IList<T> PopTelemetryRecord(ICmdletContext cmdletContext)
+        {
+            if (cmdletContext != null && cmdletContext.IsValid && telemetryDataAccquirer.ContainsKey(cmdletContext.CmdletId))
+            {
+                telemetryDataAccquirer.TryRemove(cmdletContext.CmdletId, out IList<T> records);
+                return records;
+            }
+            return null;
+        }
+    }
+}

src/Authentication.Abstractions/Interfaces/IClientFactory.cs

@@ -13,6 +13,7 @@
 // ----------------------------------------------------------------------------------
 
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Interfaces;
 using Microsoft.Azure.Commands.Common.Authentication.Models;
 using Microsoft.Rest;
 using System;
@@ -34,9 +35,10 @@ public interface IClientFactory: IHyakClientFactory
         /// <typeparam name="TClient">The client type to create</typeparam>
         /// <param name="context">The azure context to target</param>
         /// <param name="endpoint">The named endpoint the client shoulld target</param>
+        /// <param name="cmdletContext">Cmdlet Context</param>
         /// <returns>A client properly authenticated in the given context, properly configured for use with Azure PowerShell, 
         /// targeting the given named endpoint in the targeted environment</returns>
-        TClient CreateArmClient<TClient>(IAzureContext context, string endpoint) where TClient : ServiceClient<TClient>;
+        TClient CreateArmClient<TClient>(IAzureContext context, string endpoint, ICmdletContext cmdletContext) where TClient : ServiceClient<TClient>;
 
         /// <summary>
         /// Create a properly configured AutoRest client using custom client parameters

src/Authentication.Abstractions/Interfaces/ICmdletContext.cs

@@ -0,0 +1,22 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+namespace Microsoft.Azure.Commands.Common.Authentication.Abstractions.Interfaces
+{
+    public interface ICmdletContext
+    {
+        string CmdletId { get; }
+
+        bool IsValid { get; }
+    }
+}

src/Common/AzurePSCmdlet.cs

@@ -14,7 +14,7 @@
 
 using Microsoft.Azure.Commands.Common.Authentication;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
-using Microsoft.Azure.Commands.Common.Authentication.Utilities;
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Interfaces;
 using Microsoft.Azure.PowerShell.Common.Config;
 using Microsoft.Azure.PowerShell.Common.Share.Survey;
 using Microsoft.Azure.PowerShell.Common.UpgradeNotification;
@@ -94,6 +94,7 @@ protected AzurePSDataCollectionProfile _dataCollectionProfile
         protected static string _sessionId = Guid.NewGuid().ToString();
         protected const string _fileTimeStampSuffixFormat = "yyyy-MM-dd-THH-mm-ss-fff";
         protected string _clientRequestId = Guid.NewGuid().ToString();
+        protected ICmdletContext _cmdletContext;
         protected static DateTimeOffset? _previousEndTime = null;
         protected MetricHelper _metricHelper;
         protected AzurePSQoSEvent _qosEvent;
@@ -377,16 +378,12 @@ protected override void BeginProcessing()
             InitDebuggingFilter();
             SetupDebuggingTraces();
             SetupHttpClientPipeline();
+            _cmdletContext = new AzureCmdletContext(_clientRequestId);
             base.BeginProcessing();
 
             //Now see if the cmdlet has any Breaking change attributes on it and process them if it does
             //This will print any breaking change attribute messages that are applied to the cmdlet
             WriteBreakingChangeOrPreviewMessage();
-
-            if (AzureSession.Instance.TryGetComponent(nameof(ThreadCmdldetMap), out ThreadCmdldetMap threadCmdletMap))
-            {
-                threadCmdletMap.PushCmdletId(this._clientRequestId);
-            }
         }
 
         private void WriteBreakingChangeOrPreviewMessage()
@@ -849,7 +846,12 @@ protected void LogQosEvent()
 
             _qosEvent.ParameterSetName = this.ParameterSetName;
             _qosEvent.FinishQosEvent();
-            _qosEvent.AuthTelemetry = AzureSession.Instance.AuthenticationFactory.GetDataForTelemetry(_qosEvent.ClientRequestId);
+
+            if (!AzureSession.Instance.TryGetComponent(nameof(AuthenticationTelemetry), out AuthenticationTelemetry authenticationTelemetry))
+            {
+                throw new NullReferenceException("AuthenticationTelemetry not registered");
+            }
+            _qosEvent.AuthTelemetry = authenticationTelemetry.GetTelemetryRecord(_cmdletContext);
 
             if (!IsUsageMetricEnabled && (!IsErrorMetricEnabled || _qosEvent.IsSuccess))
             {

src/Common/MetricHelper.cs

@@ -292,7 +292,7 @@ public void SetPSHost(PSHost host)
         {
         }
 
-        private static void PopulateAuthenticationPropertiesFromQos(AuthenticationTelemetry telemetry, IDictionary<string, string> eventProperties)
+        private static void PopulateAuthenticationPropertiesFromQos(AuthenticationTelemetryData telemetry, IDictionary<string, string> eventProperties)
         {
             var record = telemetry.Head;
             eventProperties[$"{AuthTelemetryRecord.AuthTelemetryPropertyHeadPrefix}-{nameof(record.TokenCredentialName).ToLower()}"] = record.TokenCredentialName;
@@ -693,7 +693,7 @@ public class AzurePSQoSEvent
 
     public SanitizerTelemetry SanitizerInfo { get; set; }
 
-    public AuthenticationTelemetry AuthTelemetry { get; set; }
+    public AuthenticationTelemetryData AuthTelemetry { get; set; }
 
     public AzurePSQoSEvent()
     {