PPAPI: Make TrackedCallback more threadsafe

ppapi/proxy/tracked_callback_unittest.cc

Index: ppapi/proxy/tracked_callback_unittest.cc
diff --git a/ppapi/shared_impl/tracked_callback_unittest.cc b/ppapi/proxy/tracked_callback_unittest.cc
similarity index 40%
rename from ppapi/shared_impl/tracked_callback_unittest.cc
rename to ppapi/proxy/tracked_callback_unittest.cc
index 12ad4235bdb66341e0c82a7b922e4864d38caf8b..15dea53c8c1fb52563e51cc4d7866f183326956b 100644
--- a/ppapi/shared_impl/tracked_callback_unittest.cc
+++ b/ppapi/proxy/tracked_callback_unittest.cc
@@ -5,39 +5,81 @@
 #include "base/bind.h"
 #include "base/memory/ref_counted.h"
 #include "base/message_loop/message_loop.h"
+#include "base/synchronization/waitable_event.h"
+#include "base/threading/simple_thread.h"
 #include "ppapi/c/pp_completion_callback.h"
 #include "ppapi/c/pp_errors.h"
+#include "ppapi/proxy/ppapi_proxy_test.h"
+#include "ppapi/proxy/ppb_message_loop_proxy.h"
 #include "ppapi/shared_impl/callback_tracker.h"
 #include "ppapi/shared_impl/proxy_lock.h"
 #include "ppapi/shared_impl/resource.h"
 #include "ppapi/shared_impl/resource_tracker.h"
+#include "ppapi/shared_impl/scoped_pp_resource.h"
 #include "ppapi/shared_impl/test_globals.h"
 #include "ppapi/shared_impl/tracked_callback.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
