Move fail on suspend logic from URLRequestJob to TcpClientSocket.

net/socket/tcp_client_socket.cc

Index: net/socket/tcp_client_socket.cc
diff --git a/net/socket/tcp_client_socket.cc b/net/socket/tcp_client_socket.cc
index f8f612295295c40b54b2d19cbd4981508648b731..30dae6b19a19df4b6ef5e8ba775a50d6a4003772 100644
--- a/net/socket/tcp_client_socket.cc
+++ b/net/socket/tcp_client_socket.cc
@@ -16,6 +16,10 @@
 #include "net/base/net_errors.h"
 #include "net/socket/socket_performance_watcher.h"
 
+#if !defined(OS_NACL)
+#include "base/power_monitor/power_monitor.h"
+#endif
+
 namespace net {
 
 class NetLogWithSource;
@@ -33,7 +37,8 @@ TCPClientSocket::TCPClientSocket(
       current_address_index_(-1),
       next_connect_state_(CONNECT_STATE_NONE),
       previously_disconnected_(false),
-      total_received_bytes_(0) {}
+      total_received_bytes_(0),
+      weak_ptr_factory_(this) {}
 
 TCPClientSocket::TCPClientSocket(std::unique_ptr<TCPSocket> connected_socket,
                                  const IPEndPoint& peer_address)
@@ -43,7 +48,8 @@ TCPClientSocket::TCPClientSocket(std::unique_ptr<TCPSocket> connected_socket,
       current_address_index_(0),
       next_connect_state_(CONNECT_STATE_NONE),
       previously_disconnected_(false),
-      total_received_bytes_(0) {
+      total_received_bytes_(0),
+      weak_ptr_factory_(this) {
   DCHECK(socket_);
 
   socket_->SetDefaultOptionsForClient();
@@ -51,6 +57,7 @@ TCPClientSocket::TCPClientSocket(std::unique_ptr<TCPSocket> connected_socket,
 }
 
 TCPClientSocket::~TCPClientSocket() {
+  SetFailOnSuspend(false);
   Disconnect();
 }
 
@@ -83,6 +90,8 @@ int TCPClientSocket::Connect(const CompletionCallback& callback) {
   if (socket_->IsValid() && current_address_index_ >= 0)
     return OK;
 
+  was_disconnected_on_suspend_ = false;
+
   socket_->StartLoggingMultipleConnectAttempts(addresses_);
 
   // We will try to connect to each address in addresses_. Start with the
@@ -200,6 +209,14 @@ void TCPClientSocket::Disconnect() {
   DoDisconnect();
   current_address_index_ = -1;
   bind_address_.reset();
+
+  // Cancel any pending callbacks. Not done in Disconnect() because that's
+  // called on connection failure, when the connect callback will need to be
+  // invoked.
+  was_disconnected_on_suspend_ = false;
+  connect_callback_.Reset();
+  read_callback_.Reset();
+  write_callback_.Reset();
 }
 
 void TCPClientSocket::DoDisconnect() {
@@ -209,6 +226,10 @@ void TCPClientSocket::DoDisconnect() {
   // disconnected.
   previously_disconnected_ = socket_->IsValid() && current_address_index_ >= 0;
   socket_->Close();
+
+  // Invalidate weak pointers, so if in the middle of a callback in OnSuspend,
+  // and something destroys this, no other callback is invoked.
+  weak_ptr_factory_.InvalidateWeakPtrs();
 }
 
 bool TCPClientSocket::IsConnected() const {
@@ -269,17 +290,41 @@ bool TCPClientSocket::GetSSLInfo(SSLInfo* ssl_info) {
   return false;
 }
 
+void TCPClientSocket::SetFailOnSuspend(bool disconnect_on_suspend) {
+// Do nothing if building under NaCl.
+#if !defined(OS_NACL)
+  // Do nothing if state is unchanged.
+  if (disconnect_on_suspend_ == disconnect_on_suspend)
+    return;
+
+  // Otherwise, start/stop observing if there's a PowerMonitor configured, as
+  // needed.
+  base::PowerMonitor* power_monitor = base::PowerMonitor::Get();
+  if (!power_monitor)
+    return;
+  disconnect_on_suspend_ = disconnect_on_suspend;
+  if (disconnect_on_suspend_) {
+    power_monitor->AddObserver(this);
+  } else {
+    power_monitor->RemoveObserver(this);
+  }
+#endif  // !defined(OS_NACL)
+}
+
 int TCPClientSocket::Read(IOBuffer* buf,
                           int buf_len,
                           const CompletionCallback& callback) {
   DCHECK(!callback.is_null());
+  DCHECK(read_callback_.is_null());
 
   // |socket_| is owned by this class and the callback won't be run once
   // |socket_| is gone. Therefore, it is safe to use base::Unretained() here.
-  CompletionCallback read_callback = base::Bind(
-      &TCPClientSocket::DidCompleteRead, base::Unretained(this), callback);
+  CompletionCallback read_callback =
+      base::Bind(&TCPClientSocket::DidCompleteRead, base::Unretained(this));
   int result = socket_->Read(buf, buf_len, read_callback);
-  if (result > 0) {
+  if (result == ERR_IO_PENDING) {
+    read_callback_ = callback;
+  } else if (result > 0) {
     use_history_.set_was_used_to_convey_data();
     total_received_bytes_ += result;
   }
@@ -291,14 +336,18 @@ int TCPClientSocket::Write(IOBuffer* buf,
                            int buf_len,
                            const CompletionCallback& callback) {
   DCHECK(!callback.is_null());
+  DCHECK(write_callback_.is_null());
 
   // |socket_| is owned by this class and the callback won't be run once
   // |socket_| is gone. Therefore, it is safe to use base::Unretained() here.
-  CompletionCallback write_callback = base::Bind(
-      &TCPClientSocket::DidCompleteWrite, base::Unretained(this), callback);
+  CompletionCallback write_callback =
+      base::Bind(&TCPClientSocket::DidCompleteWrite, base::Unretained(this));
   int result = socket_->Write(buf, buf_len, write_callback);
-  if (result > 0)
+  if (result == ERR_IO_PENDING) {
+    write_callback_ = callback;
+  } else if (result > 0) {
     use_history_.set_was_used_to_convey_data();
+  }
 
   return result;
 }
@@ -308,7 +357,7 @@ int TCPClientSocket::SetReceiveBufferSize(int32_t size) {
 }
 
 int TCPClientSocket::SetSendBufferSize(int32_t size) {
-    return socket_->SetSendBufferSize(size);
+  return socket_->SetSendBufferSize(size);
 }
 
 bool TCPClientSocket::SetKeepAlive(bool enable, int delay) {
@@ -337,6 +386,38 @@ int64_t TCPClientSocket::GetTotalReceivedBytes() const {
   return total_received_bytes_;
 }
 
+void TCPClientSocket::OnSuspend() {
+  // If the socket is connected, or connecting, act as if current and future
+  // operations on the socket fail with ERR_NETWORK_IO_SUSPENDED, until the
+  // socket is reconnected.
+  // TODO(mmenke):  This doesn't cover sockets created after OnSuspend runs,
+  // just before suspend mode starts. Would it make more sense to do this on
+  // resume?
+
+  if (next_connect_state_ != CONNECT_STATE_NONE) {
+    DidCompleteConnect(ERR_NETWORK_IO_SUSPENDED);
+    return;
+  }
+
+  // Nothing to do.
+  if (!IsConnected())
+    return;
+
+  Disconnect();
+
+  was_disconnected_on_suspend_ = true;
+
+  // Grab a weap pointer just in case calling read callback results in |this|
+  // being destroyed, or disconnected. In either case, should not run the write
+  // callback.
+  base::WeakPtr<TCPClientSocket> weak_this = weak_ptr_factory_.GetWeakPtr();
+
+  if (read_callback_)
+    DidCompleteRead(ERR_NETWORK_IO_SUSPENDED);
+  if (weak_this && write_callback_)
+    DidCompleteWrite(ERR_NETWORK_IO_SUSPENDED);
+}
+
 void TCPClientSocket::DidCompleteConnect(int result) {
   DCHECK_EQ(next_connect_state_, CONNECT_STATE_CONNECT_COMPLETE);
   DCHECK_NE(result, ERR_IO_PENDING);
@@ -349,17 +430,17 @@ void TCPClientSocket::DidCompleteConnect(int result) {
   }
 }
 
-void TCPClientSocket::DidCompleteRead(const CompletionCallback& callback,
-                                      int result) {
+void TCPClientSocket::DidCompleteRead(int result) {
+  DCHECK(!read_callback_.is_null());
   if (result > 0)
     total_received_bytes_ += result;
 
-  DidCompleteReadWrite(callback, result);
+  DidCompleteReadWrite(base::ResetAndReturn(&read_callback_), result);
 }
 
-void TCPClientSocket::DidCompleteWrite(const CompletionCallback& callback,
-                                       int result) {
-  DidCompleteReadWrite(callback, result);
+void TCPClientSocket::DidCompleteWrite(int result) {
+  DCHECK(!write_callback_.is_null());
+  DidCompleteReadWrite(base::ResetAndReturn(&write_callback_), result);
 }
 
 void TCPClientSocket::DidCompleteReadWrite(const CompletionCallback& callback,