AzureSession.TryGetComponent throws KeyNotFoundException when AzureSession is used concurrently

[AnatoliB]

Here is the current implementation: https://github.com/Azure/azure-powershell-common/blob/ae44bad385cacd80cdf252d0be6a519b97486dd8/src/Authentication.Abstractions/AzureSession.cs#L214

Even though a ConcurrentDictionary is used for _componentRegistry, line 220 can still throw KeyNotFoundException if the collection is changed after invoking ContainsKey (line 218) but before invoking [] (line 220), which I believe is not intentional.

I recommend invoking _componentRegistry.TryGetValue instead.

It may be hard to reproduce this because very precise timing is required, but it actually happens when used at scale (for example, in the context of Azure Functions). A typical exception surfaced to PowerShell code looks like this:

Exception             : 
    Type       : System.Collections.Generic.KeyNotFoundException
    TargetSite : 
        Name          : ThrowKeyNotFoundException
        DeclaringType : [System.Collections.Concurrent.ConcurrentDictionary`2[TKey,TValue]]
        MemberType    : Method
        Module        : System.Collections.Concurrent.dll
    Message    : The given key 'WriteInformation-System.EventHandler`1[[Microsoft.Azure.Commands.Common.Authentication.Abstractions.StreamEventArgs, Microsoft.Azure.PowerShell.Authentication.Abstractions, Version=1.0.0.0, Culture=neutral, PublicKeyToken=31bf3856ad364e35]]' was not present in the dictionary.
    Source     : System.Collections.Concurrent
    HResult    : -2146232969
    StackTrace : 
   at System.Collections.Concurrent.ConcurrentDictionary`2.ThrowKeyNotFoundException(TKey key)
   at Microsoft.Azure.Commands.Common.Authentication.AzureSession.TryGetComponent[T](String componentName, T& component)
   at Microsoft.Azure.Commands.Profile.ConnectAzureRmAccountCommand.BeginProcessing()
   at System.Management.Automation.Cmdlet.DoBeginProcessing()
   at System.Management.Automation.CommandProcessorBase.DoBegin()
CategoryInfo          : NotSpecified: (:) [Connect-AzAccount], KeyNotFoundException
FullyQualifiedErrorId : System.Collections.Generic.KeyNotFoundException,Microsoft.Azure.Commands.Profile.ConnectAzureRmAccountCommand
InvocationInfo        : 
    MyCommand        : Connect-AzAccount
    ScriptLineNumber : 7
    OffsetInLine     : 13
    HistoryId        : 1
    ScriptName       : C:\home\site\wwwroot\MyFunction\run.ps1
    Line             : $context = (Connect-AzAccount -Identity -AccountId ... -Subscription ...

    Statement        : Connect-AzAccount -Identity -AccountId ... -Subscription ...
    PositionMessage  : At C:\home\site\wwwroot\MyFunction\run.ps1:7 char:13
                       + $context = (Connect-AzAccount -Identity -AccountId ... -Su …
                       +             ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    PSScriptRoot     : C:\home\site\wwwroot\MyFunction
    PSCommandPath    : C:\home\site\wwwroot\MyFunction\run.ps1
    InvocationName   : Connect-AzAccount
    CommandOrigin    : Internal

[isra-fel]

Transferred from common repo -

Thanks a lot for @AnatoliB 's investigation! The RCA makes a lot of sense since considering that the writeWaring key is frequently read and written. I'll try if we can snug this bug fix in the major release.