From e688eadd4abba7de13ca7fd54684c8f8f0de2e7e Mon Sep 17 00:00:00 2001 From: Matthias Werning Date: Thu, 25 Jul 2024 17:28:26 +0200 Subject: [PATCH 1/9] Fix Timer leakage in TimerBasedCredentialRefresher by maintaining only one registration across all connection recoveries --- .../client/api/ICredentialsRefresher.cs | 138 +++++++++++++----- 1 file changed, 100 insertions(+), 38 deletions(-) diff --git a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs index ea16d26b1f..ea3aa8cd18 100644 --- a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs +++ b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs @@ -31,9 +31,11 @@ using System; using System.Collections.Concurrent; -using System.Diagnostics.CodeAnalysis; +using System.Collections.Generic; using System.Diagnostics.Tracing; using System.Threading.Tasks; +using System.Timers; + namespace RabbitMQ.Client { public interface ICredentialsRefresher @@ -71,8 +73,10 @@ public class TimerBasedCredentialRefresherEventSource : EventSource public class TimerBasedCredentialRefresher : ICredentialsRefresher { - private readonly ConcurrentDictionary _registrations = new ConcurrentDictionary(); - + private readonly IDictionary _registrations = + new Dictionary(); + private readonly object _lockObj = new(); + public ICredentialsProvider Register(ICredentialsProvider provider, ICredentialsRefresher.NotifyCredentialRefreshedAsync callback) { if (!provider.ValidUntil.HasValue || provider.ValidUntil.Value.Equals(TimeSpan.Zero)) @@ -80,65 +84,123 @@ public ICredentialsProvider Register(ICredentialsProvider provider, ICredentials return provider; } - if (_registrations.TryAdd(provider, scheduleTimer(provider, callback))) + lock (_lockObj) { + if (_registrations.TryGetValue(provider, out var registration)) + { + registration.Callback = callback; + TimerBasedCredentialRefresherEventSource.Log.AlreadyRegistered(provider.Name); + return provider; + } + + registration = new TimerRegistration(_lockObj, callback); + _registrations.Add(provider, registration); + registration.ScheduleTimer(provider); + TimerBasedCredentialRefresherEventSource.Log.Registered(provider.Name); } - else - { - TimerBasedCredentialRefresherEventSource.Log.AlreadyRegistered(provider.Name); - } return provider; } public bool Unregister(ICredentialsProvider provider) { - if (_registrations.TryRemove(provider, out System.Timers.Timer? timer)) + lock (_lockObj) { - try + if (_registrations.TryGetValue(provider, out var registration)) { + _registrations.Remove(provider); + TimerBasedCredentialRefresherEventSource.Log.Unregistered(provider.Name); - timer.Stop(); - } - finally - { - timer.Dispose(); - } - return true; - } - else - { - return false; + registration.Dispose(); + return true; + } } + + return false; } - private System.Timers.Timer scheduleTimer(ICredentialsProvider provider, ICredentialsRefresher.NotifyCredentialRefreshedAsync callback) + private class TimerRegistration : IDisposable { - System.Timers.Timer timer = new System.Timers.Timer(); - timer.Interval = provider.ValidUntil!.Value.TotalMilliseconds * (1.0 - (1 / 3.0)); - timer.Elapsed += (o, e) => + + private readonly object _lockObj; + private System.Timers.Timer? _timer; + private bool _disposed; + + public ICredentialsRefresher.NotifyCredentialRefreshedAsync Callback { get; set; } + + public TimerRegistration(object lockObj, ICredentialsRefresher.NotifyCredentialRefreshedAsync callback) + { + _lockObj = lockObj; + Callback = callback; + } + + public void ScheduleTimer(ICredentialsProvider provider) + { + if (provider.ValidUntil == null) + { + throw new ArgumentNullException("ValidUntil of " + nameof(provider) + " was null"); + } + if (_disposed) + { + throw new InvalidOperationException("Registration already disposed"); + } + + var newTimer = new Timer(); + newTimer.Interval = provider.ValidUntil.Value.TotalMilliseconds * (1.0 - 1 / 3.0); + newTimer.Elapsed += (o, e) => + { + TimerBasedCredentialRefresherEventSource.Log.TriggeredTimer(provider.Name); + + lock (_lockObj) + { + try + { + if (_disposed) + { + // We were waiting and the registration has been disposed in meanwhile + return; + } + + provider.Refresh(); + ScheduleTimer(provider); + Callback.Invoke(provider.Password != null); + TimerBasedCredentialRefresherEventSource.Log.RefreshedCredentials(provider.Name, true); + } + catch (Exception) + { + Callback.Invoke(false); + TimerBasedCredentialRefresherEventSource.Log.RefreshedCredentials(provider.Name, false); + } + } + }; + newTimer.Enabled = true; + newTimer.AutoReset = false; + TimerBasedCredentialRefresherEventSource.Log.ScheduledTimer(provider.Name, newTimer.Interval); + _timer = newTimer; + } + + public void Dispose() { - TimerBasedCredentialRefresherEventSource.Log.TriggeredTimer(provider.Name); + if (_disposed) + { + throw new InvalidOperationException("registration already disposed"); + } + try { - provider.Refresh(); - scheduleTimer(provider, callback); - callback.Invoke(provider.Password != null); - TimerBasedCredentialRefresherEventSource.Log.RefreshedCredentials(provider.Name, true); + _timer?.Stop(); + _disposed = true; } - catch (Exception) + finally { - callback.Invoke(false); - TimerBasedCredentialRefresherEventSource.Log.RefreshedCredentials(provider.Name, false); + _timer?.Dispose(); + _timer = null; } + } - }; - timer.Enabled = true; - timer.AutoReset = false; - TimerBasedCredentialRefresherEventSource.Log.ScheduledTimer(provider.Name, timer.Interval); - return timer; } + } class NoOpCredentialsRefresher : ICredentialsRefresher From e7351cc37f6ae640aef35d789e000dc1926fd234 Mon Sep 17 00:00:00 2001 From: Matthias Werning Date: Thu, 25 Jul 2024 17:50:09 +0200 Subject: [PATCH 2/9] Re-add accidentally removed import of System.Diagnostics.CodeAnalysis --- projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs index ea3aa8cd18..0a1fc51507 100644 --- a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs +++ b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs @@ -30,8 +30,8 @@ //--------------------------------------------------------------------------- using System; -using System.Collections.Concurrent; using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; using System.Diagnostics.Tracing; using System.Threading.Tasks; using System.Timers; From 51c2acb34a5a43c9ad2f3e7c876a7d7c8413d100 Mon Sep 17 00:00:00 2001 From: Matthias Werning Date: Thu, 25 Jul 2024 17:58:52 +0200 Subject: [PATCH 3/9] Remove unnecessary whitespaces --- .../client/api/ICredentialsRefresher.cs | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs index 0a1fc51507..a229b8c038 100644 --- a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs +++ b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs @@ -76,7 +76,7 @@ public class TimerBasedCredentialRefresher : ICredentialsRefresher private readonly IDictionary _registrations = new Dictionary(); private readonly object _lockObj = new(); - + public ICredentialsProvider Register(ICredentialsProvider provider, ICredentialsRefresher.NotifyCredentialRefreshedAsync callback) { if (!provider.ValidUntil.HasValue || provider.ValidUntil.Value.Equals(TimeSpan.Zero)) @@ -96,7 +96,7 @@ public ICredentialsProvider Register(ICredentialsProvider provider, ICredentials registration = new TimerRegistration(_lockObj, callback); _registrations.Add(provider, registration); registration.ScheduleTimer(provider); - + TimerBasedCredentialRefresherEventSource.Log.Registered(provider.Name); } @@ -110,13 +110,13 @@ public bool Unregister(ICredentialsProvider provider) if (_registrations.TryGetValue(provider, out var registration)) { _registrations.Remove(provider); - + TimerBasedCredentialRefresherEventSource.Log.Unregistered(provider.Name); registration.Dispose(); return true; - } + } } - + return false; } @@ -126,7 +126,7 @@ private class TimerRegistration : IDisposable private readonly object _lockObj; private System.Timers.Timer? _timer; private bool _disposed; - + public ICredentialsRefresher.NotifyCredentialRefreshedAsync Callback { get; set; } public TimerRegistration(object lockObj, ICredentialsRefresher.NotifyCredentialRefreshedAsync callback) @@ -145,7 +145,7 @@ public void ScheduleTimer(ICredentialsProvider provider) { throw new InvalidOperationException("Registration already disposed"); } - + var newTimer = new Timer(); newTimer.Interval = provider.ValidUntil.Value.TotalMilliseconds * (1.0 - 1 / 3.0); newTimer.Elapsed += (o, e) => @@ -161,7 +161,7 @@ public void ScheduleTimer(ICredentialsProvider provider) // We were waiting and the registration has been disposed in meanwhile return; } - + provider.Refresh(); ScheduleTimer(provider); Callback.Invoke(provider.Password != null); @@ -171,7 +171,7 @@ public void ScheduleTimer(ICredentialsProvider provider) { Callback.Invoke(false); TimerBasedCredentialRefresherEventSource.Log.RefreshedCredentials(provider.Name, false); - } + } } }; newTimer.Enabled = true; @@ -179,14 +179,14 @@ public void ScheduleTimer(ICredentialsProvider provider) TimerBasedCredentialRefresherEventSource.Log.ScheduledTimer(provider.Name, newTimer.Interval); _timer = newTimer; } - + public void Dispose() { if (_disposed) { throw new InvalidOperationException("registration already disposed"); } - + try { _timer?.Stop(); @@ -200,7 +200,7 @@ public void Dispose() } } - + } class NoOpCredentialsRefresher : ICredentialsRefresher From e0885ff3bcfdc1e947925310dbc4b653a8c5d64a Mon Sep 17 00:00:00 2001 From: Matthias Werning Date: Fri, 26 Jul 2024 11:31:48 +0200 Subject: [PATCH 4/9] Add unit test to test callback update: Make sure old callback isn't called again after updating --- .../Unit/TestTimerBasedCredentialRefresher.cs | 54 ++++++++++++++++--- 1 file changed, 48 insertions(+), 6 deletions(-) diff --git a/projects/Test/Unit/TestTimerBasedCredentialRefresher.cs b/projects/Test/Unit/TestTimerBasedCredentialRefresher.cs index faed6f03f7..a3c3f1fb15 100644 --- a/projects/Test/Unit/TestTimerBasedCredentialRefresher.cs +++ b/projects/Test/Unit/TestTimerBasedCredentialRefresher.cs @@ -43,7 +43,7 @@ public class MockCredentialsProvider : ICredentialsProvider private readonly ITestOutputHelper _testOutputHelper; private readonly TimeSpan? _validUntil = TimeSpan.FromSeconds(1); private Exception _ex = null; - private bool _refreshCalled = false; + private int _refreshCalledTimes = 0; public MockCredentialsProvider(ITestOutputHelper testOutputHelper) { @@ -56,11 +56,11 @@ public MockCredentialsProvider(ITestOutputHelper testOutputHelper, TimeSpan vali _validUntil = validUntil; } - public bool RefreshCalled + public int RefreshCalledTimes { get { - return _refreshCalled; + return _refreshCalledTimes; } } @@ -87,7 +87,7 @@ public string Password public void Refresh() { - _refreshCalled = true; + _refreshCalledTimes++; } public void PasswordThrows(Exception ex) @@ -145,7 +145,49 @@ Task cb(bool arg) _refresher.Register(credentialsProvider, cb); Assert.True(await tcs.Task); - Assert.True(credentialsProvider.RefreshCalled); + Assert.True(credentialsProvider.RefreshCalledTimes > 0); + Assert.True(_refresher.Unregister(credentialsProvider)); + } + } + } + + [Fact] + public async Task TestRefreshTokenUpdateCallback() + { + var tcs1 = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var tcs2 = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + int cb1CalledTimes = 0; + int cb2CalledTimes = 0; + + using (var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5))) + { + using (CancellationTokenRegistration ctr = cts.Token.Register(() => { tcs1.TrySetCanceled(); tcs2.TrySetCanceled(); })) { + var credentialsProvider = new MockCredentialsProvider(_testOutputHelper, TimeSpan.FromSeconds(1)); + + Task cb1(bool arg) + { + cb1CalledTimes++; + tcs1.SetResult(arg); + return Task.CompletedTask; + } + + Task cb2(bool arg) + { + cb2CalledTimes++; + tcs2.SetResult(arg); + return Task.CompletedTask; + } + + _refresher.Register(credentialsProvider, cb1); + Assert.True(await tcs1.Task); + Assert.True(credentialsProvider.RefreshCalledTimes == 1); + Assert.True(cb1CalledTimes == 1); + _refresher.Register(credentialsProvider, cb2); + Assert.True(await tcs2.Task); + Assert.True(credentialsProvider.RefreshCalledTimes == 2); + Assert.True(cb2CalledTimes == 1); + Assert.True(cb1CalledTimes == 1); + Assert.True(_refresher.Unregister(credentialsProvider)); } } @@ -172,7 +214,7 @@ Task cb(bool arg) _refresher.Register(credentialsProvider, cb); Assert.False(await tcs.Task); - Assert.True(credentialsProvider.RefreshCalled); + Assert.True(credentialsProvider.RefreshCalledTimes > 0); Assert.True(_refresher.Unregister(credentialsProvider)); } } From 8794d688e8fe5233b712f1a29e93971f29142d80 Mon Sep 17 00:00:00 2001 From: Matthias Werning Date: Fri, 26 Jul 2024 11:35:39 +0200 Subject: [PATCH 5/9] Fix whitespace in test --- projects/Test/Unit/TestTimerBasedCredentialRefresher.cs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/projects/Test/Unit/TestTimerBasedCredentialRefresher.cs b/projects/Test/Unit/TestTimerBasedCredentialRefresher.cs index a3c3f1fb15..99ecbf03ca 100644 --- a/projects/Test/Unit/TestTimerBasedCredentialRefresher.cs +++ b/projects/Test/Unit/TestTimerBasedCredentialRefresher.cs @@ -161,7 +161,8 @@ public async Task TestRefreshTokenUpdateCallback() using (var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5))) { - using (CancellationTokenRegistration ctr = cts.Token.Register(() => { tcs1.TrySetCanceled(); tcs2.TrySetCanceled(); })) { + using (CancellationTokenRegistration ctr = cts.Token.Register(() => { tcs1.TrySetCanceled(); tcs2.TrySetCanceled(); })) + { var credentialsProvider = new MockCredentialsProvider(_testOutputHelper, TimeSpan.FromSeconds(1)); Task cb1(bool arg) From 717275bfd9833395dcebf0be3f0545d26e1760b1 Mon Sep 17 00:00:00 2001 From: Matthias Werning Date: Fri, 26 Jul 2024 14:37:09 +0200 Subject: [PATCH 6/9] Move actual refresh on timer elapse into non-locked scope to handle await of notification callback --- .../client/api/ICredentialsRefresher.cs | 38 +++++++++---------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs index a229b8c038..3fa9462e62 100644 --- a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs +++ b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs @@ -139,40 +139,40 @@ public void ScheduleTimer(ICredentialsProvider provider) { if (provider.ValidUntil == null) { - throw new ArgumentNullException("ValidUntil of " + nameof(provider) + " was null"); + throw new ArgumentNullException(nameof(provider.ValidUntil) + " of " + nameof(provider) + " was null"); } if (_disposed) { - throw new InvalidOperationException("Registration already disposed"); + return; } var newTimer = new Timer(); newTimer.Interval = provider.ValidUntil.Value.TotalMilliseconds * (1.0 - 1 / 3.0); - newTimer.Elapsed += (o, e) => + newTimer.Elapsed += async (o, e) => { TimerBasedCredentialRefresherEventSource.Log.TriggeredTimer(provider.Name); lock (_lockObj) { - try + if (_disposed) { - if (_disposed) - { - // We were waiting and the registration has been disposed in meanwhile - return; - } - - provider.Refresh(); - ScheduleTimer(provider); - Callback.Invoke(provider.Password != null); - TimerBasedCredentialRefresherEventSource.Log.RefreshedCredentials(provider.Name, true); - } - catch (Exception) - { - Callback.Invoke(false); - TimerBasedCredentialRefresherEventSource.Log.RefreshedCredentials(provider.Name, false); + // We were waiting and the registration has been disposed in meanwhile + return; } } + + try + { + provider.Refresh(); + ScheduleTimer(provider); + await Callback.Invoke(provider.Password != null).ConfigureAwait(false); + TimerBasedCredentialRefresherEventSource.Log.RefreshedCredentials(provider.Name, true); + } + catch (Exception) + { + await Callback.Invoke(false).ConfigureAwait(false); + TimerBasedCredentialRefresherEventSource.Log.RefreshedCredentials(provider.Name, false); + } }; newTimer.Enabled = true; newTimer.AutoReset = false; From 9abeada294d74b7fda7d73eb32b31d70f41b5604 Mon Sep 17 00:00:00 2001 From: Matthias Werning Date: Tue, 30 Jul 2024 16:53:02 +0200 Subject: [PATCH 7/9] Code review suggestions: Do not pass lock object down to timer registration; dispose old timer on elapse; general suggested refactorings --- .../client/api/ICredentialsRefresher.cs | 27 +++++++------------ 1 file changed, 10 insertions(+), 17 deletions(-) diff --git a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs index 3fa9462e62..9b34721b7a 100644 --- a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs +++ b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs @@ -93,7 +93,7 @@ public ICredentialsProvider Register(ICredentialsProvider provider, ICredentials return provider; } - registration = new TimerRegistration(_lockObj, callback); + registration = new TimerRegistration(callback); _registrations.Add(provider, registration); registration.ScheduleTimer(provider); @@ -107,10 +107,7 @@ public bool Unregister(ICredentialsProvider provider) { lock (_lockObj) { - if (_registrations.TryGetValue(provider, out var registration)) - { - _registrations.Remove(provider); - + if (_registrations.Remove(provider, out var registration)) { TimerBasedCredentialRefresherEventSource.Log.Unregistered(provider.Name); registration.Dispose(); return true; @@ -123,15 +120,13 @@ public bool Unregister(ICredentialsProvider provider) private class TimerRegistration : IDisposable { - private readonly object _lockObj; private System.Timers.Timer? _timer; private bool _disposed; public ICredentialsRefresher.NotifyCredentialRefreshedAsync Callback { get; set; } - public TimerRegistration(object lockObj, ICredentialsRefresher.NotifyCredentialRefreshedAsync callback) + public TimerRegistration(ICredentialsRefresher.NotifyCredentialRefreshedAsync callback) { - _lockObj = lockObj; Callback = callback; } @@ -139,7 +134,7 @@ public void ScheduleTimer(ICredentialsProvider provider) { if (provider.ValidUntil == null) { - throw new ArgumentNullException(nameof(provider.ValidUntil) + " of " + nameof(provider) + " was null"); + throw new ArgumentNullException(nameof(provider.ValidUntil) + " of " + provider.GetType().Name + " was null"); } if (_disposed) { @@ -151,14 +146,10 @@ public void ScheduleTimer(ICredentialsProvider provider) newTimer.Elapsed += async (o, e) => { TimerBasedCredentialRefresherEventSource.Log.TriggeredTimer(provider.Name); - - lock (_lockObj) + if (_disposed) { - if (_disposed) - { - // We were waiting and the registration has been disposed in meanwhile - return; - } + // We were waiting and the registration has been disposed in meanwhile + return; } try @@ -177,14 +168,16 @@ public void ScheduleTimer(ICredentialsProvider provider) newTimer.Enabled = true; newTimer.AutoReset = false; TimerBasedCredentialRefresherEventSource.Log.ScheduledTimer(provider.Name, newTimer.Interval); + var oldTimer = _timer; _timer = newTimer; + oldTimer?.Dispose(); } public void Dispose() { if (_disposed) { - throw new InvalidOperationException("registration already disposed"); + throw new ObjectDisposedException(GetType().FullName); } try From c1c3272334c1ad0e507a7bfff5421968485f256e Mon Sep 17 00:00:00 2001 From: Matthias Werning Date: Tue, 30 Jul 2024 17:02:29 +0200 Subject: [PATCH 8/9] Formatting --- projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs index 9b34721b7a..426a89375b 100644 --- a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs +++ b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs @@ -107,7 +107,8 @@ public bool Unregister(ICredentialsProvider provider) { lock (_lockObj) { - if (_registrations.Remove(provider, out var registration)) { + if (_registrations.Remove(provider, out var registration)) + { TimerBasedCredentialRefresherEventSource.Log.Unregistered(provider.Name); registration.Dispose(); return true; From 350be5016ff3640a177eda6ec7002935cbca8844 Mon Sep 17 00:00:00 2001 From: Matthias Werning Date: Tue, 30 Jul 2024 17:06:59 +0200 Subject: [PATCH 9/9] Formatting once again --- projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs index 426a89375b..ab68e0af1c 100644 --- a/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs +++ b/projects/RabbitMQ.Client/client/api/ICredentialsRefresher.cs @@ -107,7 +107,7 @@ public bool Unregister(ICredentialsProvider provider) { lock (_lockObj) { - if (_registrations.Remove(provider, out var registration)) + if (_registrations.Remove(provider, out var registration)) { TimerBasedCredentialRefresherEventSource.Log.Unregistered(provider.Name); registration.Dispose();