Set SequenceToken for SequenceLocalStorageMap destructor

Ensure that the SequenceLocalStorageMap and consequently objects still
stored in the sequence-local storage are destroyed while the owning
Sequence's SequenceToken is set for the current thread.
This makes sure that objects stored in the sequence-local storage which
expect to be destroyed on the same Sequence they've been created on do
not (D)CHECK in their destrutor.

Bug: 869272
Test: base_unittests && browser_tests --gtest_filter=*UserAffiliation*
Change-Id: I8c76b274a746a35c4ad038e7d8f04ea212ef1514
Reviewed-on: https://chromium-review.googlesource.com/1156510
Commit-Queue: Pavol Marko <pmarko@chromium.org>
Reviewed-by: Gabriel Charette <gab@chromium.org>
Cr-Commit-Position: refs/heads/master@{#579834}

base/sequence_token.cc

@@ -89,4 +89,21 @@ ScopedSetSequenceTokenForCurrentThread::
   tls_current_task_token.Get().Set(nullptr);
 }
 
+ScopedSetNestedSequenceTokenForDestructorForCurrentThread ::
+    ScopedSetNestedSequenceTokenForDestructorForCurrentThread(
+        const SequenceToken& sequence_token)
+    : original_sequence_token_(tls_current_sequence_token.Get().Get()),
+      original_task_token_(tls_current_task_token.Get().Get()),
+      sequence_token_(sequence_token),
+      task_token_(TaskToken::Create()) {
+  tls_current_sequence_token.Get().Set(&sequence_token_);
+  tls_current_task_token.Get().Set(&task_token_);
+}
+
+ScopedSetNestedSequenceTokenForDestructorForCurrentThread ::
+    ~ScopedSetNestedSequenceTokenForDestructorForCurrentThread() {
+  tls_current_sequence_token.Get().Set(original_sequence_token_);
+  tls_current_task_token.Get().Set(original_task_token_);
+}
+
 }  // namespace base

base/sequence_token.h

@@ -80,6 +80,7 @@ class BASE_EXPORT TaskToken {
 
  private:
   friend class ScopedSetSequenceTokenForCurrentThread;
+  friend class ScopedSetNestedSequenceTokenForDestructorForCurrentThread;
 
   explicit TaskToken(int token) : token_(token) {}
 
@@ -110,6 +111,26 @@ class BASE_EXPORT ScopedSetSequenceTokenForCurrentThread {
   DISALLOW_COPY_AND_ASSIGN(ScopedSetSequenceTokenForCurrentThread);
 };
 
+// Instantiate this in the scope where it is necessary to temporarily pretend to
+// be running on another sequence to run its clean-up tasks.
+class BASE_EXPORT ScopedSetNestedSequenceTokenForDestructorForCurrentThread {
+ public:
+  // Identical to ScopedSetSequenceTokenForCurrentThread but allows nesting.
+  ScopedSetNestedSequenceTokenForDestructorForCurrentThread(
+      const SequenceToken& sequence_token);
+  ~ScopedSetNestedSequenceTokenForDestructorForCurrentThread();
+
+ private:
+  const SequenceToken* const original_sequence_token_;
+  const TaskToken* const original_task_token_;
+
+  const SequenceToken sequence_token_;
+  const TaskToken task_token_;
+
+  DISALLOW_COPY_AND_ASSIGN(
+      ScopedSetNestedSequenceTokenForDestructorForCurrentThread);
+};
+
 }  // namespace base
 
 #endif  // BASE_SEQUENCE_TOKEN_H_

base/sequence_token_unittest.cc

@@ -70,6 +70,44 @@ TEST(SequenceTokenTest, ToInternalValue) {
   EXPECT_NE(token1.ToInternalValue(), token2.ToInternalValue());
 }
 
