1 files changed, 30 insertions, 25 deletions

diff --git a/net/websockets/websocket_job.cc b/net/websockets/websocket_job.cc
index 30c39a9..62a62b7 100644
--- a/net/websockets/websocket_job.cc
+++ b/net/websockets/websocket_job.cc
@@ -12,6 +12,30 @@
 #include "net/url_request/url_request_context.h"
 #include "net/websockets/websocket_throttle.h"
 
+namespace {
+
+class CompletionCallbackRunner
+    : public base::RefCountedThreadSafe<CompletionCallbackRunner> {
+ public:
+  explicit CompletionCallbackRunner(net::CompletionCallback* callback)
+      : callback_(callback) {
+    DCHECK(callback_);
+  }
+  void Run() {
+    callback_->Run(net::OK);
+  }
+ private:
+  friend class base::RefCountedThreadSafe<CompletionCallbackRunner>;
+
+  virtual ~CompletionCallbackRunner() {}
+
+  net::CompletionCallback* callback_;
+
+  DISALLOW_COPY_AND_ASSIGN(CompletionCallbackRunner);
+};
+
+}
+
 namespace net {
 
 // lower-case header names.
@@ -138,11 +162,6 @@ void WebSocketJob::DetachDelegate() {
   if (socket_)
     socket_->DetachDelegate();
   socket_ = NULL;
-  if (callback_) {
-    waiting_ = false;
-    callback_ = NULL;
-    Release();  // Balanced with OnStartOpenConnection().
-  }
 }
 
 int WebSocketJob::OnStartOpenConnection(
@@ -154,7 +173,6 @@ int WebSocketJob::OnStartOpenConnection(
   if (!waiting_)
     return OK;
   callback_ = callback;
-  AddRef();  // Balanced when callback_ becomes NULL.
   return ERR_IO_PENDING;
 }
 
@@ -191,11 +209,6 @@ void WebSocketJob::OnClose(SocketStream* socket) {
   SocketStream::Delegate* delegate = delegate_;
   delegate_ = NULL;
   socket_ = NULL;
-  if (callback_) {
-    waiting_ = false;
-    callback_ = NULL;
-    Release();  // Balanced with OnStartOpenConnection().
-  }
   if (delegate)
     delegate->OnClose(socket);
 }
@@ -424,24 +437,16 @@ bool WebSocketJob::IsWaiting() const {
 }
 
 void WebSocketJob::Wakeup() {
-  if (!waiting_)
-    return;
   waiting_ = false;
   DCHECK(callback_);
+  // We wrap |callback_| to keep this alive while this is released.
+  scoped_refptr<CompletionCallbackRunner> runner =
+      new CompletionCallbackRunner(callback_);
+  callback_ = NULL;
   MessageLoopForIO::current()->PostTask(
       FROM_HERE,
-      NewRunnableMethod(this,
-                        &WebSocketJob::DoCallback));
-}
-
-void WebSocketJob::DoCallback() {
-  // |callback_| may be NULL if OnClose() or DetachDelegate() was called.
-  net::CompletionCallback* callback = callback_;
-  callback_ = NULL;
-  if (callback) {
-    callback->Run(net::OK);
-    Release();  // Balanced with OnStartOpenConnection().
-  }
+      NewRunnableMethod(runner.get(),
+                        &CompletionCallbackRunner::Run));
 }
 
 }  // namespace net