DevTools: Removed refcounting from AndroidWebSocket

chrome/browser/devtools/device/android_web_socket.cc

 // Copyright 2014 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 "base/memory/weak_ptr.h"
 #include "base/message_loop/message_loop.h"
 #include "base/rand_util.h"
 #include "chrome/browser/devtools/device/android_device_manager.h"
 #include "content/public/browser/browser_thread.h"
 #include "net/base/io_buffer.h"
 #include "net/base/net_errors.h"
 #include "net/server/web_socket.h"
 #include "net/socket/stream_socket.h"
 
 using content::BrowserThread;
 using net::WebSocket;
 
 namespace {
 
 const int kBufferSize = 16 * 1024;
 
 class WebSocketImpl : public AndroidDeviceManager::AndroidWebSocket {
  public:
   typedef AndroidDeviceManager::Device Device;
   WebSocketImpl(scoped_refptr<base::MessageLoopProxy> device_message_loop,
                 scoped_refptr<Device> device,
                 const std::string& socket_name,
                 const std::string& url,
                 Delegate* delegate);
 
-  virtual void Connect() OVERRIDE;
-  virtual void Disconnect() OVERRIDE;
   virtual void SendFrame(const std::string& message) OVERRIDE;
-  virtual void ClearDelegate() OVERRIDE;
+  virtual ~WebSocketImpl();
 
  private:
-  friend class base::RefCountedThreadSafe<AndroidWebSocket>;
 
-  virtual ~WebSocketImpl();
+  // Counterpart of WebSocketImpl existing on handler thread. Constructed on UI.
+  // All other members must be accessed on handler thread. Owned by
+  // corresponding WebSocketImpl and reports back to it via weak pointer.
+  class Connection {

[pfeldman, 2014/08/11 11:41:38]

Lets inherit from NonThreadSafe?

[vkuzkokov, 2014/08/11 14:40:03]

NonThreadSafe is not recommended in this case. Added ThreadChecker instead.

+   public:
+    Connection(base::WeakPtr<WebSocketImpl> weak_web_socket,
+               net::StreamSocket* socket);
+    void StartListening();
+    void SendFrame(const std::string& message);

[pfeldman, 2014/08/11 11:41:38]

Blink line below.

[vkuzkokov, 2014/08/11 14:40:03]

Done.

+   private:
+    void OnBytesRead(scoped_refptr<net::IOBuffer> response_buffer, int result);
+    void SendPendingRequests(int result);
+    void Disconnect();
+    base::WeakPtr<WebSocketImpl> weak_web_socket_;
+    scoped_ptr<net::StreamSocket> socket_;
+    std::string response_buffer_;
+    std::string request_buffer_;
+    DISALLOW_COPY_AND_ASSIGN(Connection);
+  };
 
+  static void ConnectedWeak(base::WeakPtr<WebSocketImpl> weak_web_socket,
+                            int result, net::StreamSocket* socket);
   void Connected(int result, net::StreamSocket* socket);
-  void StartListeningOnHandlerThread();
-  void OnBytesRead(scoped_refptr<net::IOBuffer> response_buffer, int result);
-  void SendFrameOnHandlerThread(const std::string& message);
-  void SendPendingRequests(int result);
-  void DisconnectOnHandlerThread(bool closed_by_device);
 
   void OnSocketOpened();
   void OnFrameRead(const std::string& message);
-  void OnSocketClosed(bool closed_by_device);
+  void OnSocketClosed();
 
   scoped_refptr<base::MessageLoopProxy> device_message_loop_;
   scoped_refptr<Device> device_;
   std::string socket_name_;
   std::string url_;
-  scoped_ptr<net::StreamSocket> socket_;
+  Connection* connection_;
   Delegate* delegate_;
-  std::string response_buffer_;
-  std::string request_buffer_;
+  base::WeakPtrFactory<WebSocketImpl> weak_factory_;
+  DISALLOW_COPY_AND_ASSIGN(WebSocketImpl);
 };
 
 WebSocketImpl::WebSocketImpl(
     scoped_refptr<base::MessageLoopProxy> device_message_loop,
     scoped_refptr<Device> device,
     const std::string& socket_name,
     const std::string& url,
     Delegate* delegate)
     : device_message_loop_(device_message_loop),
       device_(device),
       socket_name_(socket_name),
       url_(url),
-      delegate_(delegate) {
-}
-
-void WebSocketImpl::Connect() {
-  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
+      delegate_(delegate),
+      weak_factory_(this) {
+  DCHECK(delegate_);
   device_->HttpUpgrade(
-      socket_name_, url_, base::Bind(&WebSocketImpl::Connected, this));
-}
-
-void WebSocketImpl::Disconnect() {
-  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
-  device_message_loop_->PostTask(
-      FROM_HERE,
-      base::Bind(&WebSocketImpl::DisconnectOnHandlerThread, this, false));
+      socket_name_, url_,
+      base::Bind(&WebSocketImpl::ConnectedWeak, weak_factory_.GetWeakPtr()));
 }
 
 void WebSocketImpl::SendFrame(const std::string& message) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
   device_message_loop_->PostTask(
       FROM_HERE,
-      base::Bind(&WebSocketImpl::SendFrameOnHandlerThread, this, message));
+      base::Bind(&Connection::SendFrame,
+                 base::Unretained(connection_), message));
 }
 
