Support SequenceLocalStorage from main thread of ScopedTaskEnvironment::MOCK_TIME

R=fdoray@chromium.org

Bug: 824770
Change-Id: Ib2e8cb5ae6c3aece489d7e27c256b18c52e89cb1
Reviewed-on: https://chromium-review.googlesource.com/975781
Commit-Queue: Gabriel Charette <gab@chromium.org>
Reviewed-by: François Doray <fdoray@chromium.org>
Cr-Commit-Position: refs/heads/master@{#545111}

base/test/scoped_task_environment.cc

@@ -15,6 +15,7 @@
 #include "base/task_scheduler/task_scheduler.h"
 #include "base/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_task_runner_handle.h"
 #include "base/time/time.h"
 
@@ -95,6 +96,16 @@ ScopedTaskEnvironment::ScopedTaskEnvironment(
               ? MakeRefCounted<TestMockTimeTaskRunner>(
                     TestMockTimeTaskRunner::Type::kBoundToThread)
               : nullptr),
+      slsm_for_mock_time_(
+          main_thread_type == MainThreadType::MOCK_TIME
+              ? std::make_unique<internal::SequenceLocalStorageMap>()
+              : nullptr),
+      slsm_registration_for_mock_time_(
+          main_thread_type == MainThreadType::MOCK_TIME
+              ? std::make_unique<
+                    internal::ScopedSetSequenceLocalStorageMapForCurrentThread>(
+                    slsm_for_mock_time_.get())
+              : nullptr),
       task_tracker_(new TestTaskTracker()) {
   CHECK(!TaskScheduler::GetInstance());
 

base/test/scoped_task_environment.h

@@ -12,6 +12,11 @@
 
 namespace base {
 
+namespace internal {
+class ScopedSetSequenceLocalStorageMapForCurrentThread;
+class SequenceLocalStorageMap;
+}  // namespace internal
+
 class MessageLoop;
 class TaskScheduler;
 class TestMockTimeTaskRunner;
@@ -126,6 +131,14 @@ class ScopedTaskEnvironment {
   const std::unique_ptr<MessageLoop> message_loop_;
   const scoped_refptr<TestMockTimeTaskRunner> mock_time_task_runner_;
 
+  // Non-null in MOCK_TIME, where an explicit SequenceLocalStorageMap needs to
+  // be provided. TODO(gab): This can be removed once mock time support is added
+  // to MessageLoop directly.
+  const std::unique_ptr<internal::SequenceLocalStorageMap> slsm_for_mock_time_;
+  const std::unique_ptr<
+      internal::ScopedSetSequenceLocalStorageMapForCurrentThread>
+      slsm_registration_for_mock_time_;
+
   const TaskScheduler* task_scheduler_ = nullptr;
 
   // Owned by |task_scheduler_|.

base/test/scoped_task_environment_unittest.cc

@@ -14,6 +14,7 @@
 #include "base/task_scheduler/post_task.h"
 #include "base/test/test_timeouts.h"
 #include "base/threading/platform_thread.h"
+#include "base/threading/sequence_local_storage_slot.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "base/time/tick_clock.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -230,6 +231,16 @@ TEST_P(ScopedTaskEnvironmentTest, DelayedTasks) {
   }
 }
 
+// Regression test for https://crbug.com/824770.
+TEST_P(ScopedTaskEnvironmentTest, SupportsSequenceLocalStorageOnMainThread) {
+  ScopedTaskEnvironment scoped_task_environment(
+      GetParam(), ScopedTaskEnvironment::ExecutionMode::ASYNC);
+
+  SequenceLocalStorageSlot<int> sls_slot;
+  sls_slot.Set(5);
+  EXPECT_EQ(5, sls_slot.Get());
+}
+
 // Verify that the TickClock returned by
 // |ScopedTaskEnvironment::GetMockTickClock| gets updated when the
 // FastForward(By|UntilNoTasksRemain) functions are called.