DevTools: Removed refcounting from AndroidWebSocket

Issue 387067 can be resolved by having port forwarding socket dependent on all other references to AndroidDeviceManager::Device.

This requires for lifetime of AWS to be manageable in the first place.

BUG=387067

Review URL: https://codereview.chromium.org/449883002

Cr-Commit-Position: refs/heads/master@{#296927}

11 files changed, 246 insertions, 223 deletions

diff --git a/chrome/browser/devtools/device/adb/adb_client_socket.cc b/chrome/browser/devtools/device/adb/adb_client_socket.cc
index 0845f92..41b3c9b 100644
--- a/chrome/browser/devtools/device/adb/adb_client_socket.cc
+++ b/chrome/browser/devtools/device/adb/adb_client_socket.cc
@@ -20,9 +20,6 @@ const char kOkayResponse[] = "OKAY";
 const char kHostTransportCommand[] = "host:transport:%s";
 const char kLocalhost[] = "127.0.0.1";
 
-typedef base::Callback<void(int, const std::string&)> CommandCallback;
-typedef base::Callback<void(int, net::StreamSocket*)> SocketCallback;
-
 std::string EncodeMessage(const std::string& message) {
   static const char kHexChars[] = "0123456789ABCDEF";
 
@@ -73,14 +70,14 @@ class AdbTransportSocket : public AdbClientSocket {
   void OnSocketAvailable(int result, const std::string& response) {
     if (!CheckNetResultOrDie(result))
       return;
-    callback_.Run(net::OK, socket_.release());
+    callback_.Run(net::OK, socket_.Pass());
     delete this;
   }
 
   bool CheckNetResultOrDie(int result) {
     if (result >= 0)
       return true;
-    callback_.Run(result, NULL);
+    callback_.Run(result, make_scoped_ptr<net::StreamSocket>(NULL));
     delete this;
     return false;
   }
diff --git a/chrome/browser/devtools/device/adb/adb_client_socket.h b/chrome/browser/devtools/device/adb/adb_client_socket.h
index f0d26ad..f331b02 100644
--- a/chrome/browser/devtools/device/adb/adb_client_socket.h
+++ b/chrome/browser/devtools/device/adb/adb_client_socket.h
@@ -13,7 +13,7 @@ class AdbClientSocket {
  public:
   typedef base::Callback<void(int, const std::string&)> CommandCallback;
   typedef base::Callback<void(int result,
-                              net::StreamSocket*)> SocketCallback;
+                              scoped_ptr<net::StreamSocket>)> SocketCallback;
 
   static void AdbQuery(int port,
                        const std::string& query,
diff --git a/chrome/browser/devtools/device/android_device_manager.cc b/chrome/browser/devtools/device/android_device_manager.cc
index 10978be..df99de7 100644
--- a/chrome/browser/devtools/device/android_device_manager.cc
+++ b/chrome/browser/devtools/device/android_device_manager.cc
@@ -50,9 +50,9 @@ static void PostSocketCallback(
     scoped_refptr<base::MessageLoopProxy> response_message_loop,
     const AndroidDeviceManager::SocketCallback& callback,
     int result,
-    net::StreamSocket* socket) {
-  response_message_loop->PostTask(FROM_HERE,
-                                  base::Bind(callback, result, socket));
+    scoped_ptr<net::StreamSocket> socket) {
+  response_message_loop->PostTask(
+      FROM_HERE, base::Bind(callback, result, base::Passed(&socket)));
 }
 
 class HttpRequest {
@@ -61,39 +61,41 @@ class HttpRequest {
   typedef AndroidDeviceManager::SocketCallback SocketCallback;
 
   static void CommandRequest(const std::string& request,
-                           const CommandCallback& callback,
-                           int result,
-                           net::StreamSocket* socket) {
+                             const CommandCallback& callback,
+                             int result,
+                             scoped_ptr<net::StreamSocket> socket) {
     if (result != net::OK) {
       callback.Run(result, std::string());
       return;
     }
-    new HttpRequest(socket, request, callback);
+    new HttpRequest(socket.Pass(), request, callback);
   }
 
   static void SocketRequest(const std::string& request,
-                          const SocketCallback& callback,
-                          int result,
-                          net::StreamSocket* socket) {
+                            const SocketCallback& callback,
+                            int result,
+                            scoped_ptr<net::StreamSocket> socket) {
     if (result != net::OK) {
-      callback.Run(result, NULL);
+      callback.Run(result, make_scoped_ptr<net::StreamSocket>(NULL));
       return;
     }
-    new HttpRequest(socket, request, callback);
+    new HttpRequest(socket.Pass(), request, callback);
   }
 
  private:
-  HttpRequest(net::StreamSocket* socket,
+  HttpRequest(scoped_ptr<net::StreamSocket> socket,
               const std::string& request,
               const CommandCallback& callback)
-      : socket_(socket), command_callback_(callback), body_pos_(0) {
+      : socket_(socket.Pass()),
+        command_callback_(callback),
+        body_pos_(0) {
     SendRequest(request);
   }
 
-  HttpRequest(net::StreamSocket* socket,
-                      const std::string& request,
-                      const SocketCallback& callback)
-    : socket_(socket),
+  HttpRequest(scoped_ptr<net::StreamSocket> socket,
+              const std::string& request,
+              const SocketCallback& callback)
+    : socket_(socket.Pass()),
       socket_callback_(callback),
       body_pos_(0) {
     SendRequest(request);
@@ -169,7 +171,7 @@ class HttpRequest {
       if (!command_callback_.is_null())
         command_callback_.Run(net::OK, response_.substr(body_pos_));
       else
-        socket_callback_.Run(net::OK, socket_.release());
+        socket_callback_.Run(net::OK, socket_.Pass());
       delete this;
       return;
     }
@@ -191,7 +193,7 @@ class HttpRequest {
     if (!command_callback_.is_null())
       command_callback_.Run(result, std::string());
     else
-      socket_callback_.Run(result, NULL);
+      socket_callback_.Run(result, make_scoped_ptr<net::StreamSocket>(NULL));
     delete this;
     return false;
   }
diff --git a/chrome/browser/devtools/device/android_device_manager.h b/chrome/browser/devtools/device/android_device_manager.h
index 219e5ce..fbc00cd 100644
--- a/chrome/browser/devtools/device/android_device_manager.h
+++ b/chrome/browser/devtools/device/android_device_manager.h
@@ -24,7 +24,8 @@ class AndroidDeviceManager
       public base::NonThreadSafe {
  public:
   typedef base::Callback<void(int, const std::string&)> CommandCallback;
-  typedef base::Callback<void(int result, net::StreamSocket*)> SocketCallback;
+  typedef base::Callback<void(int result, scoped_ptr<net::StreamSocket>)>
+      SocketCallback;
   typedef base::Callback<void(const std::vector<std::string>&)> SerialsCallback;
 
   struct BrowserInfo {
@@ -53,32 +54,21 @@ class AndroidDeviceManager
 
   typedef base::Callback<void(const DeviceInfo&)> DeviceInfoCallback;
 
-  class AndroidWebSocket : public base::RefCountedThreadSafe<AndroidWebSocket> {
+  class AndroidWebSocket {
    public:
     class Delegate {
      public:
       virtual void OnSocketOpened() = 0;
       virtual void OnFrameRead(const std::string& message) = 0;
-      virtual void OnSocketClosed(bool closed_by_device) = 0;
+      virtual void OnSocketClosed() = 0;
 
      protected:
       virtual ~Delegate() {}
     };
 
-    AndroidWebSocket() {}
-
-    virtual void Connect() = 0;
-    virtual void Disconnect() = 0;
-    virtual void SendFrame(const std::string& message) = 0;
-    virtual void ClearDelegate() = 0;
-
-   protected:
     virtual ~AndroidWebSocket() {}
 
-   private:
-    friend class base::RefCountedThreadSafe<AndroidWebSocket>;
-
-    DISALLOW_COPY_AND_ASSIGN(AndroidWebSocket);
+    virtual void SendFrame(const std::string& message) = 0;
   };
 
   class DeviceProvider;
@@ -103,7 +93,7 @@ class AndroidDeviceManager
                      const std::string& url,
                      const SocketCallback& callback);
 
-    scoped_refptr<AndroidWebSocket> CreateWebSocket(
+    AndroidWebSocket* CreateWebSocket(
         const std::string& socket_name,
         const std::string& url,
         AndroidWebSocket::Delegate* delegate);
diff --git a/chrome/browser/devtools/device/android_web_socket.cc b/chrome/browser/devtools/device/android_web_socket.cc
index 7635889..2a3886b 100644
--- a/chrome/browser/devtools/device/android_web_socket.cc
+++ b/chrome/browser/devtools/device/android_web_socket.cc
@@ -2,6 +2,7 @@
 // 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"
@@ -18,85 +19,127 @@ namespace {
 
 const int kBufferSize = 16 * 1024;
 
-class WebSocketImpl : public AndroidDeviceManager::AndroidWebSocket {
+class WebSocketImpl {
  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;
-
- private:
-  friend class base::RefCountedThreadSafe<AndroidWebSocket>;
+  typedef AndroidDeviceManager::AndroidWebSocket::Delegate Delegate;
 
-  virtual ~WebSocketImpl();
+  WebSocketImpl(Delegate* delegate,
+                scoped_ptr<net::StreamSocket> socket);
+  void StartListening();
+  void SendFrame(const std::string& message);
 
-  void Connected(int result, net::StreamSocket* socket);
-  void StartListeningOnHandlerThread();
+ private:
   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 Disconnect();
+
+  Delegate* delegate_;
+  scoped_ptr<net::StreamSocket> socket_;
+  std::string response_buffer_;
+  std::string request_buffer_;
+  base::ThreadChecker thread_checker_;
+  DISALLOW_COPY_AND_ASSIGN(WebSocketImpl);
+};
+
+class DelegateWrapper
+    : public AndroidDeviceManager::AndroidWebSocket::Delegate {
+ public:
+  DelegateWrapper(base::WeakPtr<Delegate> weak_delegate,
+                  scoped_refptr<base::MessageLoopProxy> message_loop)
+      : weak_delegate_(weak_delegate),
+        message_loop_(message_loop) {
+  }
 
-  void OnSocketOpened();
-  void OnFrameRead(const std::string& message);
-  void OnSocketClosed(bool closed_by_device);
+  virtual ~DelegateWrapper() {}
+
+  // AndroidWebSocket::Delegate implementation
+  virtual void OnSocketOpened() OVERRIDE {
+    message_loop_->PostTask(FROM_HERE,
+        base::Bind(&Delegate::OnSocketOpened, weak_delegate_));
+  }
+
+  virtual void OnFrameRead(const std::string& message) OVERRIDE {
+    message_loop_->PostTask(FROM_HERE,
+        base::Bind(&Delegate::OnFrameRead, weak_delegate_, message));
+  }
+
+  virtual void OnSocketClosed() OVERRIDE {
+    message_loop_->PostTask(FROM_HERE,
+        base::Bind(&Delegate::OnSocketClosed, weak_delegate_));
+  }
+
+ private:
+  base::WeakPtr<Delegate> weak_delegate_;
+  scoped_refptr<base::MessageLoopProxy> message_loop_;
+};
+
+class AndroidWebSocketImpl
+    : public AndroidDeviceManager::AndroidWebSocket,
+      public AndroidDeviceManager::AndroidWebSocket::Delegate {
+ public:
+  typedef AndroidDeviceManager::Device Device;
+  AndroidWebSocketImpl(
+      scoped_refptr<base::MessageLoopProxy> device_message_loop,
+      scoped_refptr<Device> device,
+      const std::string& socket_name,
+      const std::string& url,
+      AndroidWebSocket::Delegate* delegate);
+
+  virtual ~AndroidWebSocketImpl();
+
+  // AndroidWebSocket implementation
+  virtual void SendFrame(const std::string& message) OVERRIDE;
+
+  // AndroidWebSocket::Delegate implementation
+  virtual void OnSocketOpened() OVERRIDE;
+  virtual void OnFrameRead(const std::string& message) OVERRIDE;
+  virtual void OnSocketClosed() OVERRIDE;
+
+ private:
+  void Connected(int result, scoped_ptr<net::StreamSocket> socket);
 
   scoped_refptr<base::MessageLoopProxy> device_message_loop_;
   scoped_refptr<Device> device_;
   std::string socket_name_;
   std::string url_;
-  scoped_ptr<net::StreamSocket> socket_;
-  Delegate* delegate_;
-  std::string response_buffer_;
-  std::string request_buffer_;
+  WebSocketImpl* connection_;
+  DelegateWrapper* delegate_wrapper_;
+  AndroidWebSocket::Delegate* delegate_;
+  base::WeakPtrFactory<AndroidWebSocketImpl> weak_factory_;
+  DISALLOW_COPY_AND_ASSIGN(AndroidWebSocketImpl);
 };
 
-WebSocketImpl::WebSocketImpl(
+AndroidWebSocketImpl::AndroidWebSocketImpl(
     scoped_refptr<base::MessageLoopProxy> device_message_loop,
     scoped_refptr<Device> device,
     const std::string& socket_name,
     const std::string& url,
-    Delegate* delegate)
+    AndroidWebSocket::Delegate* delegate)
     : device_message_loop_(device_message_loop),
       device_(device),
       socket_name_(socket_name),
       url_(url),
-      delegate_(delegate) {
-}
-
-void WebSocketImpl::Connect() {
+      delegate_(delegate),
+      weak_factory_(this) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
+  DCHECK(delegate_);
   device_->HttpUpgrade(
-      socket_name_, url_, base::Bind(&WebSocketImpl::Connected, this));
+      socket_name_, url_,
+      base::Bind(&AndroidWebSocketImpl::Connected, weak_factory_.GetWeakPtr()));
 }
 
-void WebSocketImpl::Disconnect() {
+void AndroidWebSocketImpl::SendFrame(const std::string& message) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
   device_message_loop_->PostTask(
       FROM_HERE,
-      base::Bind(&WebSocketImpl::DisconnectOnHandlerThread, this, false));
+      base::Bind(&WebSocketImpl::SendFrame,
+                 base::Unretained(connection_), message));
 }
 
 void WebSocketImpl::SendFrame(const std::string& message) {
-  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
-  device_message_loop_->PostTask(
-      FROM_HERE,
-      base::Bind(&WebSocketImpl::SendFrameOnHandlerThread, this, message));
-}
-
-void WebSocketImpl::ClearDelegate() {
-  delegate_ = NULL;
-}
-
-void WebSocketImpl::SendFrameOnHandlerThread(const std::string& message) {
-  DCHECK_EQ(device_message_loop_, base::MessageLoopProxy::current());
+  DCHECK(thread_checker_.CalledOnValidThread());
+  if (!socket_)
+    return;
   int mask = base::RandInt(0, 0x7FFFFFFF);
   std::string encoded_frame = WebSocket::EncodeFrameHybi17(message, mask);
   request_buffer_ += encoded_frame;
@@ -104,43 +147,55 @@ void WebSocketImpl::SendFrameOnHandlerThread(const std::string& message) {
     SendPendingRequests(0);
 }
 
-WebSocketImpl::~WebSocketImpl() {
+AndroidWebSocketImpl::~AndroidWebSocketImpl() {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
+  device_message_loop_->DeleteSoon(FROM_HERE, connection_);
+  device_message_loop_->DeleteSoon(FROM_HERE, delegate_wrapper_);
+}
+
+WebSocketImpl::WebSocketImpl(Delegate* delegate,
+                             scoped_ptr<net::StreamSocket> socket)
+                             : delegate_(delegate),
+                               socket_(socket.Pass()) {
+  thread_checker_.DetachFromThread();
 }
 
-void WebSocketImpl::Connected(int result, net::StreamSocket* socket) {
+void AndroidWebSocketImpl::Connected(int result,
+                                     scoped_ptr<net::StreamSocket> socket) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
   if (result != net::OK || socket == NULL) {
-    OnSocketClosed(true);
+    OnSocketClosed();
     return;
   }
-  socket_.reset(socket);
+  delegate_wrapper_ = new DelegateWrapper(weak_factory_.GetWeakPtr(),
+                                          base::MessageLoopProxy::current());
+  connection_ = new WebSocketImpl(delegate_wrapper_, socket.Pass());
   device_message_loop_->PostTask(
       FROM_HERE,
-      base::Bind(&WebSocketImpl::StartListeningOnHandlerThread, this));
+      base::Bind(&WebSocketImpl::StartListening,
+                 base::Unretained(connection_)));
   OnSocketOpened();
 }
 
-void WebSocketImpl::StartListeningOnHandlerThread() {
-  DCHECK_EQ(device_message_loop_, base::MessageLoopProxy::current());
+void WebSocketImpl::StartListening() {
+  DCHECK(thread_checker_.CalledOnValidThread());
+  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(&WebSocketImpl::OnBytesRead,
+                 base::Unretained(this), response_buffer));
   if (result != net::ERR_IO_PENDING)
     OnBytesRead(response_buffer, result);
 }
 
-void WebSocketImpl::OnBytesRead(
-    scoped_refptr<net::IOBuffer> response_buffer, int result) {
-  DCHECK_EQ(device_message_loop_, base::MessageLoopProxy::current());
-  if (!socket_)
-    return;
-
+void WebSocketImpl::OnBytesRead(scoped_refptr<net::IOBuffer> response_buffer,
+                                int result) {
+  DCHECK(thread_checker_.CalledOnValidThread());
   if (result <= 0) {
-    DisconnectOnHandlerThread(true);
+    Disconnect();
     return;
   }
 
@@ -153,32 +208,30 @@ void WebSocketImpl::OnBytesRead(
 
   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));
+    delegate_->OnFrameRead(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(&WebSocketImpl::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;
+  DCHECK(thread_checker_.CalledOnValidThread());
   if (result < 0) {
-    DisconnectOnHandlerThread(true);
+    Disconnect();
     return;
   }
   request_buffer_ = request_buffer_.substr(result);
@@ -189,43 +242,39 @@ void WebSocketImpl::SendPendingRequests(int result) {
       new net::StringIOBuffer(request_buffer_);
   result = socket_->Write(buffer.get(), buffer->size(),
                           base::Bind(&WebSocketImpl::SendPendingRequests,
-                                     this));
+                                     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();
-  BrowserThread::PostTask(BrowserThread::UI, FROM_HERE,
-      base::Bind(&WebSocketImpl::OnSocketClosed, this, closed_by_device));
+void WebSocketImpl::Disconnect() {
+  DCHECK(thread_checker_.CalledOnValidThread());
+  socket_.reset();
+  delegate_->OnSocketClosed();
 }
 
-void WebSocketImpl::OnSocketOpened() {
-  if (delegate_)
-    delegate_->OnSocketOpened();
+void AndroidWebSocketImpl::OnSocketOpened() {
+  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
+  delegate_->OnSocketOpened();
 }
 
-void WebSocketImpl::OnFrameRead(const std::string& message) {
-  if (delegate_)
-    delegate_->OnFrameRead(message);
+void AndroidWebSocketImpl::OnFrameRead(const std::string& message) {
+  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
+  delegate_->OnFrameRead(message);
 }
 
-void WebSocketImpl::OnSocketClosed(bool closed_by_device) {
-  if (delegate_)
-    delegate_->OnSocketClosed(closed_by_device);
+void AndroidWebSocketImpl::OnSocketClosed() {
+  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
+  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);
+  return new AndroidWebSocketImpl(
+      device_message_loop_, this, socket, url, delegate);
 }
diff --git a/chrome/browser/devtools/device/devtools_android_bridge.cc b/chrome/browser/devtools/device/devtools_android_bridge.cc
index e36e785..16d173b 100644
--- a/chrome/browser/devtools/device/devtools_android_bridge.cc
+++ b/chrome/browser/devtools/device/devtools_android_bridge.cc
@@ -187,11 +187,12 @@ class ProtocolCommand
  private:
   virtual void OnSocketOpened() OVERRIDE;
   virtual void OnFrameRead(const std::string& message) OVERRIDE;
-  virtual void OnSocketClosed(bool closed_by_device) OVERRIDE;
+  virtual void OnSocketClosed() OVERRIDE;
+  virtual ~ProtocolCommand();
 
   const std::string command_;
   const base::Closure callback_;
-  scoped_refptr<DevToolsAndroidBridge::AndroidWebSocket> web_socket_;
+  scoped_ptr<DevToolsAndroidBridge::AndroidWebSocket> web_socket_;
 
   DISALLOW_COPY_AND_ASSIGN(ProtocolCommand);
 };
@@ -202,9 +203,8 @@ ProtocolCommand::ProtocolCommand(
     const std::string& command,
     const base::Closure callback)
     : command_(command),
-      callback_(callback){
-  web_socket_ = browser->CreateWebSocket(debug_url, this);
-  web_socket_->Connect();
+      callback_(callback),
+      web_socket_(browser->CreateWebSocket(debug_url, this)) {
 }
 
 void ProtocolCommand::OnSocketOpened() {
@@ -212,16 +212,18 @@ void ProtocolCommand::OnSocketOpened() {
 }
 
 void ProtocolCommand::OnFrameRead(const std::string& message) {
-  web_socket_->Disconnect();
+  delete this;
 }
 
-void ProtocolCommand::OnSocketClosed(bool closed_by_device) {
-  if (!callback_.is_null()) {
-    callback_.Run();
-  }
+void ProtocolCommand::OnSocketClosed() {
   delete this;
 }
 
+ProtocolCommand::~ProtocolCommand() {
+  if (!callback_.is_null())
+    callback_.Run();
+}
+
 }  // namespace
 
 class AgentHostDelegate;
@@ -292,14 +294,15 @@ class AgentHostDelegate
       const std::string& message) OVERRIDE;
   virtual void OnSocketOpened() OVERRIDE;
   virtual void OnFrameRead(const std::string& message) OVERRIDE;
-  virtual void OnSocketClosed(bool closed_by_device) OVERRIDE;
+  virtual void OnSocketClosed() OVERRIDE;
 
   const std::string id_;
+  scoped_refptr<DevToolsAndroidBridge::RemoteBrowser> browser_;
+  const std::string debug_url_;
   bool socket_opened_;
-  bool detached_;
   bool is_web_view_;
   std::vector<std::string> pending_messages_;
-  scoped_refptr<DevToolsAndroidBridge::AndroidWebSocket> web_socket_;
+  scoped_ptr<DevToolsAndroidBridge::AndroidWebSocket> web_socket_;
   content::DevToolsAgentHost* agent_host_;
   content::DevToolsExternalAgentProxy* proxy_;
   DISALLOW_COPY_AND_ASSIGN(AgentHostDelegate);
@@ -328,10 +331,10 @@ AgentHostDelegate::AgentHostDelegate(
     scoped_refptr<DevToolsAndroidBridge::RemoteBrowser> browser,
     const std::string& debug_url)
     : id_(id),
+      browser_(browser),
+      debug_url_(debug_url),
       socket_opened_(false),
-      detached_(false),
       is_web_view_(browser->IsWebView()),
-      web_socket_(browser->CreateWebSocket(debug_url, this)),
       agent_host_(NULL),
       proxy_(NULL) {
   g_host_delegates.Get()[id] = this;
@@ -339,20 +342,17 @@ AgentHostDelegate::AgentHostDelegate(
 
 AgentHostDelegate::~AgentHostDelegate() {
   g_host_delegates.Get().erase(id_);
-  web_socket_->ClearDelegate();
 }
 
 void AgentHostDelegate::Attach(content::DevToolsExternalAgentProxy* proxy) {
   proxy_ = proxy;
   content::RecordAction(base::UserMetricsAction(is_web_view_ ?
       "DevTools_InspectAndroidWebView" : "DevTools_InspectAndroidPage"));
-  web_socket_->Connect();
+  web_socket_.reset(browser_->CreateWebSocket(debug_url_, this));
 }
 
 void AgentHostDelegate::Detach() {
-  detached_ = true;
-  if (socket_opened_)
-    web_socket_->Disconnect();
+  web_socket_.reset();
 }
 
 void AgentHostDelegate::SendMessageToBackend(const std::string& message) {
@@ -363,11 +363,6 @@ void AgentHostDelegate::SendMessageToBackend(const std::string& message) {
 }
 
 void AgentHostDelegate::OnSocketOpened() {
-  if (detached_) {
-    web_socket_->Disconnect();
-    return;
-  }
-
   socket_opened_ = true;
   for (std::vector<std::string>::iterator it = pending_messages_.begin();
        it != pending_messages_.end(); ++it) {
@@ -381,8 +376,8 @@ void AgentHostDelegate::OnFrameRead(const std::string& message) {
       proxy_->DispatchOnClientHost(message);
 }
 
-void AgentHostDelegate::OnSocketClosed(bool closed_by_device) {
-  if (proxy_ && closed_by_device)
+void AgentHostDelegate::OnSocketClosed() {
+  if (proxy_)
     proxy_->ConnectionClosed();
 }
 
@@ -621,7 +616,7 @@ DevToolsAndroidBridge::RemoteBrowser::GetAgentHost() {
       "adb:" + device_->serial() + ":" + socket_, this, kBrowserTargetSocket);
 }
 
-scoped_refptr<DevToolsAndroidBridge::AndroidWebSocket>
+DevToolsAndroidBridge::AndroidWebSocket*
 DevToolsAndroidBridge::RemoteBrowser::CreateWebSocket(
     const std::string& url,
     DevToolsAndroidBridge::AndroidWebSocket::Delegate* delegate) {
diff --git a/chrome/browser/devtools/device/devtools_android_bridge.h b/chrome/browser/devtools/device/devtools_android_bridge.h
index 651d1aa..2f283a1 100644
--- a/chrome/browser/devtools/device/devtools_android_bridge.h
+++ b/chrome/browser/devtools/device/devtools_android_bridge.h
@@ -122,7 +122,7 @@ class DevToolsAndroidBridge
 
     scoped_refptr<content::DevToolsAgentHost> GetAgentHost();
 
-    scoped_refptr<AndroidWebSocket> CreateWebSocket(
+    AndroidWebSocket* CreateWebSocket(
         const std::string& url,
         DevToolsAndroidBridge::AndroidWebSocket::Delegate* delegate);
 
diff --git a/chrome/browser/devtools/device/port_forwarding_controller.cc b/chrome/browser/devtools/device/port_forwarding_controller.cc
index 775fb6c..0bc9c17 100644
--- a/chrome/browser/devtools/device/port_forwarding_controller.cc
+++ b/chrome/browser/devtools/device/port_forwarding_controller.cc
@@ -56,11 +56,11 @@ class SocketTunnel : public base::NonThreadSafe {
                           int port,
                           const CounterCallback& callback,
                           int result,
-                          net::StreamSocket* socket) {
+                          scoped_ptr<net::StreamSocket> socket) {
     if (result < 0)
       return;
     SocketTunnel* tunnel = new SocketTunnel(callback);
-    tunnel->Start(socket, host, port);
+    tunnel->Start(socket.Pass(), host, port);
   }
 
  private:
@@ -72,8 +72,9 @@ class SocketTunnel : public base::NonThreadSafe {
     callback_.Run(1);
   }
 
-  void Start(net::StreamSocket* socket, const std::string& host, int port) {
-    remote_socket_.reset(socket);
+  void Start(scoped_ptr<net::StreamSocket> socket,
+             const std::string& host, int port) {
+    remote_socket_.swap(socket);
 
     host_resolver_ = net::HostResolver::CreateDefaultResolver(NULL);
     net::HostResolver::RequestInfo request_info(net::HostPortPair(host, port));
@@ -251,28 +252,23 @@ FindBestBrowserForTethering(
 }  // namespace
 
 class PortForwardingController::Connection
-    : public DevToolsAndroidBridge::AndroidWebSocket::Delegate,
-      public base::RefCountedThreadSafe<
-          Connection,
-          content::BrowserThread::DeleteOnUIThread> {
+    : public DevToolsAndroidBridge::AndroidWebSocket::Delegate {
  public:
   Connection(Registry* registry,
              scoped_refptr<DevToolsAndroidBridge::RemoteDevice> device,
              scoped_refptr<DevToolsAndroidBridge::RemoteBrowser> browser,
              const ForwardingMap& forwarding_map);
+  virtual ~Connection();
 
   const PortStatusMap& GetPortStatusMap();
 
   void UpdateForwardingMap(const ForwardingMap& new_forwarding_map);
 
-  void Shutdown();
-
  private:
   friend struct content::BrowserThread::DeleteOnThread<
       content::BrowserThread::UI>;
   friend class base::DeleteHelper<Connection>;
 
-  virtual ~Connection();
 
   typedef std::map<int, std::string> ForwardingMap;
 
@@ -289,23 +285,25 @@ class PortForwardingController::Connection
   void ProcessBindResponse(int port, PortStatus status);
   void ProcessUnbindResponse(int port, PortStatus status);
 
-  void UpdateSocketCountOnHandlerThread(int port, int increment);
+  static void UpdateSocketCountOnHandlerThread(
+      base::WeakPtr<Connection> weak_connection, int port, int increment);
   void UpdateSocketCount(int port, int increment);
 
   // DevToolsAndroidBridge::AndroidWebSocket::Delegate implementation:
   virtual void OnSocketOpened() OVERRIDE;
   virtual void OnFrameRead(const std::string& message) OVERRIDE;
-  virtual void OnSocketClosed(bool closed_by_device) OVERRIDE;
+  virtual void OnSocketClosed() OVERRIDE;
 
   PortForwardingController::Registry* registry_;
   scoped_refptr<DevToolsAndroidBridge::RemoteDevice> device_;
   scoped_refptr<DevToolsAndroidBridge::RemoteBrowser> browser_;
-  scoped_refptr<DevToolsAndroidBridge::AndroidWebSocket> web_socket_;
+  scoped_ptr<DevToolsAndroidBridge::AndroidWebSocket> web_socket_;
   int command_id_;
   bool connected_;
   ForwardingMap forwarding_map_;
   CommandCallbackMap pending_responses_;
   PortStatusMap port_status_;
+  base::WeakPtrFactory<Connection> weak_factory_;
 
   DISALLOW_COPY_AND_ASSIGN(Connection);
 };
@@ -320,27 +318,18 @@ PortForwardingController::Connection::Connection(
       browser_(browser),
       command_id_(0),
       connected_(false),
-      forwarding_map_(forwarding_map) {
+      forwarding_map_(forwarding_map),
+      weak_factory_(this) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
   (*registry_)[device_->serial()] = this;
-  web_socket_ = browser->CreateWebSocket(kDevToolsRemoteBrowserTarget, this);
-  web_socket_->Connect();
-  AddRef();  // Balanced in OnSocketClosed();
-}
-
-void PortForwardingController::Connection::Shutdown() {
-  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
-  registry_ = NULL;
-  // This will have no effect if the socket is not connected yet.
-  web_socket_->Disconnect();
+  web_socket_.reset(
+      browser->CreateWebSocket(kDevToolsRemoteBrowserTarget, this));
 }
 
 PortForwardingController::Connection::~Connection() {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
-  if (registry_) {
-    DCHECK(registry_->find(device_->serial()) != registry_->end());
-    registry_->erase(device_->serial());
-  }
+  DCHECK(registry_->find(device_->serial()) != registry_->end());
+  registry_->erase(device_->serial());
 }
 
 void PortForwardingController::Connection::UpdateForwardingMap(
@@ -442,10 +431,12 @@ void PortForwardingController::Connection::ProcessUnbindResponse(
     port_status_.erase(it);
 }
 
+// static
 void PortForwardingController::Connection::UpdateSocketCountOnHandlerThread(
-    int port, int increment) {
+    base::WeakPtr<Connection> weak_connection, int port, int increment) {
   BrowserThread::PostTask(BrowserThread::UI, FROM_HERE,
-     base::Bind(&Connection::UpdateSocketCount, this, port, increment));
+     base::Bind(&Connection::UpdateSocketCount,
+                weak_connection, port, increment));
 }
 
 void PortForwardingController::Connection::UpdateSocketCount(
@@ -469,19 +460,12 @@ PortForwardingController::Connection::GetPortStatusMap() {
 
 void PortForwardingController::Connection::OnSocketOpened() {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
-  if (!registry_) {
-    // Socket was created after Shutdown was called. Disconnect immediately.
-    web_socket_->Disconnect();
-    return;
-  }
   connected_ = true;
   SerializeChanges(tethering::bind::kName, ForwardingMap(), forwarding_map_);
 }
 
-void PortForwardingController::Connection::OnSocketClosed(
-    bool closed_by_device) {
-  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
-  Release();  // Balanced in the constructor.
+void PortForwardingController::Connection::OnSocketClosed() {
+  delete this;
 }
 
 void PortForwardingController::Connection::OnFrameRead(
@@ -522,7 +506,8 @@ void PortForwardingController::Connection::OnFrameRead(
   std::string destination_host = tokens[0];
 
   SocketTunnel::CounterCallback callback =
-      base::Bind(&Connection::UpdateSocketCountOnHandlerThread, this, port);
+      base::Bind(&Connection::UpdateSocketCountOnHandlerThread,
+                 weak_factory_.GetWeakPtr(), port);
 
   device_->OpenSocket(
       connection_id.c_str(),
@@ -591,7 +576,12 @@ void PortForwardingController::OnPrefsChange() {
   if (!forwarding_map_.empty()) {
     UpdateConnections();
   } else {
-    ShutdownConnections();
+    std::vector<Connection*> registry_copy;
+    for (Registry::iterator it = registry_.begin();
+        it != registry_.end(); ++it) {
+      registry_copy.push_back(it->second);
+    }
+    STLDeleteElements(&registry_copy);
   }
 }
 
@@ -599,9 +589,3 @@ void PortForwardingController::UpdateConnections() {
   for (Registry::iterator it = registry_.begin(); it != registry_.end(); ++it)
     it->second->UpdateForwardingMap(forwarding_map_);
 }
-
-void PortForwardingController::ShutdownConnections() {
-  for (Registry::iterator it = registry_.begin(); it != registry_.end(); ++it)
-    it->second->Shutdown();
-  registry_.clear();
-}
diff --git a/chrome/browser/devtools/device/port_forwarding_controller.h b/chrome/browser/devtools/device/port_forwarding_controller.h
index 290c268..c1074bf 100644
--- a/chrome/browser/devtools/device/port_forwarding_controller.h
+++ b/chrome/browser/devtools/device/port_forwarding_controller.h
@@ -35,7 +35,6 @@ class PortForwardingController {
   void OnPrefsChange();
 
   void UpdateConnections();
-  void ShutdownConnections();
 
   Profile* profile_;
   PrefService* pref_service_;
diff --git a/chrome/browser/devtools/device/self_device_provider.cc b/chrome/browser/devtools/device/self_device_provider.cc
index ee0f043..f6930b5 100644
--- a/chrome/browser/devtools/device/self_device_provider.cc
+++ b/chrome/browser/devtools/device/self_device_provider.cc
@@ -17,9 +17,9 @@ const char kSerial[] = "local";
 
 static void RunSocketCallback(
     const AndroidDeviceManager::SocketCallback& callback,
-    net::StreamSocket* socket,
+    scoped_ptr<net::StreamSocket> socket,
     int result) {
-  callback.Run(result, socket);
+  callback.Run(result, socket.Pass());
 }
 
 }  // namespace
@@ -61,7 +61,8 @@ void SelfAsDeviceProvider::OpenSocket(const std::string& serial,
   base::StringToInt(socket_name, &port);
   net::AddressList address_list =
       net::AddressList::CreateFromIPAddress(ip_number, port);
-  net::TCPClientSocket* socket = new net::TCPClientSocket(
-      address_list, NULL, net::NetLog::Source());
-  socket->Connect(base::Bind(&RunSocketCallback, callback, socket));
+  scoped_ptr<net::StreamSocket> socket(new net::TCPClientSocket(
+      address_list, NULL, net::NetLog::Source()));
+  socket->Connect(
+      base::Bind(&RunSocketCallback, callback, base::Passed(&socket)));
 }
diff --git a/chrome/browser/devtools/device/usb/usb_device_provider.cc b/chrome/browser/devtools/device/usb/usb_device_provider.cc
index a67d60c..258365c 100644
--- a/chrome/browser/devtools/device/usb/usb_device_provider.cc
+++ b/chrome/browser/devtools/device/usb/usb_device_provider.cc
@@ -19,9 +19,12 @@ const char kLocalAbstractCommand[] = "localabstract:%s";
 const int kBufferSize = 16 * 1024;
 
 void OnOpenSocket(const UsbDeviceProvider::SocketCallback& callback,
-                  net::StreamSocket* socket,
+                  net::StreamSocket* socket_raw,
                   int result) {
-  callback.Run(result, result == net::OK ? socket : NULL);
+  scoped_ptr<net::StreamSocket> socket(socket_raw);
+  if (result != net::OK)
+    socket.reset();
+  callback.Run(result, socket.Pass());
 }
 
 void OnRead(net::StreamSocket* socket,
@@ -68,7 +71,8 @@ void RunCommand(scoped_refptr<AndroidUsbDevice> device,
     callback.Run(net::ERR_CONNECTION_FAILED, std::string());
     return;
   }
-  int result = socket->Connect(base::Bind(&OpenedForCommand, callback, socket));
+  int result = socket->Connect(
+      base::Bind(&OpenedForCommand, callback, socket));
   if (result != net::ERR_IO_PENDING)
     callback.Run(result, std::string());
 }
@@ -107,19 +111,21 @@ void UsbDeviceProvider::OpenSocket(const std::string& serial,
                                    const SocketCallback& callback) {
   UsbDeviceMap::iterator it = device_map_.find(serial);
   if (it == device_map_.end()) {
-    callback.Run(net::ERR_CONNECTION_FAILED, NULL);
+    callback.Run(net::ERR_CONNECTION_FAILED,
+                 make_scoped_ptr<net::StreamSocket>(NULL));
     return;
   }
   std::string socket_name =
       base::StringPrintf(kLocalAbstractCommand, name.c_str());
   net::StreamSocket* socket = it->second->CreateSocket(socket_name);
   if (!socket) {
-    callback.Run(net::ERR_CONNECTION_FAILED, NULL);
+    callback.Run(net::ERR_CONNECTION_FAILED,
+                 make_scoped_ptr<net::StreamSocket>(NULL));
     return;
   }
   int result = socket->Connect(base::Bind(&OnOpenSocket, callback, socket));
   if (result != net::ERR_IO_PENDING)
-    callback.Run(result, NULL);
+    callback.Run(result, make_scoped_ptr<net::StreamSocket>(NULL));
 }
 
 void UsbDeviceProvider::ReleaseDevice(const std::string& serial) {