[scheduler] Add ScopedTaskEnvironment::LifetimeObserver

Add an ability to subscribe to the notifications about
ScopedTaskEnviroments being created and destroyed, which should help
with content unittests, where we have multiple ScopedTaskEnvironments
changing under one SequenceManager.

R=gab@chromium.org
CC=alexclarke@chromium.org
BUG=892139

Change-Id: Ib56fd0f1b465edd3a6dd6693b36ef3d8b9234140
Reviewed-on: https://chromium-review.googlesource.com/c/1251449
Commit-Queue: Alexander Timin <altimin@chromium.org>
Reviewed-by: Gabriel Charette <gab@chromium.org>
Cr-Commit-Position: refs/heads/master@{#600097}

base/test/scoped_task_environment.cc

@@ -5,6 +5,7 @@
 #include "base/test/scoped_task_environment.h"
 
 #include "base/bind_helpers.h"
+#include "base/lazy_instance.h"
 #include "base/logging.h"
 #include "base/memory/ptr_util.h"
 #include "base/message_loop/message_loop.h"
@@ -16,6 +17,7 @@
 #include "base/task/task_scheduler/task_scheduler_impl.h"
 #include "base/test/test_mock_time_task_runner.h"
 #include "base/threading/sequence_local_storage_map.h"
+#include "base/threading/thread_local.h"
 #include "base/threading/thread_restrictions.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "base/time/time.h"
@@ -29,6 +31,9 @@ namespace test {
 
 namespace {
 
+LazyInstance<ThreadLocalPointer<ScopedTaskEnvironment::LifetimeObserver>>::Leaky
+    environment_lifetime_observer;
+
 std::unique_ptr<MessageLoop> CreateMessageLoopForMainThreadType(
     ScopedTaskEnvironment::MainThreadType main_thread_type) {
   switch (main_thread_type) {
@@ -146,6 +151,12 @@ ScopedTaskEnvironment::ScopedTaskEnvironment(
 
   if (execution_control_mode_ == ExecutionMode::QUEUED)
     CHECK(task_tracker_->DisallowRunTasks());
+
+  LifetimeObserver* observer = environment_lifetime_observer.Get().Get();
+  if (observer) {
+    observer->OnScopedTaskEnvironmentCreated(main_thread_type,
+                                             GetMainThreadTaskRunner());
+  }
 }
 
 ScopedTaskEnvironment::~ScopedTaskEnvironment() {
@@ -165,6 +176,16 @@ ScopedTaskEnvironment::~ScopedTaskEnvironment() {
   // on their main thread.
   ScopedAllowBaseSyncPrimitivesForTesting allow_waits_to_destroy_task_tracker;
   TaskScheduler::SetInstance(nullptr);
+
+  LifetimeObserver* observer = environment_lifetime_observer.Get().Get();
+  if (observer)
+    observer->OnScopedTaskEnvironmentDestroyed();
+}
+
+void ScopedTaskEnvironment::SetLifetimeObserver(
+    ScopedTaskEnvironment::LifetimeObserver* lifetime_observer) {
+  DCHECK_NE(!!environment_lifetime_observer.Get().Get(), !!lifetime_observer);
+  environment_lifetime_observer.Get().Set(lifetime_observer);
 }
 
 scoped_refptr<base::SingleThreadTaskRunner>

base/test/scoped_task_environment.h

@@ -98,6 +98,22 @@ class ScopedTaskEnvironment {
   // TaskScheduler and the (Thread|Sequenced)TaskRunnerHandle.
   ~ScopedTaskEnvironment();
 
+  class LifetimeObserver {
+   public:
+    virtual ~LifetimeObserver() = default;
+
+    virtual void OnScopedTaskEnvironmentCreated(
+        MainThreadType main_thread_type,
+        scoped_refptr<SingleThreadTaskRunner> task_runner) = 0;
+    virtual void OnScopedTaskEnvironmentDestroyed() = 0;
+  };
+
+  // Set a thread-local observer which will get notifications when
+  // a new ScopedTaskEnvironment is created or destroyed.
+  // This is needed due to peculiarities of Blink initialisation
+  // (Blink is per-test suite and ScopedTaskEnvironment is per-test).
+  static void SetLifetimeObserver(LifetimeObserver* lifetime_observer);
+
   // Returns a TaskRunner that schedules tasks on the main thread.
   scoped_refptr<base::SingleThreadTaskRunner> GetMainThreadTaskRunner();
 

base/test/scoped_task_environment_unittest.cc

@@ -18,6 +18,7 @@
 #include "base/threading/thread_task_runner_handle.h"
 #include "base/time/tick_clock.h"
 #include "build/build_config.h"
+#include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 #if defined(OS_POSIX)
@@ -303,6 +304,36 @@ TEST_F(ScopedTaskEnvironmentTest, FastForwardAdvanceTickClock) {
   EXPECT_EQ(kLongTaskDelay * 2, tick_clock->NowTicks() - tick_clock_ref);
 }
 
+namespace {
+
+class MockLifetimeObserver : public ScopedTaskEnvironment::LifetimeObserver {
+ public:
+  MockLifetimeObserver() = default;
+  ~MockLifetimeObserver() override = default;
+
+  MOCK_METHOD2(OnScopedTaskEnvironmentCreated,
+               void(ScopedTaskEnvironment::MainThreadType,
+                    scoped_refptr<SingleThreadTaskRunner>));
+  MOCK_METHOD0(OnScopedTaskEnvironmentDestroyed, void());
+};
+
+}  // namespace
+
+TEST_F(ScopedTaskEnvironmentTest, LifetimeObserver) {
+  testing::StrictMock<MockLifetimeObserver> lifetime_observer;
+  ScopedTaskEnvironment::SetLifetimeObserver(&lifetime_observer);
+
+  EXPECT_CALL(lifetime_observer,
+              OnScopedTaskEnvironmentCreated(testing::_, testing::_));
+  std::unique_ptr<ScopedTaskEnvironment> task_environment(
+      std::make_unique<ScopedTaskEnvironment>());
+  testing::Mock::VerifyAndClearExpectations(&lifetime_observer);
+
+  EXPECT_CALL(lifetime_observer, OnScopedTaskEnvironmentDestroyed());
+  task_environment.reset();
+  testing::Mock::VerifyAndClearExpectations(&lifetime_observer);
+}
+
 INSTANTIATE_TEST_CASE_P(
     MainThreadDefault,
     ScopedTaskEnvironmentTest,