+// Note, this file tests TrackedCallback which lives in ppapi/shared_impl.
+// Unfortunately, we need the test to live in ppapi/proxy so that it can use
+// the thread support there.
 namespace ppapi {
+namespace proxy {
 
 namespace {
 
-class TrackedCallbackTest : public testing::Test {
+class CallbackThread : public base::SimpleThread {
  public:
-  TrackedCallbackTest() : pp_instance_(1234) {}
-
-  PP_Instance pp_instance() const { return pp_instance_; }
-
-  virtual void SetUp() override {
-    ProxyAutoLock lock;
-    globals_.GetResourceTracker()->DidCreateInstance(pp_instance_);
+  CallbackThread(PP_Instance instance)

[bbudge, 2015/03/05 01:29:35]

nit: explicit

[dmichael (off chromium), 2015/03/18 15:51:48]

Done.

+      : SimpleThread("CallbackThread"), instance_(instance) {}
+  ~CallbackThread() override {
+    ProxyAutoLock acquire;
+    loop_ = 0;

[bbudge, 2015/03/05 01:29:35]

Could we use nullptr here?

[dmichael (off chromium), 2015/03/18 15:51:48]

That's better, not sure what I was thinking.

+  };
+
+  // base::SimpleThread overrides.
+  void Start() override {
+    {
+      ProxyAutoLock acquire;
+      // Create the loop here, after PpapiGlobals has been created.
+      loop_ = new MessageLoopResource(instance_);
+    }
+    base::SimpleThread::Start();
   }
-  virtual void TearDown() override {
-    ProxyAutoLock lock;
-    globals_.GetResourceTracker()->DidDeleteInstance(pp_instance_);
+  void Run() override {
+    ProxyAutoLock acquire;
+    loop_->AttachToCurrentThread();
+    // Note, run releases the lock to run events.
+    loop_->Run();
   }
 
+  MessageLoopResource* loop() { return loop_.get(); }

[bbudge, 2015/03/05 01:29:35]

nit: s/loop/message_loop
It confused me first time I read the file.

[dmichael (off chromium), 2015/03/18 15:51:48]

Done.

+
+ private:
+  PP_Instance instance_;
+  scoped_refptr<MessageLoopResource> loop_;
+};
+
+class TrackedCallbackTest : public PluginProxyTest {
+ public:
+  TrackedCallbackTest() : thread_(pp_instance()) {}
+  CallbackThread& thread() { return thread_; }
+
  private:
-  base::MessageLoop message_loop_;
-  TestGlobals globals_;
-  PP_Instance pp_instance_;
+  // PluginProxyTest overrides.
+  void SetUp() override {
+    PluginProxyTest::SetUp();
+    thread_.Start();
+  }
+  void TearDown() override {
+    {
+      ProxyAutoLock acquire;
+      thread_.loop()->PostQuit(PP_TRUE);
+    }
+    thread_.Join();

[bbudge, 2015/03/05 01:29:35]

Would it be cleaner to pull the above code into a CallbackThread::Stop() method,
to go with the Start() method?

[dmichael (off chromium), 2015/03/18 15:51:48]

How about I override Join()? That seems like a nice parallel.

+    PluginProxyTest::TearDown();
+  }
+  CallbackThread thread_;
 };
 
 // All valid results (PP_OK, PP_ERROR_...) are nonpositive.
@@ -45,33 +87,54 @@ const int32_t kInitializedResultValue = 1;
 const int32_t kOverrideResultValue = 2;
 
 struct CallbackRunInfo {
-  CallbackRunInfo()
-      : run_count(0),
-        result(kInitializedResultValue),
-        completion_task_run_count(0),
-        completion_task_result(kInitializedResultValue) {}
-  unsigned run_count;
-  int32_t result;
-  unsigned completion_task_run_count;
-  int32_t completion_task_result;
+  explicit CallbackRunInfo(base::ThreadChecker* thread_checker)

[bbudge, 2015/03/05 01:29:35]

I don't see any calls of this. It looks like you have a default ctor that is
getting used.

[dmichael (off chromium), 2015/03/18 15:51:48]

See the initializer lists for CallbackShutdownTest and CallbackMockResource. You
don't get a default constructor for free if you define any other constructors.

[bbudge, 2015/03/19 00:21:10]

Missed those, you're right.

+      : run_count_(0),
+        result_(kInitializedResultValue),
+        completion_task_run_count_(0),
+        completion_task_result_(kInitializedResultValue),
+        thread_checker_(thread_checker),
+        callback_did_run_event_(true, false) {}
+  void CallbackDidRun(int32_t result) {
+    CHECK(thread_checker_->CalledOnValidThread());
+    if (!run_count_)
+      result_ = result;
+    ++run_count_;
+    callback_did_run_event_.Signal();
+  }
+  void CompletionTaskDidRun(int32_t result) {
+    CHECK(thread_checker_->CalledOnValidThread());
+    if (!completion_task_run_count_)
+      completion_task_result_ = result;
+    ++completion_task_run_count_;
+  }
+  void WaitUntilCompleted() { callback_did_run_event_.Wait(); }
+  const unsigned run_count() { return run_count_; }
+  int32_t result() { return result_; }
+  unsigned completion_task_run_count() { return completion_task_run_count_; }
+  int32_t completion_task_result() { return completion_task_result_; }
+ private:
+  unsigned run_count_;
+  int32_t result_;
+  unsigned completion_task_run_count_;
+  int32_t completion_task_result_;
+  // Weak; owned by the creator of CallbackRunInfo.
+  base::ThreadChecker* thread_checker_;
+
+  base::WaitableEvent callback_did_run_event_;
 };
 
 void TestCallback(void* user_data, int32_t result) {
-  CallbackRunInfo* info = reinterpret_cast<CallbackRunInfo*>(user_data);
-  info->run_count++;
-  if (info->run_count == 1)
-    info->result = result;
+  CallbackRunInfo* info = static_cast<CallbackRunInfo*>(user_data);
+  info->CallbackDidRun(result);
 }
 
-}  // namespace
-
 // CallbackShutdownTest --------------------------------------------------------
 
-namespace {
-
 class CallbackShutdownTest : public TrackedCallbackTest {
  public:
-  CallbackShutdownTest() {}
+  CallbackShutdownTest() : info_did_run_(&thread_checker_),
+                           info_did_abort_(&thread_checker_),
+                           info_didnt_run_(&thread_checker_) {}
 
   // Cases:
   // (1) A callback which is run (so shouldn't be aborted on shutdown).
@@ -82,6 +145,7 @@ class CallbackShutdownTest : public TrackedCallbackTest {
   CallbackRunInfo& info_didnt_run() { return info_didnt_run_; }  // (3)
 
  private:
+  base::ThreadChecker thread_checker_;
   CallbackRunInfo info_did_run_;
   CallbackRunInfo info_did_abort_;
   CallbackRunInfo info_didnt_run_;
@@ -92,46 +156,48 @@ class CallbackShutdownTest : public TrackedCallbackTest {
 // Tests that callbacks are properly aborted on module shutdown.
 TEST_F(CallbackShutdownTest, AbortOnShutdown) {
   ProxyAutoLock lock;
-  scoped_refptr<Resource> resource(new Resource(OBJECT_IS_IMPL, pp_instance()));
+  scoped_refptr<Resource> resource(
+      new Resource(OBJECT_IS_PROXY, pp_instance()));
 
   // Set up case (1) (see above).
-  EXPECT_EQ(0U, info_did_run().run_count);
+  EXPECT_EQ(0U, info_did_run().run_count());
+  // TODO(dmichael): Test this on a background thread?
   scoped_refptr<TrackedCallback> callback_did_run = new TrackedCallback(
       resource.get(),
       PP_MakeCompletionCallback(&TestCallback, &info_did_run()));
-  EXPECT_EQ(0U, info_did_run().run_count);
+  EXPECT_EQ(0U, info_did_run().run_count());
   callback_did_run->Run(PP_OK);
-  EXPECT_EQ(1U, info_did_run().run_count);
-  EXPECT_EQ(PP_OK, info_did_run().result);
+  EXPECT_EQ(1U, info_did_run().run_count());
+  EXPECT_EQ(PP_OK, info_did_run().result());
 
   // Set up case (2).
-  EXPECT_EQ(0U, info_did_abort().run_count);
+  EXPECT_EQ(0U, info_did_abort().run_count());
   scoped_refptr<TrackedCallback> callback_did_abort = new TrackedCallback(
       resource.get(),
       PP_MakeCompletionCallback(&TestCallback, &info_did_abort()));
-  EXPECT_EQ(0U, info_did_abort().run_count);
+  EXPECT_EQ(0U, info_did_abort().run_count());
   callback_did_abort->Abort();
-  EXPECT_EQ(1U, info_did_abort().run_count);
-  EXPECT_EQ(PP_ERROR_ABORTED, info_did_abort().result);
+  EXPECT_EQ(1U, info_did_abort().run_count());
+  EXPECT_EQ(PP_ERROR_ABORTED, info_did_abort().result());
 
   // Set up case (3).
-  EXPECT_EQ(0U, info_didnt_run().run_count);
+  EXPECT_EQ(0U, info_didnt_run().run_count());
   scoped_refptr<TrackedCallback> callback_didnt_run = new TrackedCallback(
       resource.get(),
       PP_MakeCompletionCallback(&TestCallback, &info_didnt_run()));
-  EXPECT_EQ(0U, info_didnt_run().run_count);
+  EXPECT_EQ(0U, info_didnt_run().run_count());
 
-  PpapiGlobals::Get()->GetCallbackTrackerForInstance(pp_instance())->AbortAll();
+  GetGlobals()->GetCallbackTrackerForInstance(pp_instance())->AbortAll();
 
   // Check case (1).
-  EXPECT_EQ(1U, info_did_run().run_count);
+  EXPECT_EQ(1U, info_did_run().run_count());
 
   // Check case (2).
-  EXPECT_EQ(1U, info_did_abort().run_count);
+  EXPECT_EQ(1U, info_did_abort().run_count());
 
   // Check case (3).
-  EXPECT_EQ(1U, info_didnt_run().run_count);
-  EXPECT_EQ(PP_ERROR_ABORTED, info_didnt_run().result);
+  EXPECT_EQ(1U, info_didnt_run().run_count());
+  EXPECT_EQ(PP_ERROR_ABORTED, info_didnt_run().result());
 }
 
 // CallbackResourceTest --------------------------------------------------------
@@ -145,18 +211,124 @@ class CallbackResourceTest : public TrackedCallbackTest {
 
 class CallbackMockResource : public Resource {
  public:
-  CallbackMockResource(PP_Instance instance)
-      : Resource(OBJECT_IS_IMPL, instance) {}
+  static scoped_refptr<CallbackMockResource> Create(PP_Instance instance) {
+    ProxyAutoLock acquire;
+    return scoped_refptr<CallbackMockResource>(
+        new CallbackMockResource(instance));
+  }
   ~CallbackMockResource() {}
 
-  PP_Resource SetupForTest() {
-    PP_Resource resource_id = GetReference();
-    EXPECT_NE(0, resource_id);
+  // Take a reference to this resource, which will add it to the tracker.
+  void TakeRef() {
+    ProxyAutoLock acquire;
+    ScopedPPResource temp_resource(ScopedPPResource::PassRef(), GetReference());
+    EXPECT_NE(0, temp_resource.get());
+    reference_holder_ = temp_resource;
+  }
+  // Release it, removing it from the tracker.
+  void ReleaseRef() {
+    ProxyAutoLock acquire;
+    reference_holder_ = 0;
+  }
+
+  // Create the test callbacks on a background thread, so that we can verify
+  // they are run on the same thread where they were created.
+  void CreateCallbacksOnLoop(MessageLoopResource* loop_resource) {
+    ProxyAutoLock acquire;
+    // |thread_checker_| will bind to the background thread.
+    thread_checker_.DetachFromThread();
+    loop_resource->message_loop_proxy()->PostTask(FROM_HERE,
+        RunWhileLocked(
+            base::Bind(&CallbackMockResource::CreateCallbacks, this)));
+  }
+
+  int32_t CompletionTask(CallbackRunInfo* info, int32_t result) {
+    // The completion task must run on the thread where the callback was
+    // created, and must hold the proxy lock.
+    CHECK(thread_checker_.CalledOnValidThread());
+    ProxyLock::AssertAcquired();
+
+    // We should run before the callback.
+    CHECK_EQ(0U, info->run_count());
+    info->CompletionTaskDidRun(result);
+    return kOverrideResultValue;
+  }
+
+  void CheckInitialState() {
+    callbacks_created_event_.Wait();
+    EXPECT_EQ(0U, info_did_run_.run_count());
+    EXPECT_EQ(0U, info_did_run_.completion_task_run_count());
+
+    EXPECT_EQ(0U, info_did_run_with_completion_task_.run_count());
+    EXPECT_EQ(0U,
+              info_did_run_with_completion_task_.completion_task_run_count());
+
+    EXPECT_EQ(0U, info_did_abort_.run_count());
+    EXPECT_EQ(0U, info_did_abort_.completion_task_run_count());
+
+    EXPECT_EQ(0U, info_didnt_run_.run_count());
+    EXPECT_EQ(0U, info_didnt_run_.completion_task_run_count());
+  }
+
+  void RunCallbacks() {
+    callback_did_run_->Run(PP_OK);
+    callback_did_run_with_completion_task_->Run(PP_OK);
+    callback_did_abort_->Abort();
+    info_did_run_.WaitUntilCompleted();
+    info_did_run_with_completion_task_.WaitUntilCompleted();
+    info_did_abort_.WaitUntilCompleted();
+  }
+
+  void CheckIntermediateState() {
+    EXPECT_EQ(1U, info_did_run_.run_count());
+    EXPECT_EQ(PP_OK, info_did_run_.result());
+    EXPECT_EQ(0U, info_did_run_.completion_task_run_count());
+
+    EXPECT_EQ(1U, info_did_run_with_completion_task_.run_count());
+    // completion task should override the result.
+    EXPECT_EQ(kOverrideResultValue,
+              info_did_run_with_completion_task_.result());
+    EXPECT_EQ(1U,
+              info_did_run_with_completion_task_.completion_task_run_count());
+    EXPECT_EQ(PP_OK,
+              info_did_run_with_completion_task_.completion_task_result());
+
+    EXPECT_EQ(1U, info_did_abort_.run_count());
+    // completion task shouldn't override an abort.
+    EXPECT_EQ(PP_ERROR_ABORTED, info_did_abort_.result());
+    EXPECT_EQ(1U, info_did_abort_.completion_task_run_count());
+    EXPECT_EQ(PP_ERROR_ABORTED, info_did_abort_.completion_task_result());
+
+    EXPECT_EQ(0U, info_didnt_run_.completion_task_run_count());
+    EXPECT_EQ(0U, info_didnt_run_.run_count());
+  }
+
+  void CheckFinalState() {
+    info_didnt_run_.WaitUntilCompleted();
+    EXPECT_EQ(1U, info_did_run_.run_count());
+    EXPECT_EQ(PP_OK, info_did_run_.result());
+    EXPECT_EQ(1U, info_did_abort_.run_count());
+    EXPECT_EQ(PP_ERROR_ABORTED, info_did_abort_.result());
+    EXPECT_EQ(1U, info_didnt_run_.run_count());
+    EXPECT_EQ(PP_ERROR_ABORTED, info_didnt_run_.result());
+  }
+
+ private:
+  explicit CallbackMockResource(PP_Instance instance)
+      : Resource(OBJECT_IS_IMPL, instance),
+        info_did_run_(&thread_checker_),
+        info_did_run_with_completion_task_(&thread_checker_),
+        info_did_abort_(&thread_checker_),
+        info_didnt_run_(&thread_checker_),
+        callbacks_created_event_(true, false) {}
+  void CreateCallbacks() {
+    // Bind thread_checker_ to the thread where we create the callbacks.
+    // Later, when the callback runs, it will check that it was invoked on this
+    // same thread.
+    CHECK(thread_checker_.CalledOnValidThread());
 
     callback_did_run_ = new TrackedCallback(
         this, PP_MakeCompletionCallback(&TestCallback, &info_did_run_));
-    EXPECT_EQ(0U, info_did_run_.run_count);
-    EXPECT_EQ(0U, info_did_run_.completion_task_run_count);
 
     // In order to test that the completion task can override the callback
     // result, we need to test callbacks with and without a completion task.
@@ -168,70 +340,23 @@ class CallbackMockResource : public Resource {
         Bind(&CallbackMockResource::CompletionTask,
              this,
              &info_did_run_with_completion_task_));
-    EXPECT_EQ(0U, info_did_run_with_completion_task_.run_count);
-    EXPECT_EQ(0U, info_did_run_with_completion_task_.completion_task_run_count);
 
     callback_did_abort_ = new TrackedCallback(
         this, PP_MakeCompletionCallback(&TestCallback, &info_did_abort_));
     callback_did_abort_->set_completion_task(
         Bind(&CallbackMockResource::CompletionTask, this, &info_did_abort_));
-    EXPECT_EQ(0U, info_did_abort_.run_count);
-    EXPECT_EQ(0U, info_did_abort_.completion_task_run_count);
 
     callback_didnt_run_ = new TrackedCallback(
         this, PP_MakeCompletionCallback(&TestCallback, &info_didnt_run_));
     callback_didnt_run_->set_completion_task(
         Bind(&CallbackMockResource::CompletionTask, this, &info_didnt_run_));
-    EXPECT_EQ(0U, info_didnt_run_.run_count);
-    EXPECT_EQ(0U, info_didnt_run_.completion_task_run_count);
-
-    callback_did_run_->Run(PP_OK);
-    callback_did_run_with_completion_task_->Run(PP_OK);
-    callback_did_abort_->Abort();
-
-    CheckIntermediateState();
-
-    return resource_id;
-  }
-
-  int32_t CompletionTask(CallbackRunInfo* info, int32_t result) {
-    // We should run before the callback.
-    EXPECT_EQ(0U, info->run_count);
-    info->completion_task_run_count++;
-    if (info->completion_task_run_count == 1)
-      info->completion_task_result = result;
-    return kOverrideResultValue;
-  }
-
-  void CheckIntermediateState() {
-    EXPECT_EQ(1U, info_did_run_.run_count);
-    EXPECT_EQ(PP_OK, info_did_run_.result);
-    EXPECT_EQ(0U, info_did_run_.completion_task_run_count);
-
-    EXPECT_EQ(1U, info_did_run_with_completion_task_.run_count);
-    // completion task should override the result.
-    EXPECT_EQ(kOverrideResultValue, info_did_run_with_completion_task_.result);
-    EXPECT_EQ(1U, info_did_run_with_completion_task_.completion_task_run_count);
-    EXPECT_EQ(PP_OK, info_did_run_with_completion_task_.completion_task_result);
-
-    EXPECT_EQ(1U, info_did_abort_.run_count);
-    // completion task shouldn't override an abort.
-    EXPECT_EQ(PP_ERROR_ABORTED, info_did_abort_.result);
-    EXPECT_EQ(1U, info_did_abort_.completion_task_run_count);
-    EXPECT_EQ(PP_ERROR_ABORTED, info_did_abort_.completion_task_result);
 
-    EXPECT_EQ(0U, info_didnt_run_.completion_task_run_count);
-    EXPECT_EQ(0U, info_didnt_run_.run_count);
+    callbacks_created_event_.Signal();
   }
 
-  void CheckFinalState() {
-    EXPECT_EQ(1U, info_did_run_.run_count);
-    EXPECT_EQ(PP_OK, info_did_run_.result);
-    EXPECT_EQ(1U, info_did_abort_.run_count);
-    EXPECT_EQ(PP_ERROR_ABORTED, info_did_abort_.result);
-    EXPECT_EQ(1U, info_didnt_run_.run_count);
-    EXPECT_EQ(PP_ERROR_ABORTED, info_didnt_run_.result);
-  }
+  // Used to verify that the callback runs on the same thread where it is
+  // created.
+  base::ThreadChecker thread_checker_;
 
   scoped_refptr<TrackedCallback> callback_did_run_;
   CallbackRunInfo info_did_run_;
@@ -244,98 +369,78 @@ class CallbackMockResource : public Resource {
 
   scoped_refptr<TrackedCallback> callback_didnt_run_;
   CallbackRunInfo info_didnt_run_;
+
+  base::WaitableEvent callbacks_created_event_;
+
+  ScopedPPResource reference_holder_;
 };
 
 }  // namespace
 
 // Test that callbacks get aborted on the last resource unref.
 TEST_F(CallbackResourceTest, AbortOnNoRef) {
-  ProxyAutoLock lock;
-  ResourceTracker* resource_tracker = PpapiGlobals::Get()->GetResourceTracker();
-
   // Test several things: Unref-ing a resource (to zero refs) with callbacks
   // which (1) have been run, (2) have been aborted, (3) haven't been completed.
   // Check that the uncompleted one gets aborted, and that the others don't get
   // called again.
   scoped_refptr<CallbackMockResource> resource_1(
-      new CallbackMockResource(pp_instance()));
-  PP_Resource resource_1_id = resource_1->SetupForTest();
+      CallbackMockResource::Create(pp_instance()));
+  resource_1->CreateCallbacksOnLoop(thread().loop());
+  resource_1->CheckInitialState();
+  resource_1->RunCallbacks();
+  resource_1->TakeRef();
+  resource_1->CheckIntermediateState();
 
   // Also do the same for a second resource, and make sure that unref-ing the
   // first resource doesn't much up the second resource.
   scoped_refptr<CallbackMockResource> resource_2(
-      new CallbackMockResource(pp_instance()));
-  PP_Resource resource_2_id = resource_2->SetupForTest();
+      CallbackMockResource::Create(pp_instance()));
+  resource_2->CreateCallbacksOnLoop(thread().loop());
+  resource_2->CheckInitialState();
+  resource_2->RunCallbacks();
+  resource_2->TakeRef();
+  resource_2->CheckIntermediateState();
 
   // Double-check that resource #1 is still okay.
   resource_1->CheckIntermediateState();
 
   // Kill resource #1, spin the message loop to run posted calls, and check that
   // things are in the expected states.
-  resource_tracker->ReleaseResource(resource_1_id);
-  {
-    ProxyAutoUnlock unlock;
-    base::MessageLoop::current()->RunUntilIdle();
-  }
+  resource_1->ReleaseRef();
+
   resource_1->CheckFinalState();
   resource_2->CheckIntermediateState();
 
   // Kill resource #2.
-  resource_tracker->ReleaseResource(resource_2_id);
-  {
-    ProxyAutoUnlock unlock;
-    base::MessageLoop::current()->RunUntilIdle();
-  }
+  resource_2->ReleaseRef();
+
   resource_1->CheckFinalState();
   resource_2->CheckFinalState();
-
-  // This shouldn't be needed, but make sure there are no stranded tasks.
-  {
-    ProxyAutoUnlock unlock;
-    base::MessageLoop::current()->RunUntilIdle();
-  }
 }
 
 // Test that "resurrecting" a resource (getting a new ID for a |Resource|)
 // doesn't resurrect callbacks.
 TEST_F(CallbackResourceTest, Resurrection) {
-  ProxyAutoLock lock;
-  ResourceTracker* resource_tracker = PpapiGlobals::Get()->GetResourceTracker();
-
   scoped_refptr<CallbackMockResource> resource(
-      new CallbackMockResource(pp_instance()));
-  PP_Resource resource_id = resource->SetupForTest();
-
-  // Unref it, spin the message loop to run posted calls, and check that things
-  // are in the expected states.
-  resource_tracker->ReleaseResource(resource_id);
-  {
-    ProxyAutoUnlock unlock;
-    base::MessageLoop::current()->RunUntilIdle();
-  }
+      CallbackMockResource::Create(pp_instance()));
+  resource->CreateCallbacksOnLoop(thread().loop());
+  resource->CheckInitialState();
+  resource->RunCallbacks();
+  resource->TakeRef();
+  resource->CheckIntermediateState();
+
+  // Unref it and check that things are in the expected states.
+  resource->ReleaseRef();
   resource->CheckFinalState();
 
   // "Resurrect" it and check that the callbacks are still dead.
-  PP_Resource new_resource_id = resource->GetReference();
-  {
-    ProxyAutoUnlock unlock;
-    base::MessageLoop::current()->RunUntilIdle();
-  }
+  resource->TakeRef();
   resource->CheckFinalState();
 
   // Unref it again and do the same.
-  resource_tracker->ReleaseResource(new_resource_id);
-  {
-    ProxyAutoUnlock unlock;
-    base::MessageLoop::current()->RunUntilIdle();
-  }
+  resource->ReleaseRef();
   resource->CheckFinalState();
-
-  // This shouldn't be needed, but make sure there are no stranded tasks.
-  {
-    ProxyAutoUnlock unlock;
-    base::MessageLoop::current()->RunUntilIdle();
-  }
 }
 
+}  // namespace proxy
 }  // namespace ppapi