Regression: NetworkDataTask's ThreadSafeWeakPtrControlBlock are leaking

https://bugs.webkit.org/show_bug.cgi?id=259808
rdar://112162921

Reviewed by Alex Christensen.

The NetworkProcess memory was growing unboundedly during long-running browsing
benchmarks, which the number of NetworkDataTask's ThreadSafeWeakPtrControlBlock
objects growing very large.

The issue was with the `ThreadSafeWeakHashSet<NetworkDataTask> m_dataTaskSet`
container on the NetworkSession. We would add NetworkDataTask objects to this
WeakHashSet and never explicitly remove them. We had a comment in the code
indicating that the ThreadSafeWeakHashSet's amortized clean up would take care
of removing them. The issue though is that amortized clean up would never happen
and m_dataTaskSet's internal Set would just grow unboundedly, keeping control
blocks alive.

The reason amortized clean up never triggered is that we only add to m_dataTaskSet,
we never remove from it, never clear it and we only very rarely call forEach() on
it. As a result, ThreadSafeWeakHashSet::m_operationCountSinceLastCleanup would only
increment due to the add() calls. The condition for amortized clean up was:
`++m_operationCountSinceLastCleanup / 2 > m_set.size()`.

Since m_operationCountSinceLastCleanup would only increase due to the add() calls
and since nothing would ever get removed from the set, `m_operationCountSinceLastCleanup / 2`
could never reach the size of the set.

I am making several fixes in this patch:
1. I added a `|| m_operationCountSinceLastCleanup > m_maxOperationCountWithoutCleanup`
   check to amortized clean up to guarantee that amortized clean up eventually happens
   and that such weak containers can never grow unboundedly, no matter how they are used.
   m_maxOperationCountWithoutCleanup gets multiplied by 2 whenever the clean up finds
   nothing for proper amortization.
2. I made it so that NetworkDataTask objects remove themselves from the ThreadSafeWeakHashSet
   in their destructor. This is good practice no matter what and makes such the set stays
   small.
3. I actually found a bug in `ThreadSafeWeakHashSet::remove()` when implementing the fix in
   (2). ThreadSafeWeakHashSet::remove() would fail to remove the object from the set if the
   call if made after the object destructor has started running! The reason for this is that
   the ThreadSafeWeakHashSet was using a `std::pair<ControlBlockRefPtr, const T*>` as set
   key internally. This means that we needed to create a ControlBlockRefPtr of the object's
   control block in remove() in order to remove the object from the map. However, 265344@main
   made it so that ref'ing a control block returns nullptr after the object has started
   destruction. As a result, the first item in the key pair would be nullptr and we'd fail
   to remove. To address the issue, I am dropping the ControlBlockRefPtr from the key and
   using a `HashMap<const T*, ControlBlockRefPtr>` instead of a HashSet.

* Source/WTF/wtf/ThreadSafeWeakHashSet.h:
* Source/WTF/wtf/WeakHashMap.h:
* Source/WTF/wtf/WeakHashSet.h:
* Source/WTF/wtf/WeakListHashSet.h:
* Source/WebKit/NetworkProcess/NetworkDataTask.cpp:
(WebKit::NetworkDataTask::~NetworkDataTask):
* Source/WebKit/NetworkProcess/NetworkSession.cpp:
(WebKit::NetworkSession::registerNetworkDataTask):
(WebKit::NetworkSession::unregisterNetworkDataTask):
* Source/WebKit/NetworkProcess/NetworkSession.h:
* Tools/TestWebKitAPI/Tests/WTF/WeakPtr.cpp:
(TestWebKitAPI::ObjectAddingAndRemovingItself::create):
(TestWebKitAPI::ObjectAddingAndRemovingItself::~ObjectAddingAndRemovingItself):
(TestWebKitAPI::ObjectAddingAndRemovingItself::ObjectAddingAndRemovingItself):
(TestWebKitAPI::TEST):
(TestWebKitAPI::ObjectAddingItselfOnly::create):
(TestWebKitAPI::ObjectAddingItselfOnly::ObjectAddingItselfOnly):

Canonical link: https://commits.webkit.org/266594@main

Source/WTF/wtf/ThreadSafeWeakHashSet.h

@@ -26,7 +26,7 @@
 #pragma once
 
 #include <wtf/Algorithms.h>
