Refactor WebSocket throttling feature.

net/websockets/websocket_job.cc

 // Copyright (c) 2010 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.
 
 #include "net/websockets/websocket_job.h"
 
 #include "googleurl/src/gurl.h"
 #include "net/base/net_errors.h"
 #include "net/base/cookie_policy.h"
 #include "net/base/cookie_store.h"
 #include "net/http/http_util.h"
 #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.
 static const char* const kCookieHeaders[] = {
   "cookie", "cookie2"
 };
 static const char* const kSetCookieHeaders[] = {
   "set-cookie", "set-cookie2"
 };
@@ skipping 45 matching lines @@
 }
 
 // static
 void WebSocketJob::EnsureInit() {
   Singleton<WebSocketJobInitSingleton>::get();
 }
 
 WebSocketJob::WebSocketJob(SocketStream::Delegate* delegate)
     : delegate_(delegate),
       state_(INITIALIZED),
+      waiting_(false),
+      callback_(NULL),
       handshake_request_sent_(0),
       handshake_response_header_length_(0),
       response_cookies_save_index_(0),
       ALLOW_THIS_IN_INITIALIZER_LIST(can_get_cookies_callback_(
           this, &WebSocketJob::OnCanGetCookiesCompleted)),
       ALLOW_THIS_IN_INITIALIZER_LIST(can_set_cookie_callback_(
           this, &WebSocketJob::OnCanSetCookieCompleted)) {
 }
 
 WebSocketJob::~WebSocketJob() {
@@ skipping 33 matching lines @@
 
 void WebSocketJob::RestartWithAuth(
     const std::wstring& username,
     const std::wstring& password) {
   state_ = CONNECTING;
   socket_->RestartWithAuth(username, password);
 }
 
 void WebSocketJob::DetachDelegate() {
   state_ = CLOSED;
+  Singleton<WebSocketThrottle>::get()->RemoveFromQueue(this);
+  Singleton<WebSocketThrottle>::get()->WakeupSocketIfNecessary();
+
   delegate_ = NULL;
   if (socket_)
     socket_->DetachDelegate();
   socket_ = NULL;
 }
 
+int WebSocketJob::OnStartOpenConnection(
+    SocketStream* socket, CompletionCallback* callback) {
+  DCHECK(!callback_);
+  state_ = CONNECTING;
+  addresses_.Copy(socket->address_list().head(), true);
+  Singleton<WebSocketThrottle>::get()->PutInQueue(this);
+  if (!waiting_)
+    return OK;
+  callback_ = callback;
+  return ERR_IO_PENDING;
+}
+
 void WebSocketJob::OnConnected(
     SocketStream* socket, int max_pending_send_allowed) {
   if (delegate_)
     delegate_->OnConnected(socket, max_pending_send_allowed);
 }
 
 void WebSocketJob::OnSentData(SocketStream* socket, int amount_sent) {
   if (state_ == CONNECTING) {
     OnSentHandshakeRequest(socket, amount_sent);
     return;
   }
   if (delegate_)
     delegate_->OnSentData(socket, amount_sent);
 }
 
 void WebSocketJob::OnReceivedData(
     SocketStream* socket, const char* data, int len) {
   if (state_ == CONNECTING) {
     OnReceivedHandshakeResponse(socket, data, len);
     return;
   }
   if (delegate_)
     delegate_->OnReceivedData(socket, data, len);
 }
 
 void WebSocketJob::OnClose(SocketStream* socket) {
   state_ = CLOSED;
+  Singleton<WebSocketThrottle>::get()->RemoveFromQueue(this);
+  Singleton<WebSocketThrottle>::get()->WakeupSocketIfNecessary();
+
   SocketStream::Delegate* delegate = delegate_;
   delegate_ = NULL;
   socket_ = NULL;
   if (delegate)
     delegate->OnClose(socket);
 }
 
 void WebSocketJob::OnAuthRequired(
     SocketStream* socket, AuthChallengeInfo* auth_info) {
   if (delegate_)
@@ skipping 144 matching lines @@
         std::string(handshake_response_.data() +
                     handshake_response_header_length_,
                     handshake_response_.size() -
                     handshake_response_header_length_);
     std::string received_data =
         handshake_response_status_line +
         filtered_handshake_response_header +
         "\r\n" +
         remaining_data;
     state_ = OPEN;
+    Singleton<WebSocketThrottle>::get()->RemoveFromQueue(this);
+    Singleton<WebSocketThrottle>::get()->WakeupSocketIfNecessary();
+
     if (delegate_)
       delegate_->OnReceivedData(socket_,
                                 received_data.data(), received_data.size());
     return;
   }
 
   AddRef();  // Balanced in OnCanSetCookieCompleted
 
   int policy = OK;
   if (socket_->context()->cookie_policy()) {
@@ skipping 31 matching lines @@
 
 GURL WebSocketJob::GetURLForCookies() const {
   GURL url = socket_->url();
   std::string scheme = socket_->is_secure() ? "https" : "http";
   url_canon::Replacements<char> replacements;
   replacements.SetScheme(scheme.c_str(),
                          url_parse::Component(0, scheme.length()));
   return url.ReplaceComponents(replacements);
 }
 
+const AddressList& WebSocketJob::address_list() const {
+  return addresses_;
+}
+
+void WebSocketJob::SetWaiting() {
+  waiting_ = true;
+}
+
+bool WebSocketJob::IsWaiting() const {
+  return waiting_;
+}
+
+void WebSocketJob::Wakeup() {
+  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));
+}
+
 }  // namespace net