Add Socket::ReadIfReady()

net/socket/socket_test_util.cc

Index: net/socket/socket_test_util.cc
diff --git a/net/socket/socket_test_util.cc b/net/socket/socket_test_util.cc
index aeea26ad7a96fb75be2179c064f45ba5d080b8d8..74bbda51b155b797e96ced74413433451c1c4df2 100644
--- a/net/socket/socket_test_util.cc
+++ b/net/socket/socket_test_util.cc
@@ -10,6 +10,7 @@
 
 #include "base/bind.h"
 #include "base/bind_helpers.h"
+#include "base/callback_helpers.h"
 #include "base/compiler_specific.h"
 #include "base/location.h"
 #include "base/logging.h"
@@ -699,7 +700,8 @@ void SequencedSocketData::OnWriteComplete() {
 SequencedSocketData::~SequencedSocketData() {
 }
 
-MockClientSocketFactory::MockClientSocketFactory() {}
+MockClientSocketFactory::MockClientSocketFactory()
+    : enable_read_if_ready_(false) {}
 
 MockClientSocketFactory::~MockClientSocketFactory() {}
 
@@ -743,6 +745,8 @@ MockClientSocketFactory::CreateTransportClientSocket(
   SocketDataProvider* data_provider = mock_data_.GetNext();
   std::unique_ptr<MockTCPClientSocket> socket(
       new MockTCPClientSocket(addresses, net_log, data_provider));
+  if (enable_read_if_ready_)
+    socket->set_enable_read_if_ready(enable_read_if_ready_);
   return std::move(socket);
 }
 
@@ -882,7 +886,9 @@ MockTCPClientSocket::MockTCPClientSocket(const AddressList& addresses,
       peer_closed_connection_(false),
       pending_read_buf_(NULL),
       pending_read_buf_len_(0),
-      was_used_to_convey_data_(false) {
+      was_used_to_convey_data_(false),
+      enable_read_if_ready_(false),
+      weak_factory_(this) {
   DCHECK(data_);
   peer_addr_ = data->connect_data().peer_addr;
   data_->Initialize(this);
@@ -895,41 +901,28 @@ MockTCPClientSocket::~MockTCPClientSocket() {
 
 int MockTCPClientSocket::Read(IOBuffer* buf, int buf_len,
                               const CompletionCallback& callback) {
-  if (!connected_ || !data_)
-    return ERR_UNEXPECTED;
-
   // If the buffer is already in use, a read is already in progress!
   DCHECK(!pending_read_buf_);
 
-  // Store our async IO data.
-  pending_read_buf_ = buf;
-  pending_read_buf_len_ = buf_len;
-  pending_read_callback_ = callback;
-
-  if (need_read_data_) {
-    read_data_ = data_->OnRead();
-    if (read_data_.result == ERR_CONNECTION_CLOSED) {
-      // This MockRead is just a marker to instruct us to set
-      // peer_closed_connection_.
-      peer_closed_connection_ = true;
-    }
-    if (read_data_.result == ERR_TEST_PEER_CLOSE_AFTER_NEXT_MOCK_READ) {
-      // This MockRead is just a marker to instruct us to set
-      // peer_closed_connection_.  Skip it and get the next one.
-      read_data_ = data_->OnRead();
-      peer_closed_connection_ = true;
-    }
-    // ERR_IO_PENDING means that the SocketDataProvider is taking responsibility
-    // to complete the async IO manually later (via OnReadComplete).
-    if (read_data_.result == ERR_IO_PENDING) {
-      // We need to be using async IO in this case.
-      DCHECK(!callback.is_null());
-      return ERR_IO_PENDING;
-    }
-    need_read_data_ = false;
+  int rv = ReadIfReadyHelper(
+      buf, buf_len,
+      base::Bind(&MockTCPClientSocket::RetryRead, weak_factory_.GetWeakPtr()));

[davidben, 2017/02/10 23:33:48]

Ditto re whether you actually need the WeakPtr.

[xunjieli, 2017/02/13 20:28:18]

Done.

+  if (rv == ERR_IO_PENDING) {
+    pending_read_buf_ = buf;
+    pending_read_buf_len_ = buf_len;
+    pending_read_callback_ = callback;
   }
+  return rv;
+}
 
-  return CompleteRead();
+int MockTCPClientSocket::ReadIfReady(IOBuffer* buf,
+                                     int buf_len,
+                                     const CompletionCallback& callback) {
+  DCHECK(!pending_read_if_ready_callback_);
+
+  if (!enable_read_if_ready_)
+    return ERR_READ_IF_READY_NOT_IMPLEMENTED;
+  return ReadIfReadyHelper(buf, buf_len, callback);
 }
 
 int MockTCPClientSocket::Write(IOBuffer* buf, int buf_len,
@@ -1048,7 +1041,7 @@ void MockTCPClientSocket::OnReadComplete(const MockRead& data) {
     return;
 
   // There must be a read pending.
-  DCHECK(pending_read_buf_.get());
+  DCHECK(pending_read_if_ready_callback_);
   // You can't complete a read with another ERR_IO_PENDING status code.
   DCHECK_NE(ERR_IO_PENDING, data.result);
   // Since we've been waiting for data, need_read_data_ should be true.
@@ -1060,10 +1053,7 @@ void MockTCPClientSocket::OnReadComplete(const MockRead& data) {
   // The caller is simulating that this IO completes right now.  Don't
   // let CompleteRead() schedule a callback.
   read_data_.mode = SYNCHRONOUS;
-
-  CompletionCallback callback = pending_read_callback_;
-  int rv = CompleteRead();
-  RunCallback(callback, rv);
+  RunReadIfReadyCallback(read_data_.result > 0 ? OK : read_data_.result);
 }
 
 void MockTCPClientSocket::OnWriteComplete(int rv) {
@@ -1090,22 +1080,59 @@ void MockTCPClientSocket::OnDataProviderDestroyed() {
   data_ = nullptr;
 }
 
-int MockTCPClientSocket::CompleteRead() {
+void MockTCPClientSocket::RetryRead(int rv) {
+  DCHECK(pending_read_callback_);
   DCHECK(pending_read_buf_.get());
-  DCHECK(pending_read_buf_len_ > 0);
-
-  was_used_to_convey_data_ = true;
+  DCHECK_LT(0, pending_read_buf_len_);
 
-  // Save the pending async IO data and reset our |pending_| state.
-  scoped_refptr<IOBuffer> buf = pending_read_buf_;
-  int buf_len = pending_read_buf_len_;
-  CompletionCallback callback = pending_read_callback_;
-  pending_read_buf_ = NULL;
+  if (rv == OK) {
+    rv = ReadIfReadyHelper(pending_read_buf_.get(), pending_read_buf_len_,
+                           base::Bind(&MockTCPClientSocket::RetryRead,
+                                      weak_factory_.GetWeakPtr()));
+    if (rv == ERR_IO_PENDING)
+      return;
+  }
+  pending_read_buf_ = nullptr;
   pending_read_buf_len_ = 0;
-  pending_read_callback_.Reset();
+  base::ResetAndReturn(&pending_read_callback_).Run(rv);
+}
+
+int MockTCPClientSocket::ReadIfReadyHelper(IOBuffer* buf,
+                                           int buf_len,
+                                           const CompletionCallback& callback) {
+  if (!connected_ || !data_)
+    return ERR_UNEXPECTED;
+
+  DCHECK(!pending_read_if_ready_callback_);
+
+  if (need_read_data_) {
+    read_data_ = data_->OnRead();
+    if (read_data_.result == ERR_CONNECTION_CLOSED) {
+      // This MockRead is just a marker to instruct us to set
+      // peer_closed_connection_.
+      peer_closed_connection_ = true;
+    }
+    if (read_data_.result == ERR_TEST_PEER_CLOSE_AFTER_NEXT_MOCK_READ) {
+      // This MockRead is just a marker to instruct us to set
+      // peer_closed_connection_.  Skip it and get the next one.
+      read_data_ = data_->OnRead();
+      peer_closed_connection_ = true;
+    }
+    // ERR_IO_PENDING means that the SocketDataProvider is taking responsibility
+    // to complete the async IO manually later (via OnReadComplete).

[davidben, 2017/02/10 23:33:48]

Hrm. Does this feature work for ReadIfReady? (I don't really understand how
these classes work.)

[xunjieli, 2017/02/13 20:28:18]

I believe it does, but I might be missing something.

MockTCPClientSocket::OnReadComplete() will be called by SocketDataProvider.
OnReadComplete() calls RetryRead() or the real user read callback. 

+    if (read_data_.result == ERR_IO_PENDING) {
+      // We need to be using async IO in this case.
+      DCHECK(!callback.is_null());
+      pending_read_if_ready_callback_ = callback;
+      return ERR_IO_PENDING;
+    }
+    need_read_data_ = false;
+  }
+
+  was_used_to_convey_data_ = true;
 
   int result = read_data_.result;
-  DCHECK(result != ERR_IO_PENDING);
+  DCHECK_NE(ERR_IO_PENDING, result);
 
   if (read_data_.data) {
     if (read_data_.data_len - read_offset_ > 0) {
@@ -1123,12 +1150,21 @@ int MockTCPClientSocket::CompleteRead() {
 
   if (read_data_.mode == ASYNC) {
     DCHECK(!callback.is_null());
-    RunCallbackAsync(callback, result);
+    pending_read_if_ready_callback_ = callback;
+    read_data_.mode = SYNCHRONOUS;
+    base::ThreadTaskRunnerHandle::Get()->PostTask(
+        FROM_HERE, base::Bind(&MockTCPClientSocket::RunReadIfReadyCallback,
+                              weak_factory_.GetWeakPtr(), result));

[davidben, 2017/02/10 23:33:48]

Does RunCallbackAsync(callback, result) not work?

[xunjieli, 2017/02/13 20:28:18]

Done.

     return ERR_IO_PENDING;
   }
   return result;
 }
 
+void MockTCPClientSocket::RunReadIfReadyCallback(int result) {
+  DCHECK(pending_read_if_ready_callback_);
+  base::ResetAndReturn(&pending_read_if_ready_callback_).Run(result);
+}
+
 // static
 void MockSSLClientSocket::ConnectCallback(
     MockSSLClientSocket* ssl_client_socket,
@@ -1164,6 +1200,12 @@ int MockSSLClientSocket::Read(IOBuffer* buf, int buf_len,
   return transport_->socket()->Read(buf, buf_len, callback);
 }
 
+int MockSSLClientSocket::ReadIfReady(IOBuffer* buf,
+                                     int buf_len,
+                                     const CompletionCallback& callback) {
+  return transport_->socket()->ReadIfReady(buf, buf_len, callback);
+}
+
 int MockSSLClientSocket::Write(IOBuffer* buf, int buf_len,
                                const CompletionCallback& callback) {
   return transport_->socket()->Write(buf, buf_len, callback);