-void WebSocketImpl::ClearDelegate() {
-  delegate_ = NULL;
-}
-
-void WebSocketImpl::SendFrameOnHandlerThread(const std::string& message) {
-  DCHECK_EQ(device_message_loop_, base::MessageLoopProxy::current());
+void WebSocketImpl::Connection::SendFrame(const std::string& message) {
+  if (!socket_)
+    return;
   int mask = base::RandInt(0, 0x7FFFFFFF);
   std::string encoded_frame = WebSocket::EncodeFrameHybi17(message, mask);
   request_buffer_ += encoded_frame;
   if (request_buffer_.length() == encoded_frame.length())
     SendPendingRequests(0);
 }
 
 WebSocketImpl::~WebSocketImpl() {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
+  device_message_loop_->DeleteSoon(FROM_HERE, connection_);
 }
 
-void WebSocketImpl::Connected(int result, net::StreamSocket* socket) {
+WebSocketImpl::Connection::Connection(
+    base::WeakPtr<WebSocketImpl> weak_web_socket,
+    net::StreamSocket* socket)
+    : weak_web_socket_(weak_web_socket),
+      socket_(socket) {
+}
+
+// static
+void WebSocketImpl::ConnectedWeak(base::WeakPtr<WebSocketImpl> weak_web_socket,
+                                  int result, net::StreamSocket* socket_raw) {
+  // The sole purpose of a static method here is to avoid leaking StreamSocket

[pfeldman, 2014/08/11 11:41:38]

Doesn't base::Owned solve this for you?

[vkuzkokov, 2014/08/11 14:40:03]

It doesn't: we don't want StreamSocket destroyed after we exit this method.

base::Passed does help, however. It must be added on the other side
(AndroidDeviceManager::SocketCallback) to work, though.

+  // in case when requesting WebSocketImpl has already been destroyed.
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
+  scoped_ptr<net::StreamSocket> socket(socket_raw);
+  if (weak_web_socket)
+    weak_web_socket->Connected(result, socket.release());
+}
+
+void WebSocketImpl::Connected(int result, net::StreamSocket* socket_raw) {

[pfeldman, 2014/08/11 11:41:38]

Please DCHECK the threads in all methods (here and below)

[vkuzkokov, 2014/08/11 14:40:03]

Done.

+  scoped_ptr<net::StreamSocket> socket(socket_raw);
   if (result != net::OK || socket == NULL) {
-    OnSocketClosed(true);
+    OnSocketClosed();
     return;
   }
-  socket_.reset(socket);
+  connection_ = new Connection(weak_factory_.GetWeakPtr(), socket.release());
   device_message_loop_->PostTask(
       FROM_HERE,
-      base::Bind(&WebSocketImpl::StartListeningOnHandlerThread, this));
+      base::Bind(&Connection::StartListening, base::Unretained(connection_)));
   OnSocketOpened();
 }
 
-void WebSocketImpl::StartListeningOnHandlerThread() {
-  DCHECK_EQ(device_message_loop_, base::MessageLoopProxy::current());
+void WebSocketImpl::Connection::StartListening() {
+  DCHECK(socket_);
   scoped_refptr<net::IOBuffer> response_buffer =
       new net::IOBuffer(kBufferSize);
   int result = socket_->Read(
       response_buffer.get(),
       kBufferSize,
-      base::Bind(&WebSocketImpl::OnBytesRead, this, response_buffer));
+      base::Bind(&Connection::OnBytesRead,
+                 base::Unretained(this), response_buffer));
   if (result != net::ERR_IO_PENDING)
     OnBytesRead(response_buffer, result);
 }
 
-void WebSocketImpl::OnBytesRead(
+void WebSocketImpl::Connection::OnBytesRead(
     scoped_refptr<net::IOBuffer> response_buffer, int result) {
-  DCHECK_EQ(device_message_loop_, base::MessageLoopProxy::current());
-  if (!socket_)
-    return;
-
   if (result <= 0) {
-    DisconnectOnHandlerThread(true);
+    Disconnect();
     return;
   }
 
   std::string data = std::string(response_buffer->data(), result);
   response_buffer_ += data;
 
   int bytes_consumed;
   std::string output;
   WebSocket::ParseResult parse_result = WebSocket::DecodeFrameHybi17(
       response_buffer_, false, &bytes_consumed, &output);
 
   while (parse_result == WebSocket::FRAME_OK) {
     response_buffer_ = response_buffer_.substr(bytes_consumed);
     BrowserThread::PostTask(BrowserThread::UI, FROM_HERE,
-        base::Bind(&WebSocketImpl::OnFrameRead, this, output));
+        base::Bind(&WebSocketImpl::OnFrameRead, weak_web_socket_, output));
     parse_result = WebSocket::DecodeFrameHybi17(
         response_buffer_, false, &bytes_consumed, &output);
   }
 
   if (parse_result == WebSocket::FRAME_ERROR ||
       parse_result == WebSocket::FRAME_CLOSE) {
-    DisconnectOnHandlerThread(true);
+    Disconnect();
     return;
   }
 
   result = socket_->Read(
       response_buffer.get(),
       kBufferSize,
-      base::Bind(&WebSocketImpl::OnBytesRead, this, response_buffer));
+      base::Bind(&Connection::OnBytesRead,
+                 base::Unretained(this), response_buffer));
   if (result != net::ERR_IO_PENDING)
     OnBytesRead(response_buffer, result);
 }
 