-#include <wtf/HashSet.h>
+#include <wtf/HashMap.h>
 #include <wtf/ThreadSafeWeakPtr.h>
 
 namespace WTF {
@@ -100,44 +100,37 @@ class ThreadSafeWeakHashSet final {
         if (!retainedControlBlock)
             return;
         amortizedCleanupIfNeeded();
-        m_set.add({ WTFMove(retainedControlBlock), static_cast<const T*>(&value) });
+        m_map.add(static_cast<const T*>(&value), WTFMove(retainedControlBlock));
     }
 
     template<typename U, std::enable_if_t<std::is_convertible_v<U*, T*>>* = nullptr>
     bool remove(const U& value)
     {
         Locker locker { m_lock };
-        ControlBlockRefPtr retainedControlBlock { &value.controlBlock() };
-        if (!retainedControlBlock)
-            return false;
         amortizedCleanupIfNeeded();
-        return m_set.remove({ WTFMove(retainedControlBlock), static_cast<const T*>(&value) });
+        return m_map.remove(static_cast<const T*>(&value));
     }
 
     void clear()
     {
         Locker locker { m_lock };
-        m_set.clear();
+        m_map.clear();
         m_operationCountSinceLastCleanup = 0;
     }
 
     template<typename U, std::enable_if_t<std::is_convertible_v<U*, T*>>* = nullptr>
     bool contains(const U& value) const
     {
         Locker locker { m_lock };
-        ControlBlockRefPtr retainedControlBlock { &value.controlBlock() };
-        if (!retainedControlBlock)
-            return false;
         amortizedCleanupIfNeeded();
-        return m_set.contains({ WTFMove(retainedControlBlock), static_cast<const T*>(&value) });
+        return m_map.contains(static_cast<const T*>(&value));
     }
 
     bool isEmptyIgnoringNullReferences() const
     {
         Locker locker { m_lock };
         amortizedCleanupIfNeeded();
-        for (auto& pair : m_set) {
-            auto& controlBlock = pair.first;
+        for (auto& controlBlock : m_map.values()) {
             if (!controlBlock->objectHasStartedDeletion())
                 return false;
         }
@@ -149,10 +142,10 @@ class ThreadSafeWeakHashSet final {
         Vector<Ref<T>> strongReferences;
         {
             Locker locker { m_lock };
-            strongReferences.reserveInitialCapacity(m_set.size());
-            m_set.removeIf([&] (auto& pair) {
-                auto& controlBlock = pair.first;
-                auto* objectOfCorrectType = pair.second;
+            strongReferences.reserveInitialCapacity(m_map.size());
+            m_map.removeIf([&] (auto& pair) {
+                auto& controlBlock = pair.value;
+                auto* objectOfCorrectType = pair.key;
                 if (auto refPtr = controlBlock->template makeStrongReferenceIfPossible<T>(objectOfCorrectType)) {
                     strongReferences.uncheckedAppend(refPtr.releaseNonNull());
                     return false;
@@ -171,27 +164,36 @@ class ThreadSafeWeakHashSet final {
             callback(item.get());
     }
 
+    unsigned sizeIncludingEmptyEntriesForTesting()
+    {
+        Locker locker { m_lock };
+        return m_map.size();
+    }
+
 private:
     ALWAYS_INLINE void moveFrom(ThreadSafeWeakHashSet&& other)
     {
         Locker locker { m_lock };
         Locker otherLocker { other.m_lock };
-        m_set = std::exchange(other.m_set, { });
+        m_map = std::exchange(other.m_map, { });
         m_operationCountSinceLastCleanup = std::exchange(other.m_operationCountSinceLastCleanup, 0);
     }
 
+    static constexpr unsigned initialMaxOperationCountWithoutCleanup = 512;
     ALWAYS_INLINE void amortizedCleanupIfNeeded() const WTF_REQUIRES_LOCK(m_lock)
     {
-        if (++m_operationCountSinceLastCleanup / 2 > m_set.size()) {
-            m_set.removeIf([] (auto& pair) {
-                return pair.first->objectHasStartedDeletion();
+        if (++m_operationCountSinceLastCleanup / 2 > m_map.size() || m_operationCountSinceLastCleanup > m_maxOperationCountWithoutCleanup) {
+            bool didRemove = m_map.removeIf([] (auto& pair) {
+                return pair.value->objectHasStartedDeletion();
             });
+            m_maxOperationCountWithoutCleanup = didRemove ? std::max(initialMaxOperationCountWithoutCleanup, m_maxOperationCountWithoutCleanup / 2) : m_maxOperationCountWithoutCleanup * 2;
             m_operationCountSinceLastCleanup = 0;
         }
     }
 
-    mutable HashSet<std::pair<ControlBlockRefPtr, const T*>> m_set WTF_GUARDED_BY_LOCK(m_lock);
+    mutable HashMap<const T*, ControlBlockRefPtr> m_map WTF_GUARDED_BY_LOCK(m_lock);
     mutable unsigned m_operationCountSinceLastCleanup WTF_GUARDED_BY_LOCK(m_lock) { 0 };
+    mutable unsigned m_maxOperationCountWithoutCleanup WTF_GUARDED_BY_LOCK(m_lock) { initialMaxOperationCountWithoutCleanup };
     mutable Lock m_lock;
 };
 

Source/WTF/wtf/WeakHashMap.h

@@ -351,11 +351,14 @@ class WeakHashMap final {
         return currentCount;
     }
 
+    static constexpr unsigned initialMaxOperationCountWithoutCleanup = 512;
     ALWAYS_INLINE void amortizedCleanupIfNeeded(unsigned operationsPerformed = 1) const
     {
         unsigned currentCount = increaseOperationCountSinceLastCleanup(operationsPerformed);
-        if (currentCount / 2 > m_map.size())
-            const_cast<WeakHashMap&>(*this).removeNullReferences();
+        if (currentCount / 2 > m_map.size() || currentCount > m_maxOperationCountWithoutCleanup) {
+            bool didRemove = const_cast<WeakHashMap&>(*this).removeNullReferences();
+            m_maxOperationCountWithoutCleanup = didRemove ? std::max(initialMaxOperationCountWithoutCleanup, m_maxOperationCountWithoutCleanup / 2) : m_maxOperationCountWithoutCleanup * 2;
+        }
     }
 
     template <typename T>
@@ -375,6 +378,7 @@ class WeakHashMap final {
 
     WeakHashImplMap m_map;
     mutable unsigned m_operationCountSinceLastCleanup { 0 }; // FIXME: Store this as a HashTable meta data.
+    mutable unsigned m_maxOperationCountWithoutCleanup { initialMaxOperationCountWithoutCleanup };
 
     template <typename, typename, typename, typename> friend class WeakHashMapIteratorBase;
 };

Source/WTF/wtf/WeakHashSet.h

@@ -177,10 +177,11 @@ class WeakHashSet final {
 #endif
 
 private:
-    ALWAYS_INLINE void removeNullReferences()
+    ALWAYS_INLINE bool removeNullReferences()
     {
-        m_set.removeIf([] (auto& value) { return !value.get(); });
+        bool didRemove = m_set.removeIf([] (auto& value) { return !value.get(); });
         m_operationCountSinceLastCleanup = 0;
+        return didRemove;
     }
 
     ALWAYS_INLINE unsigned increaseOperationCountSinceLastCleanup(unsigned count = 1) const
@@ -189,15 +190,19 @@ class WeakHashSet final {
         return currentCount;
     }
 
+    static constexpr unsigned initialMaxOperationCountWithoutCleanup = 512;
     ALWAYS_INLINE void amortizedCleanupIfNeeded() const
     {
         unsigned currentCount = increaseOperationCountSinceLastCleanup();
-        if (currentCount / 2 > m_set.size())
-            const_cast<WeakHashSet&>(*this).removeNullReferences();
+        if (currentCount / 2 > m_set.size() || currentCount > m_maxOperationCountWithoutCleanup) {
+            bool didRemove = const_cast<WeakHashSet&>(*this).removeNullReferences();
+            m_maxOperationCountWithoutCleanup = didRemove ? std::max(initialMaxOperationCountWithoutCleanup, m_maxOperationCountWithoutCleanup / 2) : m_maxOperationCountWithoutCleanup * 2;
+        }
     }
 
     WeakPtrImplSet m_set;
     mutable unsigned m_operationCountSinceLastCleanup { 0 };
+    mutable unsigned m_maxOperationCountWithoutCleanup { initialMaxOperationCountWithoutCleanup };
 };
 
 template<typename MapFunction, typename T, typename WeakMapImpl>

Source/WTF/wtf/WeakListHashSet.h

@@ -294,16 +294,20 @@ class WeakListHashSet final {
 
     void checkConsistency() { } // To be implemented.
 
-    void removeNullReferences()
+    bool removeNullReferences()
     {
+        bool didRemove = false;
         auto it = m_set.begin();
         while (it != m_set.end()) {
             auto currentIt = it;
             ++it;
-            if (!currentIt->get())
+            if (!currentIt->get()) {
                 m_set.remove(currentIt);
+                didRemove = true;
+            }
         }
         m_operationCountSinceLastCleanup = 0;
+        return didRemove;
     }
 
 private:
@@ -313,15 +317,19 @@ class WeakListHashSet final {
         return currentCount;
     }
 
+    static constexpr unsigned initialMaxOperationCountWithoutCleanup = 512;
     ALWAYS_INLINE void amortizedCleanupIfNeeded(unsigned count = 1) const
     {
         unsigned currentCount = increaseOperationCountSinceLastCleanup(count);
-        if (currentCount / 2 > m_set.size())
-            const_cast<WeakListHashSet&>(*this).removeNullReferences();
+        if (currentCount / 2 > m_set.size() || currentCount > m_maxOperationCountWithoutCleanup) {
+            bool didRemove = const_cast<WeakListHashSet&>(*this).removeNullReferences();
+            m_maxOperationCountWithoutCleanup = didRemove ? std::max(initialMaxOperationCountWithoutCleanup, m_maxOperationCountWithoutCleanup / 2) : m_maxOperationCountWithoutCleanup * 2;
+        }
     }
 
     WeakPtrImplSet m_set;
     mutable unsigned m_operationCountSinceLastCleanup { 0 };
+    mutable unsigned m_maxOperationCountWithoutCleanup { initialMaxOperationCountWithoutCleanup };
 };
 
 template<typename MapFunction, typename T, typename WeakMapImpl>

Source/WebKit/NetworkProcess/NetworkDataTask.cpp

@@ -54,18 +54,26 @@ using namespace WebCore;
 Ref<NetworkDataTask> NetworkDataTask::create(NetworkSession& session, NetworkDataTaskClient& client, const NetworkLoadParameters& parameters)
 {
     ASSERT(!parameters.request.url().protocolIsBlob());
+    auto dataTask = [&] {
 #if PLATFORM(COCOA)
-    return NetworkDataTaskCocoa::create(session, client, parameters);
+        return NetworkDataTaskCocoa::create(session, client, parameters);
 #else
-    if (parameters.request.url().protocolIsData())
-        return NetworkDataTaskDataURL::create(session, client, parameters);
+        if (parameters.request.url().protocolIsData())
+            return NetworkDataTaskDataURL::create(session, client, parameters);
 #if USE(SOUP)
-    return NetworkDataTaskSoup::create(session, client, parameters);
+        return NetworkDataTaskSoup::create(session, client, parameters);
 #endif
 #if USE(CURL)
-    return NetworkDataTaskCurl::create(session, client, parameters);
+        return NetworkDataTaskCurl::create(session, client, parameters);
 #endif
 #endif
+    }();
+
+#if ENABLE(INSPECTOR_NETWORK_THROTTLING)
+    dataTask->setEmulatedConditions(session.bytesPerSecondLimit());
+#endif
+
+    return dataTask;
 }
 
 NetworkDataTask::NetworkDataTask(NetworkSession& session, NetworkDataTaskClient& client, const ResourceRequest& requestWithCredentials, StoredCredentialsPolicy storedCredentialsPolicy, bool shouldClearReferrerOnHTTPSToHTTPRedirect, bool dataTaskIsForMainFrameNavigation)
@@ -95,12 +103,16 @@ NetworkDataTask::NetworkDataTask(NetworkSession& session, NetworkDataTaskClient&
         scheduleFailure(FailureType::FTPDisabled);
         return;
     }
+
+    m_session->registerNetworkDataTask(*this);
 }
 
 NetworkDataTask::~NetworkDataTask()
 {
     ASSERT(RunLoop::isMain());
     ASSERT(!m_client);
+
+    m_session->unregisterNetworkDataTask(*this);
 }
 
 void NetworkDataTask::scheduleFailure(FailureType type)

Source/WebKit/NetworkProcess/NetworkDataTaskBlob.cpp

@@ -73,7 +73,6 @@ NetworkDataTaskBlob::NetworkDataTaskBlob(NetworkSession& session, NetworkDataTas
 
     m_blobData = session.blobRegistry().getBlobDataFromURL(request.url());
 
-    m_session->registerNetworkDataTask(*this);
     LOG(NetworkSession, "%p - Created NetworkDataTaskBlob for %s", this, request.url().string().utf8().data());
 }
 

Source/WebKit/NetworkProcess/NetworkDataTaskDataURL.cpp

@@ -62,7 +62,6 @@ Ref<NetworkDataTask> NetworkDataTaskDataURL::create(NetworkSession& session, Net
 NetworkDataTaskDataURL::NetworkDataTaskDataURL(NetworkSession& session, NetworkDataTaskClient& client, const NetworkLoadParameters& parameters)
     : NetworkDataTask(session, client, parameters.request, parameters.storedCredentialsPolicy, parameters.shouldClearReferrerOnHTTPSToHTTPRedirect, parameters.isMainFrameNavigation)
 {
-    m_session->registerNetworkDataTask(*this);
 }
 
 NetworkDataTaskDataURL::~NetworkDataTaskDataURL()

Source/WebKit/NetworkProcess/NetworkSession.cpp

@@ -592,13 +592,13 @@ std::unique_ptr<WebSocketTask> NetworkSession::createWebSocketTask(WebPageProxyI
 
 void NetworkSession::registerNetworkDataTask(NetworkDataTask& task)
 {
-    // Unregistration happens automatically in ThreadSafeWeakHashSet::amortizedCleanupIfNeeded.
+    ASSERT(!m_dataTaskSet.contains(task));
     m_dataTaskSet.add(task);
+}
 
-    // FIXME: This is not in a good place. It should probably be in the NetworkDataTask constructor.
-#if ENABLE(INSPECTOR_NETWORK_THROTTLING)
-    task.setEmulatedConditions(m_bytesPerSecondLimit);
-#endif
+void NetworkSession::unregisterNetworkDataTask(NetworkDataTask& task)
+{
+    m_dataTaskSet.remove(task);
 }
 
 NetworkLoadScheduler& NetworkSession::networkLoadScheduler()

Source/WebKit/NetworkProcess/NetworkSession.h

@@ -125,6 +125,7 @@ class NetworkSession
     WebCore::NetworkStorageSession* networkStorageSession() const;
 
     void registerNetworkDataTask(NetworkDataTask&);
+    void unregisterNetworkDataTask(NetworkDataTask&);
 
     void destroyPrivateClickMeasurementStore(CompletionHandler<void()>&&);
 
@@ -267,6 +268,7 @@ class NetworkSession
 #endif
 
 #if ENABLE(INSPECTOR_NETWORK_THROTTLING)
+    std::optional<int64_t> bytesPerSecondLimit() const { return m_bytesPerSecondLimit; }
     void setEmulatedConditions(std::optional<int64_t>&& bytesPerSecondLimit);
 #endif
 

Source/WebKit/NetworkProcess/cocoa/NetworkDataTaskCocoa.mm

@@ -327,8 +327,6 @@ static float toNSURLSessionTaskPriority(WebCore::ResourceLoadPriority priority)
     if (parameters.networkActivityTracker)
         m_task.get()._nw_activity = parameters.networkActivityTracker->getPlatformObject();
 #endif
-
-    m_session->registerNetworkDataTask(*this);
 }
 
 NetworkDataTaskCocoa::~NetworkDataTaskCocoa()

Source/WebKit/NetworkProcess/curl/NetworkDataTaskCurl.cpp

@@ -61,8 +61,6 @@ NetworkDataTaskCurl::NetworkDataTaskCurl(NetworkSession& session, NetworkDataTas
     , m_shouldRelaxThirdPartyCookieBlocking(parameters.shouldRelaxThirdPartyCookieBlocking)
     , m_sourceOrigin(parameters.sourceOrigin)
 {
-    m_session->registerNetworkDataTask(*this);
-
     auto request = parameters.request;
     if (request.url().protocolIsInHTTPFamily()) {
         if (m_storedCredentialsPolicy == StoredCredentialsPolicy::Use) {

Source/WebKit/NetworkProcess/soup/NetworkDataTaskSoup.cpp

@@ -65,8 +65,6 @@ NetworkDataTaskSoup::NetworkDataTaskSoup(NetworkSession& session, NetworkDataTas
     , m_sourceOrigin(parameters.sourceOrigin)
     , m_timeoutSource(RunLoop::main(), this, &NetworkDataTaskSoup::timeoutFired)
 {
-    m_session->registerNetworkDataTask(*this);
-
     auto request = parameters.request;
     if (request.url().protocolIsInHTTPFamily()) {
 #if USE(SOUP2)