Revert of DevTools: Removed refcounting from AndroidWebSocket

chrome/browser/devtools/device/port_forwarding_controller.cc

Index: chrome/browser/devtools/device/port_forwarding_controller.cc
diff --git a/chrome/browser/devtools/device/port_forwarding_controller.cc b/chrome/browser/devtools/device/port_forwarding_controller.cc
index 35c365212b3be24419d735a427eedab6aa860a9b..df610d199221fb35357ad699cf2d8bae1ccd4be0 100644
--- a/chrome/browser/devtools/device/port_forwarding_controller.cc
+++ b/chrome/browser/devtools/device/port_forwarding_controller.cc
@@ -59,11 +59,11 @@
                           int port,
                           const CounterCallback& callback,
                           int result,
-                          scoped_ptr<net::StreamSocket> socket) {
+                          net::StreamSocket* socket) {
     if (result < 0)
       return;
     SocketTunnel* tunnel = new SocketTunnel(callback);
-    tunnel->Start(socket.Pass(), host, port);
+    tunnel->Start(socket, host, port);
   }
 
  private:
@@ -75,9 +75,8 @@
     callback_.Run(1);
   }
 
-  void Start(scoped_ptr<net::StreamSocket> socket,
-             const std::string& host, int port) {
-    remote_socket_.swap(socket);
+  void Start(net::StreamSocket* socket, const std::string& host, int port) {
+    remote_socket_.reset(socket);
 
     host_resolver_ = net::HostResolver::CreateDefaultResolver(NULL);
     net::HostResolver::RequestInfo request_info(net::HostPortPair(host, port));
@@ -255,13 +254,15 @@
 }  // namespace
 
 class PortForwardingController::Connection
-    : public DevToolsAndroidBridge::AndroidWebSocket::Delegate {
+    : public DevToolsAndroidBridge::AndroidWebSocket::Delegate,
+      public base::RefCountedThreadSafe<
+          Connection,
+          content::BrowserThread::DeleteOnUIThread> {
  public:
   Connection(Registry* registry,
              scoped_refptr<DevToolsAndroidBridge::RemoteDevice> device,
              scoped_refptr<DevToolsAndroidBridge::RemoteBrowser> browser,
              const ForwardingMap& forwarding_map);
-  virtual ~Connection();
 
   const PortStatusMap& GetPortStatusMap();
 
@@ -274,6 +275,7 @@
       content::BrowserThread::UI>;
   friend class base::DeleteHelper<Connection>;
 
+  virtual ~Connection();
 
   typedef std::map<int, std::string> ForwardingMap;
 
@@ -290,25 +292,23 @@
   void ProcessBindResponse(int port, PortStatus status);
   void ProcessUnbindResponse(int port, PortStatus status);
 
-  static void UpdateSocketCountOnHandlerThread(
-      base::WeakPtr<Connection> weak_connection, int port, int increment);
+  void UpdateSocketCountOnHandlerThread(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() OVERRIDE;
+  virtual void OnSocketClosed(bool closed_by_device) OVERRIDE;
 
   PortForwardingController::Registry* registry_;
   scoped_refptr<DevToolsAndroidBridge::RemoteDevice> device_;
   scoped_refptr<DevToolsAndroidBridge::RemoteBrowser> browser_;
-  scoped_ptr<DevToolsAndroidBridge::AndroidWebSocket> web_socket_;
+  scoped_refptr<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);
 };
@@ -323,18 +323,27 @@
       browser_(browser),
       command_id_(0),
       connected_(false),
-      forwarding_map_(forwarding_map),
-      weak_factory_(this) {
+      forwarding_map_(forwarding_map) {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
   (*registry_)[device_->serial()] = this;
-  web_socket_.reset(
-      browser->CreateWebSocket(kDevToolsRemoteBrowserTarget, 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();
 }
 
 PortForwardingController::Connection::~Connection() {
   DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
-  DCHECK(registry_->find(device_->serial()) != registry_->end());
-  registry_->erase(device_->serial());
+  if (registry_) {
+    DCHECK(registry_->find(device_->serial()) != registry_->end());
+    registry_->erase(device_->serial());
+  }
 }
 
 void PortForwardingController::Connection::UpdateForwardingMap(
@@ -431,12 +440,10 @@
     port_status_.erase(it);
 }
 
-// static
 void PortForwardingController::Connection::UpdateSocketCountOnHandlerThread(
-    base::WeakPtr<Connection> weak_connection, int port, int increment) {
+    int port, int increment) {
   BrowserThread::PostTask(BrowserThread::UI, FROM_HERE,
-     base::Bind(&Connection::UpdateSocketCount,
-                weak_connection, port, increment));
+     base::Bind(&Connection::UpdateSocketCount, this, port, increment));
 }
 
 void PortForwardingController::Connection::UpdateSocketCount(
@@ -460,12 +467,19 @@
 
 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(kTetheringBind, ForwardingMap(), forwarding_map_);
 }
 
-void PortForwardingController::Connection::OnSocketClosed() {
-  delete this;
+void PortForwardingController::Connection::OnSocketClosed(
+    bool closed_by_device) {
+  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
+  Release();  // Balanced in the constructor.
 }
 
 void PortForwardingController::Connection::OnFrameRead(
@@ -505,8 +519,7 @@
   std::string destination_host = tokens[0];
 
   SocketTunnel::CounterCallback callback =
-      base::Bind(&Connection::UpdateSocketCountOnHandlerThread,
-                 weak_factory_.GetWeakPtr(), port);
+      base::Bind(&Connection::UpdateSocketCountOnHandlerThread, this, port);
 
   device_->OpenSocket(
       connection_id.c_str(),
@@ -593,7 +606,7 @@
     UpdateConnections();
   } else {
     StopListening();
-    STLDeleteValues(&registry_);
+    ShutdownConnections();
     NotifyListeners(DevicesStatus());
   }
 }
@@ -624,6 +637,12 @@
     it->second->UpdateForwardingMap(forwarding_map_);
 }
 
+void PortForwardingController::ShutdownConnections() {
+  for (Registry::iterator it = registry_.begin(); it != registry_.end(); ++it)
+    it->second->Shutdown();
+  registry_.clear();
+}
+
 void PortForwardingController::NotifyListeners(
     const DevicesStatus& status) const {
   Listeners copy(listeners_);  // Iterate over copy.