3 files changed, 38 insertions, 37 deletions
diff --git a/net/websockets/websocket_job.cc b/net/websockets/websocket_job.cc
index 62a62b7..30c39a9 100644
--- a/net/websockets/websocket_job.cc
+++ b/net/websockets/websocket_job.cc
@@ -12,30 +12,6 @@
 #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.
@@ -162,6 +138,11 @@ void WebSocketJob::DetachDelegate() {
   if (socket_)
     socket_->DetachDelegate();
   socket_ = NULL;
+  if (callback_) {
+    waiting_ = false;
+    callback_ = NULL;
+    Release();  // Balanced with OnStartOpenConnection().
+  }
 }
 
 int WebSocketJob::OnStartOpenConnection(
@@ -173,6 +154,7 @@ int WebSocketJob::OnStartOpenConnection(
   if (!waiting_)
     return OK;
   callback_ = callback;
+  AddRef();  // Balanced when callback_ becomes NULL.
   return ERR_IO_PENDING;
 }
 
@@ -209,6 +191,11 @@ 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);
 }
@@ -437,16 +424,24 @@ 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(runner.get(),
-                        &CompletionCallbackRunner::Run));
+      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().
+  }
 }
 
 }  // namespace net
diff --git a/net/websockets/websocket_job.h b/net/websockets/websocket_job.h
index 23140c5..b3ccd34 100644
--- a/net/websockets/websocket_job.h
+++ b/net/websockets/websocket_job.h
@@ -82,6 +82,7 @@ class WebSocketJob : public SocketStreamJob, public SocketStream::Delegate {
   void SetWaiting();
   bool IsWaiting() const;
   void Wakeup();
+  void DoCallback();
 
   SocketStream::Delegate* delegate_;
   State state_;
diff --git a/net/websockets/websocket_throttle_unittest.cc b/net/websockets/websocket_throttle_unittest.cc
index d568292..61e5e4b 100644
--- a/net/websockets/websocket_throttle_unittest.cc
+++ b/net/websockets/websocket_throttle_unittest.cc
@@ -61,8 +61,13 @@ class WebSocketThrottleTest : public PlatformTest {
     }
   }
 
-  static void SetAddressList(SocketStream* socket, struct addrinfo* head) {
+  static void MockSocketStreamConnect(
+      SocketStream* socket, struct addrinfo* head) {
     socket->CopyAddrInfo(head);
+    // Add reference to socket as done in SocketStream::Connect().
+    // Balanced with Release() in SocketStream::Finish() which will be
+    // called by SocketStream::DetachDelegate().
+    socket->AddRef();
   }
 };
 
@@ -77,7 +82,7 @@ TEST_F(WebSocketThrottleTest, Throttle) {
   scoped_refptr<SocketStream> s1 =
       new SocketStream(GURL("ws://host1/"), w1.get());
   w1->InitSocketStream(s1.get());
-  WebSocketThrottleTest::SetAddressList(s1, addr);
+  WebSocketThrottleTest::MockSocketStreamConnect(s1, addr);
   DeleteAddrInfo(addr);
 
   DLOG(INFO) << "socket1";
@@ -97,7 +102,7 @@ TEST_F(WebSocketThrottleTest, Throttle) {
   scoped_refptr<SocketStream> s2 =
       new SocketStream(GURL("ws://host2/"), w2.get());
   w2->InitSocketStream(s2.get());
-  WebSocketThrottleTest::SetAddressList(s2, addr);
+  WebSocketThrottleTest::MockSocketStreamConnect(s2, addr);
   DeleteAddrInfo(addr);
 
   DLOG(INFO) << "socket2";
@@ -116,7 +121,7 @@ TEST_F(WebSocketThrottleTest, Throttle) {
   scoped_refptr<SocketStream> s3 =
       new SocketStream(GURL("ws://host3/"), w3.get());
   w3->InitSocketStream(s3.get());
-  WebSocketThrottleTest::SetAddressList(s3, addr);
+  WebSocketThrottleTest::MockSocketStreamConnect(s3, addr);
   DeleteAddrInfo(addr);
 
   DLOG(INFO) << "socket3";
@@ -135,7 +140,7 @@ TEST_F(WebSocketThrottleTest, Throttle) {
   scoped_refptr<SocketStream> s4 =
       new SocketStream(GURL("ws://host4/"), w4.get());
   w4->InitSocketStream(s4.get());
-  WebSocketThrottleTest::SetAddressList(s4, addr);
+  WebSocketThrottleTest::MockSocketStreamConnect(s4, addr);
   DeleteAddrInfo(addr);
 
   DLOG(INFO) << "socket4";
@@ -153,7 +158,7 @@ TEST_F(WebSocketThrottleTest, Throttle) {
   scoped_refptr<SocketStream> s5 =
       new SocketStream(GURL("ws://host5/"), w5.get());
   w5->InitSocketStream(s5.get());
-  WebSocketThrottleTest::SetAddressList(s5, addr);
+  WebSocketThrottleTest::MockSocketStreamConnect(s5, addr);
   DeleteAddrInfo(addr);
 
   DLOG(INFO) << "socket5";
@@ -171,7 +176,7 @@ TEST_F(WebSocketThrottleTest, Throttle) {
   scoped_refptr<SocketStream> s6 =
       new SocketStream(GURL("ws://host6/"), w6.get());
   w6->InitSocketStream(s6.get());
-  WebSocketThrottleTest::SetAddressList(s6, addr);
+  WebSocketThrottleTest::MockSocketStreamConnect(s6, addr);
   DeleteAddrInfo(addr);
 
   DLOG(INFO) << "socket6";