Add Socket::ReadIfReady()

net/socket/socket_test_util.cc

 // Copyright (c) 2012 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 "net/socket/socket_test_util.h"
 
 #include <algorithm>
 #include <utility>
 #include <vector>
 
 #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"
 #include "base/memory/ptr_util.h"
 #include "base/run_loop.h"
 #include "base/single_thread_task_runner.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "base/time/time.h"
 #include "net/base/address_family.h"
 #include "net/base/address_list.h"
@@ skipping 669 matching lines @@
 
   NET_TRACE(1, " *** ") << " Completing socket write for: "
                         << data.sequence_number;
   socket()->OnWriteComplete(rv);
   NET_TRACE(1, " *** ") << "Done";
 }
 
 SequencedSocketData::~SequencedSocketData() {
 }
 
-MockClientSocketFactory::MockClientSocketFactory() {}
+MockClientSocketFactory::MockClientSocketFactory()
+    : enable_read_if_ready_(false) {}
 
 MockClientSocketFactory::~MockClientSocketFactory() {}
 
 void MockClientSocketFactory::AddSocketDataProvider(
     SocketDataProvider* data) {
   mock_data_.Add(data);
 }
 
 void MockClientSocketFactory::AddSSLSocketDataProvider(
     SSLSocketDataProvider* data) {
@@ skipping 23 matching lines @@
 
 std::unique_ptr<StreamSocket>
 MockClientSocketFactory::CreateTransportClientSocket(
     const AddressList& addresses,
     std::unique_ptr<SocketPerformanceWatcher> socket_performance_watcher,
     NetLog* net_log,
     const NetLogSource& source) {
   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);
 }
 
 std::unique_ptr<SSLClientSocket> MockClientSocketFactory::CreateSSLClientSocket(
     std::unique_ptr<ClientSocketHandle> transport_socket,
     const HostPortPair& host_and_port,
     const SSLConfig& ssl_config,
     const SSLClientSocketContext& context) {
   SSLSocketDataProvider* next_ssl_data = mock_ssl_data_.GetNext();
   if (!next_ssl_data->next_protos_expected_in_ssl_config.empty()) {
@@ skipping 119 matching lines @@
                                          SocketDataProvider* data)
     : MockClientSocket(NetLogWithSource::Make(net_log, NetLogSourceType::NONE)),
       addresses_(addresses),
       data_(data),
       read_offset_(0),
       read_data_(SYNCHRONOUS, ERR_UNEXPECTED),
       need_read_data_(true),
       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);
 }
 
 MockTCPClientSocket::~MockTCPClientSocket() {
   if (data_)
     data_->DetachSocket();
 }
 
 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;
+  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;
+}
 
-  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 MockTCPClientSocket::ReadIfReady(IOBuffer* buf,
+                                     int buf_len,
+                                     const CompletionCallback& callback) {
+  DCHECK(!pending_read_if_ready_callback_);
 
-  return CompleteRead();
+  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,
                                const CompletionCallback& callback) {
   DCHECK(buf);
   DCHECK_GT(buf_len, 0);
 
   if (!connected_ || !data_)
     return ERR_UNEXPECTED;
 
@@ skipping 98 matching lines @@
 bool MockTCPClientSocket::GetSSLInfo(SSLInfo* ssl_info) {
   return false;
 }
 
 void MockTCPClientSocket::OnReadComplete(const MockRead& data) {
   // If |data_| has been destroyed, safest to just do nothing.
   if (!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.
   DCHECK(need_read_data_);
 
   read_data_ = data;
   need_read_data_ = false;
 
   // 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) {
   // If |data_| has been destroyed, safest to just do nothing.
   if (!data_)
     return;
 
   // There must be a read pending.
   DCHECK(!pending_write_callback_.is_null());
   CompletionCallback callback = pending_write_callback_;
   RunCallback(callback, rv);
 }
 
 void MockTCPClientSocket::OnConnectComplete(const MockConnect& data) {
   // If |data_| has been destroyed, safest to just do nothing.
   if (!data_)
     return;
 
   CompletionCallback callback = pending_connect_callback_;
   RunCallback(callback, data.result);
 }
 
 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);
+  DCHECK_LT(0, pending_read_buf_len_);
+
+  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;
+  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;
 
-  // 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;
-  pending_read_buf_len_ = 0;
-  pending_read_callback_.Reset();
-
   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) {
       result = std::min(buf_len, read_data_.data_len - read_offset_);

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

If you've got some data in here and read_data_.mode is ASYNC, it seems like this
will consume data when it's not supposed to.

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

Done. Ah, good catch!

       memcpy(buf->data(), read_data_.data + read_offset_, result);
       read_offset_ += result;
       if (read_offset_ == read_data_.data_len) {
         need_read_data_ = true;
         read_offset_ = 0;
       }
     } else {
       result = 0;  // EOF
     }
   }
 
   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,
     const CompletionCallback& callback,
     int rv) {
   if (rv == OK)
     ssl_client_socket->connected_ = true;
   callback.Run(rv);
 }
 
@@ skipping 15 matching lines @@
 
 MockSSLClientSocket::~MockSSLClientSocket() {
   Disconnect();
 }
 
 int MockSSLClientSocket::Read(IOBuffer* buf, int buf_len,
                               const CompletionCallback& callback) {
   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);
 }
 
 int MockSSLClientSocket::Connect(const CompletionCallback& callback) {
   int rv = transport_->socket()->Connect(
       base::Bind(&ConnectCallback, base::Unretained(this), callback));
   if (rv == OK) {
     if (data_->connect.result == OK)
@@ skipping 567 matching lines @@
 }
 
 int64_t CountWriteBytes(const MockWrite writes[], size_t writes_size) {
   int64_t total = 0;
   for (const MockWrite* write = writes; write != writes + writes_size; ++write)
     total += write->data_len;
   return total;
 }
 
 }  // namespace net