Add FlushAsyncForTesting to InterfacePtr API

This will be used in
URLLoaderFactoryGetter::FlushNetworkInterfaceForTesting to prevent
nested IO message loops.

Hit this in
https://chromium-review.googlesource.com/c/chromium/src/+/1139048.

Bug: 857577
Change-Id: I1b6a53f39737ff0a6d4681b4256991569b91a2d5
Reviewed-on: https://chromium-review.googlesource.com/1142198
Commit-Queue: Clark DuVall <cduvall@chromium.org>
Reviewed-by: Ken Rockot <rockot@chromium.org>
Cr-Commit-Position: refs/heads/master@{#576350}

mojo/public/cpp/bindings/interface_endpoint_client.h

@@ -113,6 +113,7 @@ class MOJO_CPP_BINDINGS_EXPORT InterfaceEndpointClient
   void QueryVersion(const base::Callback<void(uint32_t)>& callback);
   void RequireVersion(uint32_t version);
   void FlushForTesting();
+  void FlushAsyncForTesting(base::OnceClosure callback);
 
  private:
   // Maps from the id of a response to the MessageReceiver that handles the

mojo/public/cpp/bindings/interface_ptr.h

@@ -127,6 +127,12 @@ class InterfacePtr {
   // stimulus.
   void FlushForTesting() { internal_state_.FlushForTesting(); }
 
+  // Same as |FlushForTesting()| but will call |callback| when the flush is
+  // complete.
+  void FlushAsyncForTesting(base::OnceClosure callback) {
+    internal_state_.FlushAsyncForTesting(std::move(callback));
+  }
+
   // Closes the bound message pipe, if any.
   void reset() {
     State doomed;

mojo/public/cpp/bindings/lib/control_message_proxy.cc

@@ -12,6 +12,7 @@
 #include "base/callback_helpers.h"
 #include "base/macros.h"
 #include "base/run_loop.h"
+#include "base/threading/sequenced_task_runner_handle.h"
 #include "mojo/public/cpp/bindings/lib/serialization.h"
 #include "mojo/public/cpp/bindings/lib/validation_util.h"
 #include "mojo/public/cpp/bindings/message.h"
@@ -39,12 +40,12 @@ bool ValidateControlResponse(Message* message) {
 }
 
 using RunCallback =
-    base::Callback<void(interface_control::RunResponseMessageParamsPtr)>;
+    base::OnceCallback<void(interface_control::RunResponseMessageParamsPtr)>;
 
 class RunResponseForwardToCallback : public MessageReceiver {
  public:
-  explicit RunResponseForwardToCallback(const RunCallback& callback)
-      : callback_(callback) {}
+  explicit RunResponseForwardToCallback(RunCallback callback)
+      : callback_(std::move(callback)) {}
   bool Accept(Message* message) override;
 
  private:
@@ -65,13 +66,13 @@ bool RunResponseForwardToCallback::Accept(Message* message) {
   Deserialize<interface_control::RunResponseMessageParamsDataView>(
       params, &params_ptr, &context);
 
-  callback_.Run(std::move(params_ptr));
+  std::move(callback_).Run(std::move(params_ptr));
   return true;
 }
 
 void SendRunMessage(MessageReceiverWithResponder* receiver,
                     interface_control::RunInputPtr input_ptr,
-                    const RunCallback& callback) {
+                    RunCallback callback) {
   auto params_ptr = interface_control::RunMessageParams::New();
   params_ptr->input = std::move(input_ptr);
   Message message(interface_control::kRunMessageId,
@@ -81,7 +82,7 @@ void SendRunMessage(MessageReceiverWithResponder* receiver,
   Serialize<interface_control::RunMessageParamsDataView>(
       params_ptr, message.payload_buffer(), &params, &context);
   std::unique_ptr<MessageReceiver> responder =
-      std::make_unique<RunResponseForwardToCallback>(callback);
+      std::make_unique<RunResponseForwardToCallback>(std::move(callback));
   ignore_result(receiver->AcceptWithResponder(&message, std::move(responder)));
 }
 
@@ -115,9 +116,9 @@ void RunVersionCallback(
   callback.Run(version);
 }
 
-void RunClosure(const base::Closure& callback,
+void RunClosure(base::OnceClosure callback,
                 interface_control::RunResponseMessageParamsPtr run_response) {
-  callback.Run();
+  std::move(callback).Run();
 }
 
 }  // namespace
@@ -133,7 +134,7 @@ void ControlMessageProxy::QueryVersion(
   auto input_ptr = interface_control::RunInput::New();
   input_ptr->set_query_version(interface_control::QueryVersion::New());
   SendRunMessage(receiver_, std::move(input_ptr),
-                 base::Bind(&RunVersionCallback, callback));
+                 base::BindOnce(&RunVersionCallback, callback));
 }
 
 void ControlMessageProxy::RequireVersion(uint32_t version) {
@@ -145,29 +146,38 @@ void ControlMessageProxy::RequireVersion(uint32_t version) {
 }
 
 void ControlMessageProxy::FlushForTesting() {
-  if (encountered_error_)
+  base::RunLoop run_loop(base::RunLoop::Type::kNestableTasksAllowed);
+  FlushAsyncForTesting(run_loop.QuitClosure());
+  run_loop.Run();
+}
+
+void ControlMessageProxy::FlushAsyncForTesting(base::OnceClosure callback) {
+  if (encountered_error_) {
+    base::SequencedTaskRunnerHandle::Get()->PostTask(FROM_HERE,
+                                                     std::move(callback));
     return;
+  }
 
   auto input_ptr = interface_control::RunInput::New();
   input_ptr->set_flush_for_testing(interface_control::FlushForTesting::New());
-  base::RunLoop run_loop(base::RunLoop::Type::kNestableTasksAllowed);
-  run_loop_quit_closure_ = run_loop.QuitClosure();
+  DCHECK(!pending_flush_callback_);
+  pending_flush_callback_ = std::move(callback);
   SendRunMessage(
       receiver_, std::move(input_ptr),
-      base::Bind(&RunClosure,
-                 base::Bind(&ControlMessageProxy::RunFlushForTestingClosure,
+      base::BindOnce(
+          &RunClosure,
+          base::BindOnce(&ControlMessageProxy::RunFlushForTestingClosure,
                          base::Unretained(this))));
-  run_loop.Run();
 }
 
 void ControlMessageProxy::RunFlushForTestingClosure() {
-  DCHECK(!run_loop_quit_closure_.is_null());
-  base::ResetAndReturn(&run_loop_quit_closure_).Run();
+  DCHECK(!pending_flush_callback_.is_null());
+  std::move(pending_flush_callback_).Run();
 }
 
 void ControlMessageProxy::OnConnectionError() {
   encountered_error_ = true;
-  if (!run_loop_quit_closure_.is_null())
+  if (!pending_flush_callback_.is_null())
     RunFlushForTestingClosure();
 }
 

mojo/public/cpp/bindings/lib/control_message_proxy.h

@@ -29,6 +29,7 @@ class MOJO_CPP_BINDINGS_EXPORT ControlMessageProxy {
   void RequireVersion(uint32_t version);
 
   void FlushForTesting();
+  void FlushAsyncForTesting(base::OnceClosure callback);
   void OnConnectionError();
 
  private:
@@ -38,7 +39,7 @@ class MOJO_CPP_BINDINGS_EXPORT ControlMessageProxy {
   MessageReceiverWithResponder* receiver_;
   bool encountered_error_ = false;
 
-  base::Closure run_loop_quit_closure_;
+  base::OnceClosure pending_flush_callback_;
 
   DISALLOW_COPY_AND_ASSIGN(ControlMessageProxy);
 };

mojo/public/cpp/bindings/lib/interface_endpoint_client.cc

@@ -347,6 +347,10 @@ void InterfaceEndpointClient::FlushForTesting() {
   control_message_proxy_.FlushForTesting();
 }
 
+void InterfaceEndpointClient::FlushAsyncForTesting(base::OnceClosure callback) {
+  control_message_proxy_.FlushAsyncForTesting(std::move(callback));
+}
+
 void InterfaceEndpointClient::InitControllerIfNecessary() {
   if (controller_ || handle_.pending_association())
     return;

mojo/public/cpp/bindings/lib/interface_ptr_state.h

@@ -131,6 +131,11 @@ class InterfacePtrState : public InterfacePtrStateBase {
     endpoint_client()->FlushForTesting();
   }
 
+  void FlushAsyncForTesting(base::OnceClosure callback) {
+    ConfigureProxyIfNecessary();
+    endpoint_client()->FlushAsyncForTesting(std::move(callback));
+  }
+
   void CloseWithReason(uint32_t custom_reason, const std::string& description) {
     ConfigureProxyIfNecessary();
     endpoint_client()->CloseWithReason(custom_reason, description);

mojo/public/cpp/bindings/tests/interface_ptr_unittest.cc

@@ -754,6 +754,27 @@ TEST_P(InterfacePtrTest, FlushForTesting) {
   EXPECT_EQ(10.0, calculator_ui.GetOutput());
 }
 
+TEST_P(InterfacePtrTest, FlushAsyncForTesting) {
+  math::CalculatorPtr calc;
+  MathCalculatorImpl calc_impl(MakeRequest(&calc));
+  calc.set_connection_error_handler(base::BindOnce(&Fail));
+
+  MathCalculatorUI calculator_ui(std::move(calc));
+
+  calculator_ui.Add(2.0, base::DoNothing());
+  base::RunLoop run_loop;
+  calculator_ui.GetInterfacePtr().FlushAsyncForTesting(run_loop.QuitClosure());
+  run_loop.Run();
+  EXPECT_EQ(2.0, calculator_ui.GetOutput());
+
+  calculator_ui.Multiply(5.0, base::DoNothing());
+  base::RunLoop run_loop2;
+  calculator_ui.GetInterfacePtr().FlushAsyncForTesting(run_loop2.QuitClosure());
+  run_loop2.Run();
+
+  EXPECT_EQ(10.0, calculator_ui.GetOutput());
+}
+
 void SetBool(bool* value) {
   *value = true;
 }
@@ -768,6 +789,21 @@ TEST_P(InterfacePtrTest, FlushForTestingWithClosedPeer) {
   calc.FlushForTesting();
 }
 
+TEST_P(InterfacePtrTest, FlushAsyncForTestingWithClosedPeer) {
+  math::CalculatorPtr calc;
+  MakeRequest(&calc);
+  bool called = false;
+  calc.set_connection_error_handler(base::BindOnce(&SetBool, &called));
+  base::RunLoop run_loop;
+  calc.FlushAsyncForTesting(run_loop.QuitClosure());
+  run_loop.Run();
+  EXPECT_TRUE(called);
+
+  base::RunLoop run_loop2;
+  calc.FlushAsyncForTesting(run_loop2.QuitClosure());
+  run_loop2.Run();
+}
+
 TEST_P(InterfacePtrTest, ConnectionErrorWithReason) {
   math::CalculatorPtr calc;
   MathCalculatorImpl calc_impl(MakeRequest(&calc));