From 27d5fb8dc14c0d8672190eedd04a6eb029538425 Mon Sep 17 00:00:00 2001
From: Luke Bakken <luke@bakken.io>
Date: Wed, 29 Nov 2023 12:21:31 -0800
Subject: [PATCH] Fixes #1429 for `6.x`

See #1431

* Modify test program to demonstrate the issue
* Fix underlying issue by allowing multiple Register calls for the same instance of ICredentialsProvider
---
 .../client/api/ICredentialsRefresher.cs       | 38 +++++++++++--------
 projects/TestApplications/OAuth2/Program.cs   | 31 ++++++++-------
 .../Unit/APIApproval.Approve.verified.txt     |  2 +
 3 files changed, 40 insertions(+), 31 deletions(-)

diff --git a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs
index bbf9666788..a609aff052 100644
--- a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs
+++ b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs
@@ -30,7 +30,7 @@
 //---------------------------------------------------------------------------
 
 using System;
-using System.Collections.Generic;
+using System.Collections.Concurrent;
 using System.Diagnostics.Tracing;
 using System.Timers;
 
@@ -59,11 +59,13 @@ public class TimerBasedCredentialRefresherEventSource : EventSource
         public void TriggeredTimer(string name) => WriteEvent(4, "TriggeredTimer", name);
         [Event(5)]
         public void RefreshedCredentials(string name, bool succesfully) => WriteEvent(5, "RefreshedCredentials", name, succesfully);
+        [Event(6)]
+        public void AlreadyRegistered(string name) => WriteEvent(6, "AlreadyRegistered", name);
     }
 
     public class TimerBasedCredentialRefresher : ICredentialsRefresher
     {
-        private Dictionary<ICredentialsProvider, Timer> _registrations = new Dictionary<ICredentialsProvider, Timer>();
+        private readonly ConcurrentDictionary<ICredentialsProvider, Timer> _registrations = new ConcurrentDictionary<ICredentialsProvider, Timer>();
 
         public ICredentialsProvider Register(ICredentialsProvider provider, NotifyCredentialRefreshed callback)
         {
@@ -72,25 +74,31 @@ public ICredentialsProvider Register(ICredentialsProvider provider, NotifyCreden
                 return provider;
             }
 
-            _registrations.Add(provider, scheduleTimer(provider, callback));
-            TimerBasedCredentialRefresherEventSource.Log.Registered(provider.Name);
+            if (_registrations.TryAdd(provider, scheduleTimer(provider, callback)))
+            {
+                TimerBasedCredentialRefresherEventSource.Log.Registered(provider.Name);
+            }
+            else
+            {
+                TimerBasedCredentialRefresherEventSource.Log.AlreadyRegistered(provider.Name);
+            }
+
             return provider;
         }
 
         public bool Unregister(ICredentialsProvider provider)
         {
-            if (!_registrations.ContainsKey(provider))
+            if (_registrations.TryRemove(provider, out Timer timer))
             {
-                return false;
-            }
-
-            var timer = _registrations[provider];
-            if (timer != null)
-            {
-                TimerBasedCredentialRefresherEventSource.Log.Unregistered(provider.Name);
-                timer.Stop();
-                _registrations.Remove(provider);
-                timer.Dispose();
+                try
+                {
+                    TimerBasedCredentialRefresherEventSource.Log.Unregistered(provider.Name);
+                    timer.Stop();
+                }
+                finally
+                {
+                    timer.Dispose();
+                }
                 return true;
             }
             else
diff --git a/projects/TestApplications/OAuth2/Program.cs b/projects/TestApplications/OAuth2/Program.cs
index ae76b24023..fbf9f0254a 100644
--- a/projects/TestApplications/OAuth2/Program.cs
+++ b/projects/TestApplications/OAuth2/Program.cs
@@ -44,26 +44,25 @@ public static void Main(string[] args)
                 CredentialsRefresher = GetCredentialsRefresher()
             };
 
-            using (IConnection connection = connectionFactory.CreateConnection())
+            using IConnection publishConnection = connectionFactory.CreateConnection();
+            using IConnection consumingConnection = connectionFactory.CreateConnection();
+
+            using IModel publisher = declarePublisher(publishConnection);
+            using IModel subscriber = declareConsumer(consumingConnection);
+
+            Publish(publisher);
+            Consume(subscriber);
+
+            if (oauth2Options.TokenExpiresInSeconds > 0)
             {
-                using (IModel publisher = declarePublisher(connection))
-                using (IModel subscriber = declareConsumer(connection))
+                for (int i = 0; i < 4; i++)
                 {
+                    Console.WriteLine("Wait until Token expires. Attempt #" + (i + 1));
+                    Thread.Sleep((oauth2Options.TokenExpiresInSeconds + 10) * 1000);
+                    Console.WriteLine("Resuming ..");
                     Publish(publisher);
+                    _doneEvent.Reset();
                     Consume(subscriber);
-
-                    if (oauth2Options.TokenExpiresInSeconds > 0)
-                    {
-                        for (int i = 0; i < 4; i++)
-                        {
-                            Console.WriteLine("Wait until Token expires. Attempt #" + (i + 1));
-                            Thread.Sleep((oauth2Options.TokenExpiresInSeconds + 10) * 1000);
-                            Console.WriteLine("Resuming ..");
-                            Publish(publisher);
-                            _doneEvent.Reset();
-                            Consume(subscriber);
-                        }
-                    }
                 }
             }
         }
diff --git a/projects/Unit/APIApproval.Approve.verified.txt b/projects/Unit/APIApproval.Approve.verified.txt
index d7bcdaab0c..044048d3be 100644
--- a/projects/Unit/APIApproval.Approve.verified.txt
+++ b/projects/Unit/APIApproval.Approve.verified.txt
@@ -671,6 +671,8 @@ namespace RabbitMQ.Client
     {
         public TimerBasedCredentialRefresherEventSource() { }
         public static RabbitMQ.Client.TimerBasedCredentialRefresherEventSource Log { get; }
+        [System.Diagnostics.Tracing.Event(6)]
+        public void AlreadyRegistered(string name) { }
         [System.Diagnostics.Tracing.Event(5)]
         public void RefreshedCredentials(string name, bool succesfully) { }
         [System.Diagnostics.Tracing.Event(1)]