Fix ThreadChecker bug in WeakNSObject

When a WeakNSObject was copied on a different thread from where it was
created, the state of the thread checker was lost.
As a consequence it was not possible to send a WeakNSObject on another
thread, copy it there, and send the copy back: a ThreadChecker error
was raised.

This CL adds more thread checks in the WeakContainer, which matches
what WeakPtr also does, and keeps most of the checks in the main
class (which is stricter than WeakPtr).

Review URL: https://codereview.chromium.org/853503002

Cr-Commit-Position: refs/heads/master@{#311865}

2 files changed, 85 insertions, 25 deletions

diff --git a/base/ios/weak_nsobject.h b/base/ios/weak_nsobject.h
index 46aecb5..a1984bb 100644
--- a/base/ios/weak_nsobject.h
+++ b/base/ios/weak_nsobject.h
@@ -36,17 +36,28 @@
 // NSObject, this relationship is maintained via the ObjectiveC associated
 // object API, indirectly via an ObjectiveC CRBWeakNSProtocolSentinel class.
 //
-// The implementation assumes that the tracked object will be released on the
-// same thread that the WeakNSObject is created on.
-//
+// Threading restrictions:
+// - Several WeakNSObject pointing to the same underlying object must all be
+//   created and dereferenced on the same thread;
+// - thread safety is enforced by the implementation, except in two cases:
+//   (1) it is allowed to copy a WeakNSObject on a different thread. However,
+//       that copy must return to the original thread before being dereferenced,
+//   (2) it is allowed to destroy a WeakNSObject on any thread;
+// - the implementation assumes that the tracked object will be released on the
+//   same thread that the WeakNSObject is created on.
 namespace base {
 
 // WeakContainer keeps a weak pointer to an object and clears it when it
 // receives nullify() from the object's sentinel.
 class WeakContainer : public base::RefCountedThreadSafe<WeakContainer> {
  public:
-  WeakContainer(id object) : object_(object) {}
-  id object() { return object_; }
+  explicit WeakContainer(id object) : object_(object) {}
+
+  id object() {
+    DCHECK(checker_.CalledOnValidThread());
+    return object_;
+  }
+
   void nullify() {
     DCHECK(checker_.CalledOnValidThread());
     object_ = nil;
@@ -74,53 +85,63 @@ namespace base {
 
 // Base class for all WeakNSObject derivatives.
 template <typename NST>
-class WeakNSProtocol : public base::NonThreadSafe {
+class WeakNSProtocol {
  public:
   explicit WeakNSProtocol(NST object = nil) {
     container_ = [CRBWeakNSProtocolSentinel containerForObject:object];
   }
 
   WeakNSProtocol(const WeakNSProtocol<NST>& that) {
+    // A WeakNSProtocol object can be copied on one thread and used on
+    // another.
+    checker_.DetachFromThread();
     container_ = that.container_;
   }
 
   ~WeakNSProtocol() {
-    // A WeakNSProtocol object can be allocated on one thread and released on
+    // A WeakNSProtocol object can be used on one thread and released on
     // another. This is not the case for the contained object.
-    DetachFromThread();
+    checker_.DetachFromThread();
   }
 
   void reset(NST object = nil) {
-    DCHECK(CalledOnValidThread());
+    DCHECK(checker_.CalledOnValidThread());
     container_ = [CRBWeakNSProtocolSentinel containerForObject:object];
   }
 
   NST get() const {
-    DCHECK(CalledOnValidThread());
+    DCHECK(checker_.CalledOnValidThread());
     if (!container_.get())
       return nil;
     return container_->object();
   }
 
   WeakNSProtocol& operator=(const WeakNSProtocol<NST>& that) {
-    DCHECK(CalledOnValidThread());
+    DCHECK(checker_.CalledOnValidThread());
     container_ = that.container_;
     return *this;
   }
 
   bool operator==(NST that) const {
-    DCHECK(CalledOnValidThread());
+    DCHECK(checker_.CalledOnValidThread());
     return get() == that;
   }
 
-  bool operator!=(NST that) const { return get() != that; }
+  bool operator!=(NST that) const {
+    DCHECK(checker_.CalledOnValidThread());
+    return get() != that;
+  }
 
-  operator NST() const { return get(); }
+  operator NST() const {
+    DCHECK(checker_.CalledOnValidThread());
+    return get();
+  }
 
  private:
   // Refecounted reference to the container tracking the ObjectiveC object this
   // class encapsulates.
   scoped_refptr<base::WeakContainer> container_;
+  base::ThreadChecker checker_;
 };
 
 // Free functions
diff --git a/base/ios/weak_nsobject_unittest.mm b/base/ios/weak_nsobject_unittest.mm
index 9758aed..325dcd2 100644
--- a/base/ios/weak_nsobject_unittest.mm
+++ b/base/ios/weak_nsobject_unittest.mm
@@ -3,16 +3,19 @@
 // found in the LICENSE file.
 
 #include "base/basictypes.h"
+#include "base/bind.h"
 #include "base/ios/weak_nsobject.h"
 #include "base/mac/scoped_nsobject.h"
+#include "base/message_loop/message_loop.h"
+#include "base/single_thread_task_runner.h"
+#include "base/threading/thread.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
-using base::WeakNSObject;
-
+namespace base {
 namespace {
 
 TEST(WeakNSObjectTest, WeakNSObject) {
-  base::scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
+  scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
   WeakNSObject<NSObject> w1(p1);
   EXPECT_TRUE(w1);
   p1.reset();
@@ -20,7 +23,7 @@ TEST(WeakNSObjectTest, WeakNSObject) {
 }
 
 TEST(WeakNSObjectTest, MultipleWeakNSObject) {
-  base::scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
+  scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
   WeakNSObject<NSObject> w1(p1);
   WeakNSObject<NSObject> w2(w1);
   EXPECT_TRUE(w1);
@@ -32,7 +35,7 @@ TEST(WeakNSObjectTest, MultipleWeakNSObject) {
 }
 
 TEST(WeakNSObjectTest, WeakNSObjectDies) {
-  base::scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
+  scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
   {
     WeakNSObject<NSObject> w1(p1);
     EXPECT_TRUE(w1);
@@ -40,7 +43,7 @@ TEST(WeakNSObjectTest, WeakNSObjectDies) {
 }
 
 TEST(WeakNSObjectTest, WeakNSObjectReset) {
-  base::scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
+  scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
   WeakNSObject<NSObject> w1(p1);
   EXPECT_TRUE(w1);
   w1.reset();
@@ -50,8 +53,8 @@ TEST(WeakNSObjectTest, WeakNSObjectReset) {
 }
 
 TEST(WeakNSObjectTest, WeakNSObjectResetWithObject) {
-  base::scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
-  base::scoped_nsobject<NSObject> p2([[NSObject alloc] init]);
+  scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
+  scoped_nsobject<NSObject> p2([[NSObject alloc] init]);
   WeakNSObject<NSObject> w1(p1);
   EXPECT_TRUE(w1);
   w1.reset(p2);
@@ -61,7 +64,7 @@ TEST(WeakNSObjectTest, WeakNSObjectResetWithObject) {
 }
 
 TEST(WeakNSObjectTest, WeakNSObjectEmpty) {
-  base::scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
+  scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
   WeakNSObject<NSObject> w1;
   EXPECT_FALSE(w1);
   w1.reset(p1);
@@ -71,7 +74,7 @@ TEST(WeakNSObjectTest, WeakNSObjectEmpty) {
 }
 
 TEST(WeakNSObjectTest, WeakNSObjectCopy) {
-  base::scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
+  scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
   WeakNSObject<NSObject> w1(p1);
   WeakNSObject<NSObject> w2(w1);
   EXPECT_TRUE(w1);
@@ -82,7 +85,7 @@ TEST(WeakNSObjectTest, WeakNSObjectCopy) {
 }
 
 TEST(WeakNSObjectTest, WeakNSObjectAssignment) {
-  base::scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
+  scoped_nsobject<NSObject> p1([[NSObject alloc] init]);
   WeakNSObject<NSObject> w1(p1);
   WeakNSObject<NSObject> w2;
   EXPECT_FALSE(w2);
@@ -94,4 +97,40 @@ TEST(WeakNSObjectTest, WeakNSObjectAssignment) {
   EXPECT_FALSE(w2);
 }
 
+// Touches |weak_data| by increasing its length by 1. Used to check that the
+// weak object can be dereferenced.
+void TouchWeakData(const WeakNSObject<NSMutableData>& weak_data) {
+  if (!weak_data)
+    return;
+  [weak_data increaseLengthBy:1];
+}
+
+// Makes a copy of |weak_object| on the current thread and posts a task to touch
+// the weak object on its original thread.
+void CopyWeakNSObjectAndPost(const WeakNSObject<NSMutableData>& weak_object,
+                             scoped_refptr<SingleThreadTaskRunner> runner) {
+  WeakNSObject<NSMutableData> weak_copy(weak_object);
+  runner->PostTask(FROM_HERE, Bind(&TouchWeakData, weak_copy));
+}
+
+// Tests that the weak object can be copied on a different thread.
+TEST(WeakNSObjectTest, WeakNSObjectCopyOnOtherThread) {
+  MessageLoop loop;
+  Thread other_thread("WeakNSObjectCopyOnOtherThread");
+  other_thread.Start();
+
+  scoped_nsobject<NSMutableData> data([[NSMutableData alloc] init]);
+  WeakNSObject<NSMutableData> weak(data);
+
+  scoped_refptr<SingleThreadTaskRunner> runner = loop.task_runner();
+  other_thread.task_runner()->PostTask(
+      FROM_HERE, Bind(&CopyWeakNSObjectAndPost, weak, runner));
+  other_thread.Stop();
+  loop.RunUntilIdle();
+
+  // Check that TouchWeakData was called.
+  EXPECT_EQ(1u, [data length]);
+}
+
 }  // namespace
+}  // namespace base