+TEST(SequenceTokenTest, ReturnsToOriginalAfterNestedSequenceToken) {
+  const SequenceToken token1 = SequenceToken::Create();
+
+  {
+    ScopedSetSequenceTokenForCurrentThread
+        scoped_set_sequence_token_for_current_thread(token1);
+    EXPECT_TRUE(SequenceToken::GetForCurrentThread().IsValid());
+    EXPECT_EQ(token1, SequenceToken::GetForCurrentThread());
+
+    {
+      const SequenceToken token2 = SequenceToken::Create();
+      ScopedSetNestedSequenceTokenForDestructorForCurrentThread
+          scoped_set_nested_sequence_token_for_current_thread(token2);
+      EXPECT_TRUE(SequenceToken::GetForCurrentThread().IsValid());
+      EXPECT_EQ(token2, SequenceToken::GetForCurrentThread());
+    }
+
+    EXPECT_TRUE(SequenceToken::GetForCurrentThread().IsValid());
+    EXPECT_EQ(token1, SequenceToken::GetForCurrentThread());
+  }
+
+  EXPECT_FALSE(SequenceToken::GetForCurrentThread().IsValid());
+}
+
+TEST(SequenceTokenTest, ReturnsToInvalidAfterNestedSequenceToken) {
+  EXPECT_FALSE(SequenceToken::GetForCurrentThread().IsValid());
+
+  {
+    const SequenceToken token = SequenceToken::Create();
+    ScopedSetNestedSequenceTokenForDestructorForCurrentThread
+        scoped_set_nested_sequence_token_for_current_thread(token);
+    EXPECT_TRUE(SequenceToken::GetForCurrentThread().IsValid());
+    EXPECT_EQ(token, SequenceToken::GetForCurrentThread());
+  }
+
+  EXPECT_FALSE(SequenceToken::GetForCurrentThread().IsValid());
+}
+
 // Expect a default-constructed TaskToken to be invalid and not equal to
 // another invalid TaskToken.
 TEST(TaskTokenTest, InvalidDefaultConstructed) {
@@ -130,4 +168,41 @@ TEST(TaskTokenTest, NotEqualInDifferentScopes) {
   EXPECT_NE(token_a, token_b);
 }
 
+TEST(TaskToken, ReturnsToOriginalAfterNestedTaskToken) {
+  TaskToken token_outer;
+  TaskToken token_inner;
+
+  {
+    ScopedSetSequenceTokenForCurrentThread
+        scoped_set_sequence_token_for_current_thread(SequenceToken::Create());
+    token_outer = TaskToken::GetForCurrentThread();
+    EXPECT_TRUE(token_outer.IsValid());
+
+    {
+      ScopedSetNestedSequenceTokenForDestructorForCurrentThread
+          scoped_set_nested_sequence_token_for_current_thread(
+              SequenceToken::Create());
+      token_inner = TaskToken::GetForCurrentThread();
+      EXPECT_TRUE(token_inner.IsValid());
+      EXPECT_NE(token_outer, token_inner);
+    }
+
+    EXPECT_EQ(token_outer, TaskToken::GetForCurrentThread());
+  }
+
+  EXPECT_FALSE(TaskToken::GetForCurrentThread().IsValid());
+}
+
+TEST(TaskToken, ReturnsToInvalidAfterNestedTaskToken) {
+  EXPECT_FALSE(SequenceToken::GetForCurrentThread().IsValid());
+
+  {
+    ScopedSetNestedSequenceTokenForDestructorForCurrentThread
+        scoped_set_nested_sequence_token_for_current_thread(
+            SequenceToken::Create());
+    EXPECT_TRUE(TaskToken::GetForCurrentThread().IsValid());
+  }
+
+  EXPECT_FALSE(TaskToken::GetForCurrentThread().IsValid());
+}
 }  // namespace base

base/task_scheduler/sequence.cc