-void WebSocketImpl::SendPendingRequests(int result) {
-  DCHECK_EQ(device_message_loop_, base::MessageLoopProxy::current());
-  if (!socket_)
-    return;
+void WebSocketImpl::Connection::SendPendingRequests(int result) {
   if (result < 0) {
-    DisconnectOnHandlerThread(true);
+    Disconnect();
     return;
   }
   request_buffer_ = request_buffer_.substr(result);
   if (request_buffer_.empty())
     return;
 
   scoped_refptr<net::StringIOBuffer> buffer =
       new net::StringIOBuffer(request_buffer_);
   result = socket_->Write(buffer.get(), buffer->size(),
-                          base::Bind(&WebSocketImpl::SendPendingRequests,
-                                     this));
+                          base::Bind(&Connection::SendPendingRequests,
+                                     base::Unretained(this)));
   if (result != net::ERR_IO_PENDING)
     SendPendingRequests(result);
 }
 
-void WebSocketImpl::DisconnectOnHandlerThread(bool closed_by_device) {
-  DCHECK_EQ(device_message_loop_, base::MessageLoopProxy::current());
-  if (!socket_)
-    return;
-  // Wipe out socket_ first since Disconnect can re-enter this method.
-  scoped_ptr<net::StreamSocket> socket(socket_.release());
-  socket->Disconnect();
+void WebSocketImpl::Connection::Disconnect() {
+  socket_.reset();
   BrowserThread::PostTask(BrowserThread::UI, FROM_HERE,
-      base::Bind(&WebSocketImpl::OnSocketClosed, this, closed_by_device));
+      base::Bind(&WebSocketImpl::OnSocketClosed, weak_web_socket_));

[pfeldman, 2014/08/11 11:41:38]

weak_web_socket_ could be pointing to the garbage in memory.

[vkuzkokov, 2014/08/11 14:40:03]

That's why we use weak pointer.

 }
 
 void WebSocketImpl::OnSocketOpened() {
-  if (delegate_)
-    delegate_->OnSocketOpened();
+  delegate_->OnSocketOpened();
 }
 
 void WebSocketImpl::OnFrameRead(const std::string& message) {
-  if (delegate_)
-    delegate_->OnFrameRead(message);
+  delegate_->OnFrameRead(message);
 }
 
-void WebSocketImpl::OnSocketClosed(bool closed_by_device) {
-  if (delegate_)
-    delegate_->OnSocketClosed(closed_by_device);
+void WebSocketImpl::OnSocketClosed() {
+  delegate_->OnSocketClosed();
 }
 
 }  // namespace
 
-scoped_refptr<AndroidDeviceManager::AndroidWebSocket>
+AndroidDeviceManager::AndroidWebSocket*
 AndroidDeviceManager::Device::CreateWebSocket(
     const std::string& socket,
     const std::string& url,
     AndroidDeviceManager::AndroidWebSocket::Delegate* delegate) {
   return new WebSocketImpl(device_message_loop_, this, socket, url, delegate);
 }