From 3b06092b6a70f99b5d541eb8305bc2a5336e7649 Mon Sep 17 00:00:00 2001 From: Jonas Nyrup Date: Wed, 30 Aug 2023 23:01:14 +0200 Subject: [PATCH] Fix memory leaks with AsyncDictionaryHelper (#2297) To ensure that entries from _pendingRequests are removed when removing a key from _dictionary. --- lib/PuppeteerSharp/Browser.cs | 2 +- lib/PuppeteerSharp/ChromeTargetManager.cs | 16 +++++----- lib/PuppeteerSharp/Connection.cs | 8 ++--- lib/PuppeteerSharp/FirefoxTargetManager.cs | 9 ++---- lib/PuppeteerSharp/FrameTree.cs | 16 ++++------ .../Helpers/AsyncDictionaryHelper.cs | 30 ++++++++++++++++--- lib/PuppeteerSharp/Helpers/MultiMap.cs | 6 ++++ lib/PuppeteerSharp/Target.cs | 2 +- 8 files changed, 52 insertions(+), 37 deletions(-) diff --git a/lib/PuppeteerSharp/Browser.cs b/lib/PuppeteerSharp/Browser.cs index 53faaa747..c8863dc26 100644 --- a/lib/PuppeteerSharp/Browser.cs +++ b/lib/PuppeteerSharp/Browser.cs @@ -145,7 +145,7 @@ public bool IsClosed /// public ITarget[] Targets() - => TargetManager.GetAvailableTargets().InnerDictionary.Values.ToArray(); + => TargetManager.GetAvailableTargets().Values.ToArray(); /// public async Task CreateIncognitoBrowserContextAsync() diff --git a/lib/PuppeteerSharp/ChromeTargetManager.cs b/lib/PuppeteerSharp/ChromeTargetManager.cs index 163a16947..7318c99b3 100644 --- a/lib/PuppeteerSharp/ChromeTargetManager.cs +++ b/lib/PuppeteerSharp/ChromeTargetManager.cs @@ -19,8 +19,7 @@ internal class ChromeTargetManager : ITargetManager private readonly Func _targetFactoryFunc; private readonly Func _targetFilterFunc; private readonly ILogger _logger; - private readonly ConcurrentDictionary _availableTargetsByTargetIdDictionary = new(); - private readonly AsyncDictionaryHelper _attachedTargetsByTargetId; + private readonly AsyncDictionaryHelper _attachedTargetsByTargetId = new("Target {0} not found"); private readonly ConcurrentDictionary _attachedTargetsBySessionId = new(); private readonly ConcurrentDictionary _discoveredTargetsByTargetId = new(); private readonly ConcurrentDictionary> _targetInterceptors = new(); @@ -37,7 +36,6 @@ public ChromeTargetManager( Func targetFilterFunc, int targetDiscoveryTimeout = 0) { - _attachedTargetsByTargetId = new AsyncDictionaryHelper(_availableTargetsByTargetIdDictionary, "Target {0} not found"); _connection = connection; _targetFilterFunc = targetFilterFunc; _targetFactoryFunc = targetFactoryFunc; @@ -188,7 +186,7 @@ private void OnTargetCreated(TargetCreatedResponse e) if (e.TargetInfo.Type == TargetType.Browser && e.TargetInfo.Attached) { - if (_availableTargetsByTargetIdDictionary.ContainsKey(e.TargetInfo.TargetId)) + if (_attachedTargetsByTargetId.ContainsKey(e.TargetInfo.TargetId)) { return; } @@ -206,7 +204,7 @@ private async Task OnTargetDestroyedAsync(string messageId, TargetDestroyedRespo await EnsureTargetsIdsForInitAsync().ConfigureAwait(false); FinishInitializationIfReady(e.TargetId); - if (targetInfo?.Type == TargetType.ServiceWorker && _availableTargetsByTargetIdDictionary.TryRemove(e.TargetId, out var target)) + if (targetInfo?.Type == TargetType.ServiceWorker && _attachedTargetsByTargetId.TryRemove(e.TargetId, out var target)) { TargetGone?.Invoke(this, new TargetChangedArgs { Target = target, TargetInfo = targetInfo }); } @@ -222,7 +220,7 @@ private void OnTargetInfoChanged(TargetCreatedResponse e) _discoveredTargetsByTargetId[e.TargetInfo.TargetId] = e.TargetInfo; if (_ignoredTargets.Contains(e.TargetInfo.TargetId) || - !_availableTargetsByTargetIdDictionary.TryGetValue(e.TargetInfo.TargetId, out var target) || + !_attachedTargetsByTargetId.TryGetValue(e.TargetInfo.TargetId, out var target) || !e.TargetInfo.Attached) { return; @@ -266,7 +264,7 @@ await parent.SendAsync( await EnsureTargetsIdsForInitAsync().ConfigureAwait(false); FinishInitializationIfReady(targetInfo.TargetId); await SilentDetach().ConfigureAwait(false); - if (_availableTargetsByTargetIdDictionary.ContainsKey(targetInfo.TargetId)) + if (_attachedTargetsByTargetId.ContainsKey(targetInfo.TargetId)) { return; } @@ -286,7 +284,7 @@ await parent.SendAsync( return; } - var existingTarget = _availableTargetsByTargetIdDictionary.TryGetValue(targetInfo.TargetId, out var target); + var existingTarget = _attachedTargetsByTargetId.TryGetValue(targetInfo.TargetId, out var target); if (!existingTarget) { target = _targetFactoryFunc(targetInfo, session); @@ -384,7 +382,7 @@ private void OnDetachedFromTarget(object sender, TargetDetachedFromTargetRespons return; } - _availableTargetsByTargetIdDictionary.TryRemove(target.TargetId, out _); + _attachedTargetsByTargetId.TryRemove(target.TargetId, out _); TargetGone?.Invoke(this, new TargetChangedArgs { Target = target }); } } diff --git a/lib/PuppeteerSharp/Connection.cs b/lib/PuppeteerSharp/Connection.cs index f7cb59f41..6e8179f9e 100644 --- a/lib/PuppeteerSharp/Connection.cs +++ b/lib/PuppeteerSharp/Connection.cs @@ -23,8 +23,7 @@ public class Connection : IDisposable, ICDPConnection private readonly TaskQueue _callbackQueue = new(); private readonly ConcurrentDictionary _callbacks = new(); - private readonly ConcurrentDictionary _sessions = new(); - private readonly AsyncDictionaryHelper _asyncSessions; + private readonly AsyncDictionaryHelper _sessions = new("Session {0} not found"); private readonly List _manuallyAttached = new(); private int _lastId; @@ -38,7 +37,6 @@ internal Connection(string url, int delay, bool enqueueAsyncMessages, IConnectio _logger = LoggerFactory.CreateLogger(); MessageQueue = new AsyncMessageQueue(enqueueAsyncMessages, _logger); - _asyncSessions = new AsyncDictionaryHelper(_sessions, "Session {0} not found"); Transport.MessageReceived += Transport_MessageReceived; Transport.Closed += Transport_Closed; @@ -221,7 +219,7 @@ internal void Close(string closeReason) internal CDPSession GetSession(string sessionId) => _sessions.GetValueOrDefault(sessionId); - internal Task GetSessionAsync(string sessionId) => _asyncSessions.GetItemAsync(sessionId); + internal Task GetSessionAsync(string sessionId) => _sessions.GetItemAsync(sessionId); /// /// Releases all resource used by the object. @@ -298,7 +296,7 @@ private void ProcessIncomingMessage(ConnectionResponse obj) { var sessionId = param.SessionId; var session = new CDPSession(this, param.TargetInfo.Type, sessionId); - _asyncSessions.AddItem(sessionId, session); + _sessions.AddItem(sessionId, session); SessionAttached?.Invoke(this, new SessionEventArgs { Session = session }); diff --git a/lib/PuppeteerSharp/FirefoxTargetManager.cs b/lib/PuppeteerSharp/FirefoxTargetManager.cs index 73b27aa25..76d1c0ded 100644 --- a/lib/PuppeteerSharp/FirefoxTargetManager.cs +++ b/lib/PuppeteerSharp/FirefoxTargetManager.cs @@ -17,9 +17,7 @@ internal class FirefoxTargetManager : ITargetManager private readonly Func _targetFilterFunc; private readonly ILogger _logger; private readonly ConcurrentDictionary> _targetInterceptors = new(); - - private readonly ConcurrentDictionary _availableTargetsByTargetIdDictionary = new(); - private readonly AsyncDictionaryHelper _availableTargetsByTargetId; + private readonly AsyncDictionaryHelper _availableTargetsByTargetId = new("Target {0} not found"); private readonly ConcurrentDictionary _availableTargetsBySessionId = new(); private readonly ConcurrentDictionary _discoveredTargetsByTargetId = new(); private readonly TaskCompletionSource _initializeCompletionSource = new(); @@ -31,7 +29,6 @@ public FirefoxTargetManager( Func targetFactoryFunc, Func targetFilterFunc) { - _availableTargetsByTargetId = new AsyncDictionaryHelper(_availableTargetsByTargetIdDictionary, "Target {0} not found"); _connection = connection; _targetFilterFunc = targetFilterFunc; _targetFactoryFunc = targetFactoryFunc; @@ -148,7 +145,7 @@ private void OnTargetDestroyed(TargetDestroyedResponse e) _discoveredTargetsByTargetId.TryRemove(e.TargetId, out var targetInfo); FinishInitializationIfReady(e.TargetId); - if (_availableTargetsByTargetIdDictionary.TryGetValue(e.TargetId, out var target)) + if (_availableTargetsByTargetId.TryGetValue(e.TargetId, out var target)) { TargetGone?.Invoke(this, new TargetChangedArgs { Target = target, TargetInfo = targetInfo }); } @@ -159,7 +156,7 @@ private void OnAttachedToTarget(object sender, TargetAttachedToTargetResponse e) var parent = sender as ICDPConnection; var targetInfo = e.TargetInfo; var session = _connection.GetSession(e.SessionId) ?? throw new PuppeteerException($"Session {e.SessionId} was not created."); - var existingTarget = _availableTargetsByTargetIdDictionary.TryGetValue(targetInfo.TargetId, out var target); + var existingTarget = _availableTargetsByTargetId.TryGetValue(targetInfo.TargetId, out var target); session.MessageReceived += OnMessageReceived; _availableTargetsBySessionId.TryAdd(session.Id, target); diff --git a/lib/PuppeteerSharp/FrameTree.cs b/lib/PuppeteerSharp/FrameTree.cs index 7cf1edada..ea162c66d 100644 --- a/lib/PuppeteerSharp/FrameTree.cs +++ b/lib/PuppeteerSharp/FrameTree.cs @@ -10,24 +10,18 @@ namespace PuppeteerSharp { internal class FrameTree { - private readonly ConcurrentDictionary _frames = new(); + private readonly AsyncDictionaryHelper _frames = new("Frame {0} not found"); private readonly ConcurrentDictionary _parentIds = new(); private readonly ConcurrentDictionary> _childIds = new(); private readonly ConcurrentDictionary>> _waitRequests = new(); - private readonly AsyncDictionaryHelper _asyncFrames; - - public FrameTree() - { - _asyncFrames = new AsyncDictionaryHelper(_frames, "Frame {0} not found"); - } public Frame MainFrame { get; set; } public Frame[] Frames => _frames.Values.ToArray(); - internal Task GetFrameAsync(string frameId) => _asyncFrames.GetItemAsync(frameId); + internal Task GetFrameAsync(string frameId) => _frames.GetItemAsync(frameId); - internal Task TryGetFrameAsync(string frameId) => _asyncFrames.TryGetItemAsync(frameId); + internal Task TryGetFrameAsync(string frameId) => _frames.TryGetItemAsync(frameId); internal Frame GetById(string id) { @@ -52,7 +46,7 @@ internal Task WaitForFrameAsync(string frameId) internal void AddFrame(Frame frame) { - _asyncFrames.AddItem(frame.Id, frame); + _frames.AddItem(frame.Id, frame); if (frame.ParentId != null) { _parentIds.TryAdd(frame.Id, frame.ParentId); @@ -78,7 +72,7 @@ internal void AddFrame(Frame frame) internal void RemoveFrame(Frame frame) { - _frames.TryRemove(frame.Id, out var _); + _frames.TryRemove(frame.Id, out _); _parentIds.TryRemove(frame.Id, out var _); if (frame.ParentId != null) diff --git a/lib/PuppeteerSharp/Helpers/AsyncDictionaryHelper.cs b/lib/PuppeteerSharp/Helpers/AsyncDictionaryHelper.cs index 91785c5b6..df70938e2 100644 --- a/lib/PuppeteerSharp/Helpers/AsyncDictionaryHelper.cs +++ b/lib/PuppeteerSharp/Helpers/AsyncDictionaryHelper.cs @@ -1,5 +1,6 @@ using System; using System.Collections.Concurrent; +using System.Collections.Generic; using System.Globalization; using System.Threading.Tasks; @@ -9,15 +10,14 @@ internal class AsyncDictionaryHelper { private readonly string _timeoutMessage; private readonly MultiMap> _pendingRequests = new(); - private readonly ConcurrentDictionary _dictionary; + private readonly ConcurrentDictionary _dictionary = new(); - public AsyncDictionaryHelper(ConcurrentDictionary dictionary, string timeoutMessage) + public AsyncDictionaryHelper(string timeoutMessage) { - _dictionary = dictionary; _timeoutMessage = timeoutMessage; } - internal ConcurrentDictionary InnerDictionary => _dictionary; + internal ICollection Values => _dictionary.Values; internal async Task GetItemAsync(TKey key) { @@ -58,5 +58,27 @@ internal void AddItem(TKey key, TValue value) tcs.TrySetResult(value); } } + + internal bool TryRemove(TKey key, out TValue value) + { + var result = _dictionary.TryRemove(key, out value); + _ = _pendingRequests.TryRemove(key, out _); + return result; + } + + internal void Clear() + { + _dictionary.Clear(); + _pendingRequests.Clear(); + } + + internal TValue GetValueOrDefault(TKey key) + => _dictionary.GetValueOrDefault(key); + + internal bool TryGetValue(TKey key, out TValue value) + => _dictionary.TryGetValue(key, out value); + + internal bool ContainsKey(TKey key) + => _dictionary.ContainsKey(key); } } diff --git a/lib/PuppeteerSharp/Helpers/MultiMap.cs b/lib/PuppeteerSharp/Helpers/MultiMap.cs index b72624b6a..97fb8b52e 100644 --- a/lib/PuppeteerSharp/Helpers/MultiMap.cs +++ b/lib/PuppeteerSharp/Helpers/MultiMap.cs @@ -20,7 +20,13 @@ internal bool Has(TKey key, TValue value) internal bool Delete(TKey key, TValue value) => _map.TryGetValue(key, out var set) && set.Remove(value); + internal bool TryRemove(TKey key, out ICollection value) + => _map.TryRemove(key, out value); + internal TValue FirstValue(TKey key) => _map.TryGetValue(key, out var set) ? set.FirstOrDefault() : default; + + internal void Clear() + => _map.Clear(); } } diff --git a/lib/PuppeteerSharp/Target.cs b/lib/PuppeteerSharp/Target.cs index 05d413e23..d9f931792 100644 --- a/lib/PuppeteerSharp/Target.cs +++ b/lib/PuppeteerSharp/Target.cs @@ -88,7 +88,7 @@ internal Target( /// public ITarget Opener => TargetInfo.OpenerId != null ? - ((Browser)Browser).TargetManager.GetAvailableTargets().InnerDictionary.GetValueOrDefault(TargetInfo.OpenerId) : null; + ((Browser)Browser).TargetManager.GetAvailableTargets().GetValueOrDefault(TargetInfo.OpenerId) : null; /// public IBrowser Browser => BrowserContext.Browser;