Try to fix flaky websocket tests.

Some websoket layout tests became flaky from r41818.
This is because it adds websocket throttling in WebSocketJob.

Make sure Close() will call OnClose() even if it is waiting resolving or waiting
in throttling queue, so that WebSocketJob is removed from throttling queue
and wake up next WebSocketJob.

BUG=38397
TEST=layout tests websocket/tests passes

Review URL: http://codereview.chromium.org/1096001

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@42074 0039d316-1c4b-4281-b951-d872f2087c98

4 files changed, 40 insertions, 40 deletions

diff --git a/net/socket_stream/socket_stream.cc b/net/socket_stream/socket_stream.cc
index b6d3a65..1d10ad5 100644
--- a/net/socket_stream/socket_stream.cc
+++ b/net/socket_stream/socket_stream.cc
@@ -167,9 +167,8 @@ void SocketStream::Close() {
       "The current MessageLoop must exist";
   DCHECK_EQ(MessageLoop::TYPE_IO, MessageLoop::current()->type()) <<
       "The current MessageLoop must be TYPE_IO";
-  if (!socket_.get() || !socket_->IsConnected() || next_state_ == STATE_NONE)
-    return;
-  socket_->Disconnect();
+  if (socket_.get() && socket_->IsConnected())
+    socket_->Disconnect();
   next_state_ = STATE_CLOSE;
   // Close asynchronously, so that delegate won't be called
   // back before returning Close().
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";