Removed cached reference to BrokerRegistrar.

BUG=none
TEST=none

Review URL: http://codereview.chromium.org/5349004

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@67432 0039d316-1c4b-4281-b951-d872f2087c98

7 files changed, 91 insertions, 62 deletions

diff --git a/ceee/ie/broker/broker_rpc_client.cc b/ceee/ie/broker/broker_rpc_client.cc
index 54a308b..3bf360e 100644
--- a/ceee/ie/broker/broker_rpc_client.cc
+++ b/ceee/ie/broker/broker_rpc_client.cc
@@ -9,6 +9,7 @@
 #include <atlbase.h>
 #include "base/lock.h"
 #include "base/logging.h"
+#include "base/win/scoped_comptr.h"
 #include "broker_lib.h"  // NOLINT
 #include "broker_rpc_lib.h"  // NOLINT
 #include "ceee/common/com_utils.h"
@@ -57,19 +58,21 @@ void BrokerRpcClient::ReleaseContext() {
   } RpcEndExcept
 }
 
-HRESULT BrokerRpcClient::StartServer() {
-  // TODO(vitalybuka@google.com): Start broker without COM after the last
-  // COM interface is removed.
-  CComPtr<ICeeeBrokerRegistrar> broker;
-  HRESULT hr = broker.CoCreateInstance(CLSID_CeeeBroker);
-  LOG_IF(ERROR, FAILED(hr)) << "Failed to create broker. " << com::LogHr(hr);
-  return hr;
-}
-
-HRESULT BrokerRpcClient::Connect() {
+HRESULT BrokerRpcClient::Connect(bool start_server) {
   if (is_connected())
     return S_OK;
 
+  // Keep alive until RPC is connected.
+  base::win::ScopedComPtr<ICeeeBrokerRegistrar> broker;
+  if (start_server) {
+    // TODO(vitalybuka@google.com): Start broker without COM after the last
+    // COM interface is removed.
+    HRESULT hr = broker.CreateInstance(CLSID_CeeeBroker);
+    LOG_IF(ERROR, FAILED(hr)) << "Failed to create broker. " << com::LogHr(hr);
+    if (FAILED(hr))
+      return hr;
+  }
+
   std::wstring end_point = GetRpcEndPointAddress();
   std::wstring protocol = kRpcProtocol;
   DCHECK(!protocol.empty());
@@ -130,7 +133,7 @@ HRESULT BrokerRpcClient::FireEvent(const char* event_name,
   } RpcExcept(HandleRpcException(RpcExceptionCode())) {
     LogRpcException("RPC error in FireEvent", RpcExceptionCode());
   } RpcEndExcept
-  return RPC_E_FAULT;
+    return RPC_E_FAULT;
 }
 
 bool BrokerRpcClient::SendUmaHistogramTimes(BSTR event_name, int sample) {
@@ -147,7 +150,7 @@ bool BrokerRpcClient::SendUmaHistogramData(BSTR event_name, int sample,
                                            int min, int max, int bucket_count) {
   RpcTryExcept {
     BrokerRpcClient_SendUmaHistogramData(binding_handle_, event_name, sample,
-        min, max, bucket_count);
+                                         min, max, bucket_count);
     return true;
   } RpcExcept(HandleRpcException(RpcExceptionCode())) {
     LogRpcException("RPC error in SendUmaHistogramData", RpcExceptionCode());
diff --git a/ceee/ie/broker/broker_rpc_client.h b/ceee/ie/broker/broker_rpc_client.h
index fc38c0d..0f96e0e 100644
--- a/ceee/ie/broker/broker_rpc_client.h
+++ b/ceee/ie/broker/broker_rpc_client.h
@@ -17,7 +17,7 @@ class BrokerRpcClient {
   virtual ~BrokerRpcClient();
 
   // Initialize connection with server.
-  virtual HRESULT Connect();
+  virtual HRESULT Connect(bool start_server);
 
   // Relese connection with server
   virtual void Disconnect();
@@ -41,9 +41,6 @@ class BrokerRpcClient {
                                     int bucket_count);
   // @}
 
-  // Starts new CEEE broker if nessesary.
-  static HRESULT StartServer();
-
  private:
   void LockContext();
   void ReleaseContext();
diff --git a/ceee/ie/broker/broker_rpc_unittest.cc b/ceee/ie/broker/broker_rpc_unittest.cc
index dd75832..ba24a56 100644
--- a/ceee/ie/broker/broker_rpc_unittest.cc
+++ b/ceee/ie/broker/broker_rpc_unittest.cc
@@ -48,7 +48,7 @@ class BrokerRpcTest : public testing::Test {
 TEST_F(BrokerRpcTest, ConnectNoServer) {
   BrokerRpcClient client;
   ASSERT_FALSE(client.is_connected());
-  ASSERT_FALSE(SUCCEEDED(client.Connect()));
+  ASSERT_FALSE(SUCCEEDED(client.Connect(false)));
   ASSERT_FALSE(client.is_connected());
 }
 
@@ -58,7 +58,7 @@ TEST_F(BrokerRpcTest, Connect) {
   ASSERT_TRUE(server.Start());
   ASSERT_TRUE(server.is_started());
   BrokerRpcClient client;
-  ASSERT_TRUE(SUCCEEDED(client.Connect()));
+  ASSERT_TRUE(SUCCEEDED(client.Connect(false)));
   ASSERT_TRUE(client.is_connected());
 }
 
@@ -67,14 +67,14 @@ TEST_F(BrokerRpcTest, FireEvent) {
   ASSERT_TRUE(server.Start());
 
   BrokerRpcClient client;
-  ASSERT_TRUE(SUCCEEDED(client.Connect()));
+  ASSERT_TRUE(SUCCEEDED(client.Connect(false)));
 
   const char* event_name = "event_name";
   const char* event_args = "event_args";
 
   EXPECT_CALL(broker_rpc_mock_, BrokerRpcServer_FireEvent(_, StrEq(event_name),
                                                           StrEq(event_args)))
-      .Times(1);
+          .Times(1);
 
   ASSERT_TRUE(SUCCEEDED(client.FireEvent(event_name, event_args)));
 }
diff --git a/ceee/ie/plugin/bho/browser_helper_object.cc b/ceee/ie/plugin/bho/browser_helper_object.cc
index e926adc..09c40bc 100644
--- a/ceee/ie/plugin/bho/browser_helper_object.cc
+++ b/ceee/ie/plugin/bho/browser_helper_object.cc
@@ -213,12 +213,17 @@ HRESULT BrowserHelperObject::GetParentBrowser(IWebBrowser2* browser,
   return S_OK;
 }
 
-HRESULT BrowserHelperObject::GetBrokerRegistrar(
-    ICeeeBrokerRegistrar** broker) {
-  // This is a singleton and will not create a new one.
-  return ::CoCreateInstance(CLSID_CeeeBroker, NULL, CLSCTX_ALL,
-                            IID_ICeeeBrokerRegistrar,
-                            reinterpret_cast<void**>(broker));
+HRESULT BrowserHelperObject::GetBrokerRegistrar(ICeeeBrokerRegistrar** broker) {
+  DCHECK(broker);
+  if (broker == NULL)
+    return E_INVALIDARG;
+  HRESULT hr = ::CoCreateInstance(CLSID_CeeeBroker, NULL, CLSCTX_ALL,
+                                  IID_ICeeeBrokerRegistrar,
+                                  reinterpret_cast<void**>(broker));
+  DCHECK(SUCCEEDED(hr) && *broker) << "CoCreating Broker. " << com::LogHr(hr);
+  if (FAILED(hr) || *broker == NULL)
+    return com::AlwaysError(hr);
+  return hr;
 }
 
 HRESULT BrowserHelperObject::CreateExecutor(IUnknown** executor) {
@@ -251,23 +256,53 @@ WebProgressNotifier* BrowserHelperObject::CreateWebProgressNotifier() {
   return SUCCEEDED(hr) ? web_progress_notifier.release() : NULL;
 }
 
+HRESULT BrowserHelperObject::RegisterTabExecutor(DWORD thread_id,
+                                                 IUnknown* executor) {
+  base::win::ScopedComPtr<ICeeeBrokerRegistrar> registrar;
+  HRESULT hr = GetBrokerRegistrar(registrar.Receive());
+  if (FAILED(hr))
+    return hr;
+  // TODO(mad@chromium.org): Implement the proper manual/secure
+  // registration.
+  return registrar->RegisterTabExecutor(thread_id, executor);
+}
+
+HRESULT BrowserHelperObject::UnregisterExecutor(DWORD thread_id) {
+  base::win::ScopedComPtr<ICeeeBrokerRegistrar> registrar;
+  HRESULT hr = GetBrokerRegistrar(registrar.Receive());
+  if (FAILED(hr))
+    return hr;
+  return registrar->UnregisterExecutor(thread_id);
+}
+
+HRESULT BrowserHelperObject::SetTabToolBandIdForHandle(int tool_band_id,
+                                                       long tab_handle) {
+  base::win::ScopedComPtr<ICeeeBrokerRegistrar> registrar;
+  HRESULT hr = GetBrokerRegistrar(registrar.Receive());
+  if (FAILED(hr))
+    return hr;
+  return registrar->SetTabToolBandIdForHandle(tool_band_id, tab_handle);
+}
+
+HRESULT BrowserHelperObject::SetTabIdForHandle(int tool_band_id,
+                                               long tab_handle) {
+  base::win::ScopedComPtr<ICeeeBrokerRegistrar> registrar;
+  HRESULT hr = GetBrokerRegistrar(registrar.Receive());
+  if (FAILED(hr))
+    return hr;
+  return registrar->SetTabIdForHandle(tool_band_id, tab_handle);
+}
+
 HRESULT BrowserHelperObject::Initialize(IUnknown* site) {
   TRACE_EVENT_INSTANT("ceee.bho.initialize", this, "");
 
-  // We need to make sure the Broker is CoCreated BEFORE we try RPC to it.
-  HRESULT hr = GetBrokerRegistrar(broker_registrar_.Receive());
-  DCHECK(SUCCEEDED(hr) && broker_registrar_) << "CoCreating Broker. " <<
-      com::LogHr(hr);
-  if (FAILED(hr) || !broker_registrar_)
-    return com::AlwaysError(hr);
-
   // Unfortunately, we need to connect before taking performance. Including
   // this call in our Timing would be useful. Unit tests make it
   // complicated.
   // TODO(hansl@chromium.org): Change initialization to be able to time
   // this call too.
   DCHECK(broker_rpc_ == NULL);
-  hr = ConnectRpcBrokerClient();
+  HRESULT hr = ConnectRpcBrokerClient();
   if (FAILED(hr) || !broker_rpc_->is_connected()) {
     NOTREACHED() << "Couldn't connect to the RPC server.";
     return com::AlwaysError(hr);
@@ -294,19 +329,16 @@ HRESULT BrowserHelperObject::Initialize(IUnknown* site) {
 
   // Preemptively feed the broker with an executor in our thread.
   DCHECK(executor_ == NULL);
-  hr = CreateExecutor(executor_.Receive());
-  LOG_IF(ERROR, FAILED(hr)) << "Failed to create Executor, hr=" <<
-      com::LogHr(hr);
+  base::win::ScopedComPtr<IUnknown> executor;
+  hr = CreateExecutor(executor.Receive());
   DCHECK(SUCCEEDED(hr)) << "CoCreating Executor. " << com::LogHr(hr);
-  if (SUCCEEDED(hr)) {
-    // TODO(mad@chromium.org): Implement the proper manual/secure
-    // registration.
-    hr = broker_registrar_->RegisterTabExecutor(::GetCurrentThreadId(),
-                                                executor_);
-    DCHECK(SUCCEEDED(hr)) << "Registering Executor. " << com::LogHr(hr);
-  }
   if (FAILED(hr))
     return hr;
+  hr = RegisterTabExecutor(::GetCurrentThreadId(), executor);
+  DCHECK(SUCCEEDED(hr)) << "Registering Executor. " << com::LogHr(hr);
+  if (FAILED(hr))
+    return hr;
+  executor_ = executor;
 
   // We need the service provider for both the sink connections and
   // to get a handle to the tab window.
@@ -365,21 +397,18 @@ HRESULT BrowserHelperObject::TearDown() {
 
   // Unregister our executor so that the broker don't have to rely
   // on the thread dying to release it and confuse COM in trying to release it.
-  if (broker_registrar_ != NULL) {
-    DCHECK(executor_ != NULL);
+  if (executor_ != NULL) {
     // Manually disconnect executor_,
     // so it doesn't get called while we unregister it below.
     HRESULT hr = ::CoDisconnectObject(executor_, 0);
+    executor_.Release();
     DCHECK(SUCCEEDED(hr)) << "Couldn't disconnect Executor. " << com::LogHr(hr);
 
     // TODO(mad@chromium.org): Implement the proper manual/secure
     // unregistration.
-    hr = broker_registrar_->UnregisterExecutor(::GetCurrentThreadId());
+    hr = UnregisterExecutor(::GetCurrentThreadId());
     DCHECK(SUCCEEDED(hr)) << "Unregistering Executor. " << com::LogHr(hr);
-  } else {
-    DCHECK(executor_ == NULL);
   }
-  executor_.Release();
 
   if (web_browser_)
     DispEventUnadvise(web_browser_, &DIID_DWebBrowserEvents2);
@@ -570,7 +599,7 @@ bool BrowserHelperObject::EnsureTabId() {
   }
 
   CeeeWindowHandle handle = reinterpret_cast<CeeeWindowHandle>(tab_window_);
-  hr = broker_registrar_->SetTabIdForHandle(tab_id_, handle);
+  hr = SetTabIdForHandle(tab_id_, handle);
   if (FAILED(hr)) {
     DCHECK(SUCCEEDED(hr)) << "An error occured when setting the tab_id: " <<
         com::LogHr(hr);
@@ -1192,7 +1221,7 @@ HRESULT BrowserHelperObject::CreateFrameEventHandler(
 
 HRESULT BrowserHelperObject::ConnectRpcBrokerClient() {
   broker_rpc_.reset(new BrokerRpcClient);
-  return broker_rpc_->Connect();
+  return broker_rpc_->Connect(true);
 }
 
 HRESULT BrowserHelperObject::AttachBrowser(IWebBrowser2* browser,
@@ -1553,8 +1582,7 @@ void BrowserHelperObject::UnregisterSink(Sink* sink) {
 
 STDMETHODIMP BrowserHelperObject::SetToolBandSessionId(long session_id) {
   CeeeWindowHandle handle = reinterpret_cast<CeeeWindowHandle>(tab_window_);
-  HRESULT hr = broker_registrar_->SetTabToolBandIdForHandle(session_id,
-                                                            handle);
+  HRESULT hr = SetTabToolBandIdForHandle(session_id, handle);
   if (FAILED(hr)) {
     DCHECK(SUCCEEDED(hr)) << "An error occured when setting the toolband " <<
         "tab ID: " << com::LogHr(hr);
diff --git a/ceee/ie/plugin/bho/browser_helper_object.h b/ceee/ie/plugin/bho/browser_helper_object.h
index 8433ca7..8fc6584 100644
--- a/ceee/ie/plugin/bho/browser_helper_object.h
+++ b/ceee/ie/plugin/bho/browser_helper_object.h
@@ -236,6 +236,7 @@ class ATL_NO_VTABLE BrowserHelperObject
                                    IWebBrowser2** parent_browser);
 
   // Unit testing seam to create the broker registrar.
+  // Release registrar ASAP. Cached reference may block broker from release.
   virtual HRESULT GetBrokerRegistrar(ICeeeBrokerRegistrar** broker);
 
   // Unit testing seam to create an executor.
@@ -314,6 +315,12 @@ class ATL_NO_VTABLE BrowserHelperObject
   virtual HRESULT AttachBrowserHandler(IWebBrowser2* webbrowser,
                                        IFrameEventHandler** handler);
 
+  // Helpers to call function in registrar with following release of broker.
+  HRESULT RegisterTabExecutor(DWORD thread_id, IUnknown* executor);
+  HRESULT UnregisterExecutor(DWORD thread_id);
+  HRESULT SetTabIdForHandle(int tool_band_id, long tab_handle);
+  HRESULT SetTabToolBandIdForHandle(int tool_band_id, long tab_handle);
+
   // Function info objects describing our message handlers.
   // Effectively const but can't make const because of silly ATL macro problem.
   static _ATL_FUNC_INFO handler_type_idispatch_5variantptr_boolptr_;
@@ -328,9 +335,6 @@ class ATL_NO_VTABLE BrowserHelperObject
   // The Chrome Frame host handling a Chrome Frame instance for us.
   ScopedChromeFramePtr chrome_frame_host_;
 
-  // The Broker Registrar we use to un/register executors for our thread.
-  base::win::ScopedComPtr<ICeeeBrokerRegistrar> broker_registrar_;
-
   // We keep a reference to the executor we registered so that we can
   // manually disconnect it, so it doesn't get called while we unregister it.
   ScopedIUnkPtr executor_;
diff --git a/ceee/ie/plugin/bho/events_funnel.cc b/ceee/ie/plugin/bho/events_funnel.cc
index 294ea4d..c73d1f2 100644
--- a/ceee/ie/plugin/bho/events_funnel.cc
+++ b/ceee/ie/plugin/bho/events_funnel.cc
@@ -42,10 +42,7 @@ HRESULT EventsFunnel::SendEventToBroker(const char* event_name,
     broker_rpc_client_.reset(new BrokerRpcClient());
     if (!broker_rpc_client_.get())
       return E_OUTOFMEMORY;
-    HRESULT hr = BrokerRpcClient::StartServer();
-    if (FAILED(hr))
-      return hr;
-    hr = broker_rpc_client_->Connect();
+    HRESULT hr = broker_rpc_client_->Connect(true);
     // Don't reset broker_rpc_client_. See comment in *h file.
     if (FAILED(hr))
       return hr;
diff --git a/ceee/ie/testing/mock_broker_and_friends.h b/ceee/ie/testing/mock_broker_and_friends.h
index 9efe3c0..3b52b74 100644
--- a/ceee/ie/testing/mock_broker_and_friends.h
+++ b/ceee/ie/testing/mock_broker_and_friends.h
@@ -391,7 +391,7 @@ class MockWebRequestEventsFunnel : public WebRequestEventsFunnel {
 
 class MockBrokerRpcClient : public BrokerRpcClient {
  public:
-  MOCK_METHOD0(Connect, HRESULT());
+  MOCK_METHOD1(Connect, HRESULT(bool));
   MOCK_METHOD0(Disconnect, void());
   MOCK_CONST_METHOD0(is_connected, bool());
   MOCK_METHOD2(FireEvent, HRESULT(const char*, const char*));