2 files changed, 26 insertions, 17 deletions

diff --git a/webkit/plugins/ppapi/ppb_surface_3d_impl.cc b/webkit/plugins/ppapi/ppb_surface_3d_impl.cc
index f1c9eef..81d761f 100644
--- a/webkit/plugins/ppapi/ppb_surface_3d_impl.cc
+++ b/webkit/plugins/ppapi/ppb_surface_3d_impl.cc
@@ -23,7 +23,6 @@ PPB_Surface3D_Impl::PPB_Surface3D_Impl(PluginInstance* instance)
     : Resource(instance),
       bound_to_instance_(false),
       swap_initiated_(false),
-      swap_callback_(PP_BlockUntilComplete()),
       context_(NULL) {
 }
 
@@ -58,24 +57,29 @@ int32_t PPB_Surface3D_Impl::GetAttrib(int32_t attribute, int32_t* value) {
 }
 
 int32_t PPB_Surface3D_Impl::SwapBuffers(PP_CompletionCallback callback) {
-  if (!context_)
-    return PP_ERROR_FAILED;
-
-  if (swap_callback_.func) {
-    // Already a pending SwapBuffers that hasn't returned yet.
-    return PP_ERROR_INPROGRESS;
-  }
-
   if (!callback.func) {
     // Blocking SwapBuffers isn't supported (since we have to be on the main
     // thread).
     return PP_ERROR_BADARGUMENT;
   }
 
-  swap_callback_ = callback;
+  if (swap_callback_.get() && !swap_callback_->completed()) {
+    // Already a pending SwapBuffers that hasn't returned yet.
+    return PP_ERROR_INPROGRESS;
+  }
+
+  if (!context_)
+    return PP_ERROR_FAILED;
+
+  PP_Resource resource_id = GetReferenceNoAddRef();
+  CHECK(resource_id);
+  swap_callback_ = new TrackedCompletionCallback(
+      instance()->module()->GetCallbackTracker(), resource_id, callback);
   gpu::gles2::GLES2Implementation* impl = context_->gles2_impl();
   if (impl)
     context_->gles2_impl()->SwapBuffers();
+  // |SwapBuffers()| should not call us back synchronously, but double-check.
+  DCHECK(!swap_callback_->completed());
   return PP_OK_COMPLETIONPENDING;
 }
 
@@ -119,12 +123,14 @@ void PPB_Surface3D_Impl::ViewInitiatedPaint() {
 }
 
 void PPB_Surface3D_Impl::ViewFlushedPaint() {
-  if (swap_initiated_ && swap_callback_.func) {
+  if (swap_initiated_ && swap_callback_.get() && !swap_callback_->completed()) {
     // We must clear swap_callback_ before issuing the callback. It will be
     // common for the plugin to issue another SwapBuffers in response to the
     // callback, and we don't want to think that a callback is already pending.
     swap_initiated_ = false;
-    PP_RunAndClearCompletionCallback(&swap_callback_, PP_OK);
+    scoped_refptr<TrackedCompletionCallback> callback;
+    callback.swap(swap_callback_);
+    callback->Run(PP_OK);  // Will complete abortively if necessary.
   }
 }
 
@@ -136,11 +142,13 @@ void PPB_Surface3D_Impl::OnSwapBuffers() {
   if (bound_to_instance_) {
     instance()->CommitBackingTexture();
     swap_initiated_ = true;
-  } else if (swap_callback_.func) {
+  } else if (swap_callback_.get() && !swap_callback_->completed()) {
     // If we're off-screen, no need to trigger compositing so run the callback
     // immediately.
     swap_initiated_ = false;
-    PP_RunAndClearCompletionCallback(&swap_callback_, PP_OK);
+    scoped_refptr<TrackedCompletionCallback> callback;
+    callback.swap(swap_callback_);
+    callback->Run(PP_OK);  // Will complete abortively if necessary.
   }
 }
 
diff --git a/webkit/plugins/ppapi/ppb_surface_3d_impl.h b/webkit/plugins/ppapi/ppb_surface_3d_impl.h
index 7f4a7d1..24a66e5 100644
--- a/webkit/plugins/ppapi/ppb_surface_3d_impl.h
+++ b/webkit/plugins/ppapi/ppb_surface_3d_impl.h
@@ -1,12 +1,13 @@
-// Copyright (c) 2010 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #ifndef WEBKIT_PLUGINS_PPAPI_PPB_SURFACE_3D_IMPL_H_
 #define WEBKIT_PLUGINS_PPAPI_PPB_SURFACE_3D_IMPL_H_
 
-#include "base/callback.h"
+#include "base/memory/ref_counted.h"
 #include "ppapi/thunk/ppb_surface_3d_api.h"
+#include "webkit/plugins/ppapi/callbacks.h"
 #include "webkit/plugins/ppapi/plugin_delegate.h"
 #include "webkit/plugins/ppapi/resource.h"
 
@@ -69,7 +70,7 @@ class PPB_Surface3D_Impl : public Resource,
 
   // True when the page's SwapBuffers has been issued but not returned yet.
   bool swap_initiated_;
-  PP_CompletionCallback swap_callback_;
+  scoped_refptr<TrackedCompletionCallback> swap_callback_;
 
   // The context this surface is currently bound to.
   PPB_Context3D_Impl* context_;