From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Yutaka Hirano <yhirano@chromium.org>
Date: Fri, 11 Jun 2021 08:42:15 +0000
Subject: Reduce memory consumption on LongTaskObserver::DidProcessTask

https://crrev.com/c/2909934 fixed a security issue, but it introduced a
copy operation for each DidProcessTask for a long task. We see a memory
regression on the change, and this is an attempt to mitigate the
regression.

(cherry picked from commit 8097e73295a88e64d8318d982847a5e4f2bcc4d2)

(cherry picked from commit 7be6a34fe2f01af881bb074bc616bf5b6b5f7c31)

Bug: 1210487, 1211539
Change-Id: Ib9101e29d70fadb11b7967754e847bb5cc754feb
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2915153
Reviewed-by: Benoit L <lizeb@chromium.org>
Commit-Queue: Yutaka Hirano <yhirano@chromium.org>
Cr-Original-Original-Commit-Position: refs/heads/master@{#886221}
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2944320
Bot-Commit: Rubber Stamper <rubber-stamper@appspot.gserviceaccount.com>
Cr-Original-Commit-Position: refs/branch-heads/4472@{#1460}
Cr-Original-Branched-From: 3d60439cfb36485e76a1c5bb7f513d3721b20da1-refs/heads/master@{#870763}
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2948750
Owners-Override: Victor-Gabriel Savu <vsavu@google.com>
Reviewed-by: Artem Sumaneev <asumaneev@google.com>
Commit-Queue: Victor-Gabriel Savu <vsavu@google.com>
Cr-Commit-Position: refs/branch-heads/4430@{#1520}
Cr-Branched-From: e5ce7dc4f7518237b3d9bb93cccca35d25216cbe-refs/heads/master@{#857950}

diff --git a/third_party/blink/renderer/core/loader/long_task_detector.cc b/third_party/blink/renderer/core/loader/long_task_detector.cc
index f040ae5053265fb136c629f106aeefa0b01130f1..3816779cfafaef06295734b4a8a2f033bf752691 100644
--- a/third_party/blink/renderer/core/loader/long_task_detector.cc
+++ b/third_party/blink/renderer/core/loader/long_task_detector.cc
@@ -24,6 +24,7 @@ LongTaskDetector::LongTaskDetector() = default;
 void LongTaskDetector::RegisterObserver(LongTaskObserver* observer) {
   DCHECK(IsMainThread());
   DCHECK(observer);
+  DCHECK(!iterating_);
   if (observers_.insert(observer).is_new_entry && observers_.size() == 1) {
     // Number of observers just became non-zero.
     Thread::Current()->AddTaskTimeObserver(this);
@@ -32,6 +33,10 @@ void LongTaskDetector::RegisterObserver(LongTaskObserver* observer) {
 
 void LongTaskDetector::UnregisterObserver(LongTaskObserver* observer) {
   DCHECK(IsMainThread());
+  if (iterating_) {
+    observers_to_be_removed_.push_back(observer);
+    return;
+  }
   observers_.erase(observer);
   if (observers_.size() == 0) {
     Thread::Current()->RemoveTaskTimeObserver(this);
@@ -43,16 +48,21 @@ void LongTaskDetector::DidProcessTask(base::TimeTicks start_time,
   if ((end_time - start_time) < LongTaskDetector::kLongTaskThreshold)
     return;
 
-  // We copy `observers_` because it might be mutated in OnLongTaskDetected,
-  // and container mutation is not allowed during iteration.
-  const HeapHashSet<Member<LongTaskObserver>> observers = observers_;
-  for (auto& observer : observers) {
+  iterating_ = true;
+  for (auto& observer : observers_) {
     observer->OnLongTaskDetected(start_time, end_time);
   }
+  iterating_ = false;
+
+  for (const auto& observer : observers_to_be_removed_) {
+    UnregisterObserver(observer);
+  }
+  observers_to_be_removed_.clear();
 }
 
 void LongTaskDetector::Trace(Visitor* visitor) const {
   visitor->Trace(observers_);
+  visitor->Trace(observers_to_be_removed_);
 }
 
 }  // namespace blink
diff --git a/third_party/blink/renderer/core/loader/long_task_detector.h b/third_party/blink/renderer/core/loader/long_task_detector.h
index dc6f0dbab5c059b83bfe4212f0126e9690ab1a7f..5fd4bb8d2abcc67dd4e47927daa260fa37bc4aca 100644
--- a/third_party/blink/renderer/core/loader/long_task_detector.h
+++ b/third_party/blink/renderer/core/loader/long_task_detector.h
@@ -49,6 +49,8 @@ class CORE_EXPORT LongTaskDetector final
                       base::TimeTicks end_time) override;
 
   HeapHashSet<Member<LongTaskObserver>> observers_;
+  HeapVector<Member<LongTaskObserver>> observers_to_be_removed_;
+  bool iterating_ = false;
 
   DISALLOW_COPY_AND_ASSIGN(LongTaskDetector);
 };
diff --git a/third_party/blink/renderer/core/loader/long_task_detector_test.cc b/third_party/blink/renderer/core/loader/long_task_detector_test.cc
index 04959d3b682ddbf40577adc5799fe57a9ae9d500..403ba452362a3fa2a6b24f238bad35d9eb1bac0c 100644
--- a/third_party/blink/renderer/core/loader/long_task_detector_test.cc
+++ b/third_party/blink/renderer/core/loader/long_task_detector_test.cc
@@ -39,6 +39,8 @@ class SelfUnregisteringObserver
   }
   bool IsCalled() const { return called_; }
 
+  void Reset() { called_ = false; }
+
  private:
   bool called_ = false;
 };
@@ -148,6 +150,11 @@ TEST_F(LongTaskDetectorTest, SelfUnregisteringObserver) {
   SimulateTask(LongTaskDetector::kLongTaskThreshold +
                base::TimeDelta::FromMilliseconds(10));
   EXPECT_TRUE(observer->IsCalled());
+  observer->Reset();
+
+  SimulateTask(LongTaskDetector::kLongTaskThreshold +
+               base::TimeDelta::FromMilliseconds(10));
+  EXPECT_FALSE(observer->IsCalled());
 }
 
 }  // namespace blink