Try to fix flaky websocket tests again.

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.

r42074 was reverted because failure in SocketStreamMetrics::OnClose
on Linux Builder (Views dbg).
Ignore UMA record if the connection has not been established.

TBR=tyoshino
BUG=38397
TEST=layout tests websocket/tests passes

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

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

5 files changed, 56 insertions, 51 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/socket_stream/socket_stream_metrics.cc b/net/socket_stream/socket_stream_metrics.cc
index 625a491..71239af 100644
--- a/net/socket_stream/socket_stream_metrics.cc
+++ b/net/socket_stream/socket_stream_metrics.cc
@@ -70,16 +70,18 @@ void SocketStreamMetrics::OnWrite(int len) {
 
 void SocketStreamMetrics::OnClose() {
   base::TimeTicks closed_time = base::TimeTicks::Now();
-  UMA_HISTOGRAM_LONG_TIMES("Net.SocketStream.Duration",
-                           closed_time - connect_establish_time_);
-  UMA_HISTOGRAM_COUNTS("Net.SocketStream.ReceivedBytes",
-                       received_bytes_);
-  UMA_HISTOGRAM_COUNTS("Net.SocketStream.ReceivedCounts",
-                       received_counts_);
-  UMA_HISTOGRAM_COUNTS("Net.SocketStream.SentBytes",
-                       sent_bytes_);
-  UMA_HISTOGRAM_COUNTS("Net.SocketStream.SentCounts",
-                       sent_counts_);
+  if (!connect_establish_time_.is_null()) {
+    UMA_HISTOGRAM_LONG_TIMES("Net.SocketStream.Duration",
+                             closed_time - connect_establish_time_);
+    UMA_HISTOGRAM_COUNTS("Net.SocketStream.ReceivedBytes",
+                         received_bytes_);
+    UMA_HISTOGRAM_COUNTS("Net.SocketStream.ReceivedCounts",
+                         received_counts_);
+    UMA_HISTOGRAM_COUNTS("Net.SocketStream.SentBytes",
+                         sent_bytes_);
+    UMA_HISTOGRAM_COUNTS("Net.SocketStream.SentCounts",
+                         sent_counts_);
+  }
 }
 
 void SocketStreamMetrics::CountProtocolType(ProtocolType protocol_type) {
diff --git a/net/websockets/websocket_job.cc b/net/websockets/websocket_job.cc
index 62a62b7..e43a68a 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.
@@ -158,10 +134,17 @@ void WebSocketJob::DetachDelegate() {
   Singleton<WebSocketThrottle>::get()->RemoveFromQueue(this);
   Singleton<WebSocketThrottle>::get()->WakeupSocketIfNecessary();
 
+  scoped_refptr<WebSocketJob> protect(this);
+
   delegate_ = NULL;
   if (socket_)
     socket_->DetachDelegate();
   socket_ = NULL;
+  if (callback_) {
+    waiting_ = false;
+    callback_ = NULL;
+    Release();  // Balanced with OnStartOpenConnection().
+  }
 }
 
 int WebSocketJob::OnStartOpenConnection(
@@ -173,6 +156,7 @@ int WebSocketJob::OnStartOpenConnection(
   if (!waiting_)
     return OK;
   callback_ = callback;
+  AddRef();  // Balanced when callback_ becomes NULL.
   return ERR_IO_PENDING;
 }
 
@@ -206,9 +190,16 @@ void WebSocketJob::OnClose(SocketStream* socket) {
   Singleton<WebSocketThrottle>::get()->RemoveFromQueue(this);
   Singleton<WebSocketThrottle>::get()->WakeupSocketIfNecessary();
 
+  scoped_refptr<WebSocketJob> protect(this);
+
   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 +428,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.
+  if (callback_) {
+    net::CompletionCallback* callback = callback_;
+    callback_ = NULL;
+    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..bb4ac1e 100644
--- a/net/websockets/websocket_job.h
+++ b/net/websockets/websocket_job.h
@@ -8,7 +8,6 @@
 #include <string>
 #include <vector>
 
-#include "base/ref_counted.h"
 #include "net/base/address_list.h"
 #include "net/base/completion_callback.h"
 #include "net/socket_stream/socket_stream_job.h"
@@ -82,6 +81,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";