services/shell: Fix running the connection-lost callback.

services/shell/public/cpp/lib/shell_connection.cc

 // Copyright 2014 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 "services/shell/public/cpp/shell_connection.h"
 
 #include "base/bind.h"
 #include "base/logging.h"
 #include "base/message_loop/message_loop.h"
 #include "mojo/public/cpp/bindings/interface_ptr.h"
@@ skipping 25 matching lines @@
 void ShellConnection::set_initialize_handler(const base::Closure& callback) {
   initialize_handler_ = callback;
 }
 
 void ShellConnection::SetAppTestConnectorForTesting(
     mojom::ConnectorPtr connector) {
   pending_connector_request_ = nullptr;
   connector_.reset(new ConnectorImpl(std::move(connector)));
 }
 
+void ShellConnection::SetConnectionLostClosure(const base::Closure& closure) {
+  connection_lost_closure_ = closure;
+  if (should_run_connection_lost_closure_ &&
+      !connection_lost_closure_.is_null())
+    connection_lost_closure_.Run();
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // ShellConnection, mojom::ShellClient implementation:
 
 void ShellConnection::Initialize(mojom::IdentityPtr identity,
                                  uint32_t id,
                                  const InitializeCallback& callback) {
   identity_ = identity.To<Identity>();
   if (!initialize_handler_.is_null())
     initialize_handler_.Run();
 
@@ skipping 24 matching lines @@
   incoming_connections_.push_back(std::move(registry));
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // ShellConnection, private:
 
 void ShellConnection::OnConnectionError() {
   // Note that the ShellClient doesn't technically have to quit now, it may live
   // on to service existing connections. All existing Connectors however are
   // invalid.
-  if (client_->ShellConnectionLost() && !connection_lost_closure_.is_null())
+  should_run_connection_lost_closure_ = client_->ShellConnectionLost();
+  if (should_run_connection_lost_closure_ &&
+      !connection_lost_closure_.is_null())
     connection_lost_closure_.Run();
   // We don't reset the connector as clients may have taken a raw pointer to it.
   // Connect() will return nullptr if they try to connect to anything.
 }
 
 }  // namespace shell