@@ -12,7 +12,7 @@
 namespace base {
 namespace internal {
 
-Sequence::Sequence() = default;
+Sequence::Sequence() : sequence_local_storage_(base::in_place) {}
 
 bool Sequence::PushTask(Task task) {
   // Use CHECK instead of DCHECK to crash earlier. See http://crbug.com/711167
@@ -74,7 +74,16 @@ SequenceSortKey Sequence::GetSortKey() const {
   return SequenceSortKey(priority, next_task_sequenced_time);
 }
 
-Sequence::~Sequence() = default;
+Sequence::~Sequence() {
+  // Destroy the |sequence_local_storage_| with our |token_| set for the
+  // current thread, because objects stored in the sequence-local storage may
+  // expect to be destroyed on the same sequence they've been created on.
+  // This is fine, because this sequence can not be in use anywhere else when
+  // its destructor is running.
+  ScopedSetNestedSequenceTokenForDestructorForCurrentThread
+      scoped_set_sequence_token_for_current_thread(token_);
+  sequence_local_storage_.reset();
+}
 
 }  // namespace internal
 }  // namespace base

base/task_scheduler/sequence.h

@@ -71,7 +71,7 @@ class BASE_EXPORT Sequence : public RefCountedThreadSafe<Sequence> {
   const SequenceToken& token() const { return token_; }
 
   SequenceLocalStorageMap* sequence_local_storage() {
-    return &sequence_local_storage_;
+    return &sequence_local_storage_.value();
   }
 
  private:
@@ -91,7 +91,9 @@ class BASE_EXPORT Sequence : public RefCountedThreadSafe<Sequence> {
       {};
 
   // Holds data stored through the SequenceLocalStorageSlot API.
-  SequenceLocalStorageMap sequence_local_storage_;
+  // Valid throughout this Sequence's lifetime (Optional so it can be released
+  // in a custom scope in the destructor).
+  base::Optional<SequenceLocalStorageMap> sequence_local_storage_;
 
   DISALLOW_COPY_AND_ASSIGN(Sequence);
 };

base/threading/sequence_local_storage_map.cc

@@ -18,9 +18,13 @@ LazyInstance<ThreadLocalPointer<SequenceLocalStorageMap>>::Leaky
     tls_current_sequence_local_storage = LAZY_INSTANCE_INITIALIZER;
 }  // namespace
 
-SequenceLocalStorageMap::SequenceLocalStorageMap() = default;
+SequenceLocalStorageMap::SequenceLocalStorageMap() {
+  DETACH_FROM_SEQUENCE(sequence_checker_);
+}
 
-SequenceLocalStorageMap::~SequenceLocalStorageMap() = default;
+SequenceLocalStorageMap::~SequenceLocalStorageMap() {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+}
 
 ScopedSetSequenceLocalStorageMapForCurrentThread::
     ScopedSetSequenceLocalStorageMapForCurrentThread(
@@ -48,6 +52,8 @@ SequenceLocalStorageMap& SequenceLocalStorageMap::GetForCurrentThread() {
 }
 
 void* SequenceLocalStorageMap::Get(int slot_id) {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+
   const auto it = sls_map_.find(slot_id);
   if (it == sls_map_.end())
     return nullptr;
@@ -57,6 +63,8 @@ void* SequenceLocalStorageMap::Get(int slot_id) {
 void SequenceLocalStorageMap::Set(
     int slot_id,
     SequenceLocalStorageMap::ValueDestructorPair value_destructor_pair) {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+
   auto it = sls_map_.find(slot_id);
 
   if (it == sls_map_.end())

base/threading/sequence_local_storage_map.h

@@ -8,6 +8,7 @@
 #include "base/base_export.h"
 #include "base/containers/flat_map.h"
 #include "base/macros.h"
+#include "base/sequence_checker.h"
 
 namespace base {
 namespace internal {
@@ -67,6 +68,8 @@ class BASE_EXPORT SequenceLocalStorageMap {
   // than other map implementations.
   base::flat_map<int, ValueDestructorPair> sls_map_;
 
+  SEQUENCE_CHECKER(sequence_checker_);
+
   DISALLOW_COPY_AND_ASSIGN(SequenceLocalStorageMap);
 };