Skip to content

Commit f62ec11

Browse files
jnyrupjnyrup
committed
Fix memory leaks with AsyncDictionaryHelper
To ensure that entries from _pendingRequests are removed when removing a key from _dictionary.
1 parent 6d712e9 commit f62ec11

File tree

8 files changed

+52
-37
lines changed

8 files changed

+52
-37
lines changed

‎lib/PuppeteerSharp/Browser.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -148,7 +148,7 @@ public bool IsClosed
148148

149149
/// <inheritdoc/>
150150
public ITarget[] Targets()
151-
=> TargetManager.GetAvailableTargets().InnerDictionary.Values.ToArray();
151+
=> TargetManager.GetAvailableTargets().Values.ToArray();
152152

153153
/// <inheritdoc/>
154154
public async Task<IBrowserContext> CreateIncognitoBrowserContextAsync()

‎lib/PuppeteerSharp/ChromeTargetManager.cs‎

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,7 @@ internal class ChromeTargetManager : ITargetManager
1919
private readonly Func<TargetInfo, CDPSession, Target> _targetFactoryFunc;
2020
private readonly Func<TargetInfo, bool> _targetFilterFunc;
2121
private readonly ILogger<ChromeTargetManager> _logger;
22-
private readonly ConcurrentDictionary<string, Target> _availableTargetsByTargetIdDictionary = new();
23-
private readonly AsyncDictionaryHelper<string, Target> _attachedTargetsByTargetId;
22+
private readonly AsyncDictionaryHelper<string, Target> _attachedTargetsByTargetId = new("Target {0} not found");
2423
private readonly ConcurrentDictionary<string, Target> _attachedTargetsBySessionId = new();
2524
private readonly ConcurrentDictionary<string, TargetInfo> _discoveredTargetsByTargetId = new();
2625
private readonly ConcurrentDictionary<ICDPConnection, List<TargetInterceptor>> _targetInterceptors = new();
@@ -37,7 +36,6 @@ public ChromeTargetManager(
3736
Func<TargetInfo, bool> targetFilterFunc,
3837
int targetDiscoveryTimeout = 0)
3938
{
40-
_attachedTargetsByTargetId = new AsyncDictionaryHelper<string, Target>(_availableTargetsByTargetIdDictionary, "Target {0} not found");
4139
_connection = connection;
4240
_targetFilterFunc = targetFilterFunc;
4341
_targetFactoryFunc = targetFactoryFunc;
@@ -188,7 +186,7 @@ private void OnTargetCreated(TargetCreatedResponse e)
188186

189187
if (e.TargetInfo.Type == TargetType.Browser && e.TargetInfo.Attached)
190188
{
191-
if (_availableTargetsByTargetIdDictionary.ContainsKey(e.TargetInfo.TargetId))
189+
if (_attachedTargetsByTargetId.ContainsKey(e.TargetInfo.TargetId))
192190
{
193191
return;
194192
}
@@ -204,7 +202,7 @@ private async void OnTargetDestroyed(TargetDestroyedResponse e)
204202
await EnsureTargetsIdsForInit().ConfigureAwait(false);
205203
FinishInitializationIfReady(e.TargetId);
206204

207-
if (targetInfo?.Type == TargetType.ServiceWorker && _availableTargetsByTargetIdDictionary.TryRemove(e.TargetId, out var target))
205+
if (targetInfo?.Type == TargetType.ServiceWorker && _attachedTargetsByTargetId.TryRemove(e.TargetId, out var target))
208206
{
209207
TargetGone?.Invoke(this, new TargetChangedArgs { Target = target, TargetInfo = targetInfo });
210208
}
@@ -215,7 +213,7 @@ private void OnTargetInfoChanged(TargetCreatedResponse e)
215213
_discoveredTargetsByTargetId[e.TargetInfo.TargetId] = e.TargetInfo;
216214

217215
if (_ignoredTargets.Contains(e.TargetInfo.TargetId) ||
218-
!_availableTargetsByTargetIdDictionary.TryGetValue(e.TargetInfo.TargetId, out var target) ||
216+
!_attachedTargetsByTargetId.TryGetValue(e.TargetInfo.TargetId, out var target) ||
219217
!e.TargetInfo.Attached)
220218
{
221219
return;
@@ -259,7 +257,7 @@ await parent.SendAsync(
259257
await EnsureTargetsIdsForInit().ConfigureAwait(false);
260258
FinishInitializationIfReady(targetInfo.TargetId);
261259
await SilentDetach().ConfigureAwait(false);
262-
if (_availableTargetsByTargetIdDictionary.ContainsKey(targetInfo.TargetId))
260+
if (_attachedTargetsByTargetId.ContainsKey(targetInfo.TargetId))
263261
{
264262
return;
265263
}
@@ -279,7 +277,7 @@ await parent.SendAsync(
279277
return;
280278
}
281279

282-
var existingTarget = _availableTargetsByTargetIdDictionary.TryGetValue(targetInfo.TargetId, out var target);
280+
var existingTarget = _attachedTargetsByTargetId.TryGetValue(targetInfo.TargetId, out var target);
283281
if (!existingTarget)
284282
{
285283
target = _targetFactoryFunc(targetInfo, session);
@@ -377,7 +375,7 @@ private void OnDetachedFromTarget(object sender, TargetDetachedFromTargetRespons
377375
return;
378376
}
379377

380-
_availableTargetsByTargetIdDictionary.TryRemove(target.TargetId, out _);
378+
_attachedTargetsByTargetId.TryRemove(target.TargetId, out _);
381379
TargetGone?.Invoke(this, new TargetChangedArgs { Target = target });
382380
}
383381
}

‎lib/PuppeteerSharp/Connection.cs‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,7 @@ public class Connection : IDisposable, ICDPConnection
2323
private readonly TaskQueue _callbackQueue = new();
2424

2525
private readonly ConcurrentDictionary<int, MessageTask> _callbacks = new();
26-
private readonly ConcurrentDictionary<string, CDPSession> _sessions = new();
27-
private readonly AsyncDictionaryHelper<string, CDPSession> _asyncSessions;
26+
private readonly AsyncDictionaryHelper<string, CDPSession> _sessions = new("Session {0} not found");
2827
private readonly List<string> _manuallyAttached = new();
2928
private int _lastId;
3029

@@ -38,7 +37,6 @@ internal Connection(string url, int delay, bool enqueueAsyncMessages, IConnectio
3837
_logger = LoggerFactory.CreateLogger<Connection>();
3938

4039
MessageQueue = new AsyncMessageQueue(enqueueAsyncMessages, _logger);
41-
_asyncSessions = new AsyncDictionaryHelper<string, CDPSession>(_sessions, "Session {0} not found");
4240

4341
Transport.MessageReceived += Transport_MessageReceived;
4442
Transport.Closed += Transport_Closed;
@@ -221,7 +219,7 @@ internal void Close(string closeReason)
221219

222220
internal CDPSession GetSession(string sessionId) => _sessions.GetValueOrDefault(sessionId);
223221

224-
internal Task<CDPSession> GetSessionAsync(string sessionId) => _asyncSessions.GetItemAsync(sessionId);
222+
internal Task<CDPSession> GetSessionAsync(string sessionId) => _sessions.GetItemAsync(sessionId);
225223

226224
/// <summary>
227225
/// Releases all resource used by the <see cref="Connection"/> object.
@@ -287,7 +285,7 @@ private void ProcessIncomingMessage(ConnectionResponse obj)
287285
{
288286
var sessionId = param.SessionId;
289287
var session = new CDPSession(this, param.TargetInfo.Type, sessionId);
290-
_asyncSessions.AddItem(sessionId, session);
288+
_sessions.AddItem(sessionId, session);
291289

292290
SessionAttached?.Invoke(this, new SessionEventArgs { Session = session });
293291

‎lib/PuppeteerSharp/FirefoxTargetManager.cs‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,7 @@ internal class FirefoxTargetManager : ITargetManager
1717
private readonly Func<TargetInfo, bool> _targetFilterFunc;
1818
private readonly ILogger<FirefoxTargetManager> _logger;
1919
private readonly ConcurrentDictionary<ICDPConnection, List<TargetInterceptor>> _targetInterceptors = new();
20-
21-
private readonly ConcurrentDictionary<string, Target> _availableTargetsByTargetIdDictionary = new();
22-
private readonly AsyncDictionaryHelper<string, Target> _availableTargetsByTargetId;
20+
private readonly AsyncDictionaryHelper<string, Target> _availableTargetsByTargetId = new("Target {0} not found");
2321
private readonly ConcurrentDictionary<string, Target> _availableTargetsBySessionId = new();
2422
private readonly ConcurrentDictionary<string, TargetInfo> _discoveredTargetsByTargetId = new();
2523
private readonly TaskCompletionSource<bool> _initializeCompletionSource = new();
@@ -31,7 +29,6 @@ public FirefoxTargetManager(
3129
Func<TargetInfo, CDPSession, Target> targetFactoryFunc,
3230
Func<TargetInfo, bool> targetFilterFunc)
3331
{
34-
_availableTargetsByTargetId = new AsyncDictionaryHelper<string, Target>(_availableTargetsByTargetIdDictionary, "Target {0} not found");
3532
_connection = connection;
3633
_targetFilterFunc = targetFilterFunc;
3734
_targetFactoryFunc = targetFactoryFunc;
@@ -148,7 +145,7 @@ private void OnTargetDestroyed(TargetDestroyedResponse e)
148145
_discoveredTargetsByTargetId.TryRemove(e.TargetId, out var targetInfo);
149146
FinishInitializationIfReady(e.TargetId);
150147

151-
if (_availableTargetsByTargetIdDictionary.TryGetValue(e.TargetId, out var target))
148+
if (_availableTargetsByTargetId.TryGetValue(e.TargetId, out var target))
152149
{
153150
TargetGone?.Invoke(this, new TargetChangedArgs { Target = target, TargetInfo = targetInfo });
154151
}
@@ -159,7 +156,7 @@ private void OnAttachedToTarget(object sender, TargetAttachedToTargetResponse e)
159156
var parent = sender as ICDPConnection;
160157
var targetInfo = e.TargetInfo;
161158
var session = _connection.GetSession(e.SessionId) ?? throw new PuppeteerException($"Session {e.SessionId} was not created.");
162-
var existingTarget = _availableTargetsByTargetIdDictionary.TryGetValue(targetInfo.TargetId, out var target);
159+
var existingTarget = _availableTargetsByTargetId.TryGetValue(targetInfo.TargetId, out var target);
163160
session.MessageReceived += OnMessageReceived;
164161

165162
_availableTargetsBySessionId.TryAdd(session.Id, target);

‎lib/PuppeteerSharp/FrameTree.cs‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -10,24 +10,18 @@ namespace PuppeteerSharp
1010
{
1111
internal class FrameTree
1212
{
13-
private readonly ConcurrentDictionary<string, Frame> _frames = new();
1413
private readonly ConcurrentDictionary<string, string> _parentIds = new();
1514
private readonly ConcurrentDictionary<string, List<string>> _childIds = new();
1615
private readonly ConcurrentDictionary<string, List<TaskCompletionSource<Frame>>> _waitRequests = new();
17-
private readonly AsyncDictionaryHelper<string, Frame> _asyncFrames;
18-
19-
public FrameTree()
20-
{
21-
_asyncFrames = new AsyncDictionaryHelper<string, Frame>(_frames, "Frame {0} not found");
22-
}
16+
private readonly AsyncDictionaryHelper<string, Frame> _frames = new("Frame {0} not found");
2317

2418
public Frame MainFrame { get; set; }
2519

2620
public Frame[] Frames => _frames.Values.ToArray();
2721

28-
internal Task<Frame> GetFrameAsync(string frameId) => _asyncFrames.GetItemAsync(frameId);
22+
internal Task<Frame> GetFrameAsync(string frameId) => _frames.GetItemAsync(frameId);
2923

30-
internal Task<Frame> TryGetFrameAsync(string frameId) => _asyncFrames.TryGetItemAsync(frameId);
24+
internal Task<Frame> TryGetFrameAsync(string frameId) => _frames.TryGetItemAsync(frameId);
3125

3226
internal Frame GetById(string id)
3327
{
@@ -52,7 +46,7 @@ internal Task<Frame> WaitForFrameAsync(string frameId)
5246

5347
internal void AddFrame(Frame frame)
5448
{
55-
_asyncFrames.AddItem(frame.Id, frame);
49+
_frames.AddItem(frame.Id, frame);
5650
if (frame.ParentId != null)
5751
{
5852
_parentIds.TryAdd(frame.Id, frame.ParentId);
@@ -78,7 +72,7 @@ internal void AddFrame(Frame frame)
7872

7973
internal void RemoveFrame(Frame frame)
8074
{
81-
_frames.TryRemove(frame.Id, out var _);
75+
_frames.TryRemove(frame.Id, out _);
8276
_parentIds.TryRemove(frame.Id, out var _);
8377

8478
if (frame.ParentId != null)

‎lib/PuppeteerSharp/Helpers/AsyncDictionaryHelper.cs‎

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using System;
22
using System.Collections.Concurrent;
3+
using System.Collections.Generic;
34
using System.Globalization;
45
using System.Threading.Tasks;
56

@@ -9,15 +10,14 @@ internal class AsyncDictionaryHelper<TKey, TValue>
910
{
1011
private readonly string _timeoutMessage;
1112
private readonly MultiMap<TKey, TaskCompletionSource<TValue>> _pendingRequests = new();
12-
private readonly ConcurrentDictionary<TKey, TValue> _dictionary;
13+
private readonly ConcurrentDictionary<TKey, TValue> _dictionary = new();
1314

14-
public AsyncDictionaryHelper(ConcurrentDictionary<TKey, TValue> dictionary, string timeoutMessage)
15+
public AsyncDictionaryHelper(string timeoutMessage)
1516
{
16-
_dictionary = dictionary;
1717
_timeoutMessage = timeoutMessage;
1818
}
1919

20-
internal ConcurrentDictionary<TKey, TValue> InnerDictionary => _dictionary;
20+
internal ICollection<TValue> Values => _dictionary.Values;
2121

2222
internal async Task<TValue> GetItemAsync(TKey key)
2323
{
@@ -58,5 +58,27 @@ internal void AddItem(TKey key, TValue value)
5858
tcs.TrySetResult(value);
5959
}
6060
}
61+
62+
internal bool TryRemove(TKey key, out TValue value)
63+
{
64+
var result = _dictionary.TryRemove(key, out value);
65+
_ = _pendingRequests.TryRemove(key, out _);
66+
return result;
67+
}
68+
69+
internal void Clear()
70+
{
71+
_dictionary.Clear();
72+
_pendingRequests.Clear();
73+
}
74+
75+
internal TValue GetValueOrDefault(TKey key)
76+
=> _dictionary.GetValueOrDefault(key);
77+
78+
internal bool TryGetValue(TKey key, out TValue value)
79+
=> _dictionary.TryGetValue(key, out value);
80+
81+
internal bool ContainsKey(TKey key)
82+
=> _dictionary.ContainsKey(key);
6183
}
6284
}

‎lib/PuppeteerSharp/Helpers/MultiMap.cs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,13 @@ internal bool Has(TKey key, TValue value)
2020
internal bool Delete(TKey key, TValue value)
2121
=> _map.TryGetValue(key, out var set) && set.Remove(value);
2222

23+
internal bool TryRemove(TKey key, out ICollection<TValue> value)
24+
=> _map.TryRemove(key, out value);
25+
2326
internal TValue FirstValue(TKey key)
2427
=> _map.TryGetValue(key, out var set) ? set.FirstOrDefault() : default;
28+
29+
internal void Clear()
30+
=> _map.Clear();
2531
}
2632
}

‎lib/PuppeteerSharp/Target.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ internal Target(
8888

8989
/// <inheritdoc/>
9090
public ITarget Opener => TargetInfo.OpenerId != null ?
91-
((Browser)Browser).TargetManager.GetAvailableTargets().InnerDictionary.GetValueOrDefault(TargetInfo.OpenerId) : null;
91+
((Browser)Browser).TargetManager.GetAvailableTargets().GetValueOrDefault(TargetInfo.OpenerId) : null;
9292

9393
/// <inheritdoc/>
9494
public IBrowser Browser => BrowserContext.Browser;

0 commit comments

Comments
 (0)