[Media Router] Clean up MediaSinkServiceBase.

MediaSinkServiceBase changes:

- Move the Observer interface currently in CMSSImpl to
MediaSinkServiceBase and generalizes it so it can be used for other
MediaSinkServiceBase implementations.
- Standardizes add/update/remove/get sink operations in
MediaSinkServiceBase.
Both of these make the code using MediaSinkServiceBase more testable as
additional code in MRP depend on them.

Also cleaned up logic related to the discovery timer by combining
Start/RestartTimer().

-----

DialMediaSinkSevice changes:
- DMSSImpl now pools sinks discovered in the latest round and is used
to merge with the |MSSBase::sinks_| when the discovery timer fires.
Before this patch, the current set is removed on the start of a discovery
round (because unlike Cast, there is no active connection to a device to
tell that it is still alive or gone), which makes the sink list unstable.
With this patch, the current set is only modified at the end of a discovery
round.
- TODOs in future patches:
-- move app discovery service out of DMSSImpl, similar to Cast.
-- get rid of DMSS and move UI thread dependencies (if any) into part
   of DMSSImpl. Then rename DMSSImpl -> DMSS.

-----

CastMediaSinkService changes;
- No longer keep track of sinks by ip endpoint. Instead uses
|MSSBase::sinks_| to maintain sink list.
- CastAppDiscoveryService no longer need to maintain its own map of
sinks, and instead uses |MSSBase::GetSinks()/GetSinkById()|.
- The callback used for dual disovery is converted to use
MediaSinkServiceBase::Observer.
- TODOs in future patches:
-- get rid of CMSS and move UI thread dependencies into CMSSImpl.
Rename CMSSImpl -> CMSS.
-----

Notes:
- The MediaRouteDesktop ProvideSinks callback is not converted to use
Observers. This is because the callback is expected to be temporary (until
we can turn off the extension MRPs), so we treat it as a special case.

Bug: 816628,698940
Change-Id: I11fd13f7ebf743d78c9499ad24f98cf72f571ee5
Reviewed-on: https://chromium-review.googlesource.com/1033727
Commit-Queue: Derek Cheng <imcheng@chromium.org>
Reviewed-by: Brandon Tolsch <btolsch@chromium.org>
Cr-Commit-Position: refs/heads/master@{#555854}

chrome/browser/media/router/discovery/dial/dial_media_sink_service.cc

@@ -23,8 +23,7 @@ DialMediaSinkService::~DialMediaSinkService() {
 }
 
 void DialMediaSinkService::Start(
-    const OnSinksDiscoveredCallback& sink_discovery_cb,
-    const OnDialSinkAddedCallback& dial_sink_added_cb) {
+    const OnSinksDiscoveredCallback& sink_discovery_cb) {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   DCHECK(!impl_);
 
@@ -34,24 +33,16 @@ void DialMediaSinkService::Start(
       base::BindRepeating(&DialMediaSinkService::RunSinksDiscoveredCallback,
                           weak_ptr_factory_.GetWeakPtr(), sink_discovery_cb));
 
-  impl_ = CreateImpl(sink_discovery_cb_impl, dial_sink_added_cb);
+  impl_ = CreateImpl(sink_discovery_cb_impl);
 
   impl_->task_runner()->PostTask(
       FROM_HERE, base::BindOnce(&DialMediaSinkServiceImpl::Start,
                                 base::Unretained(impl_.get())));
 }
 
-void DialMediaSinkService::OnUserGesture() {
-  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
-  impl_->task_runner()->PostTask(
-      FROM_HERE, base::BindOnce(&DialMediaSinkServiceImpl::OnUserGesture,
-                                base::Unretained(impl_.get())));
-}
-
 std::unique_ptr<DialMediaSinkServiceImpl, base::OnTaskRunnerDeleter>
 DialMediaSinkService::CreateImpl(
-    const OnSinksDiscoveredCallback& sink_discovery_cb,
-    const OnDialSinkAddedCallback& dial_sink_added_cb) {
+    const OnSinksDiscoveredCallback& sink_discovery_cb) {
   // Clone the connector so it can be used on the IO thread.
   std::unique_ptr<service_manager::Connector> connector =
       content::ServiceManagerConnection::GetForProcess()
@@ -65,7 +56,7 @@ DialMediaSinkService::CreateImpl(
           content::BrowserThread::IO);
   return std::unique_ptr<DialMediaSinkServiceImpl, base::OnTaskRunnerDeleter>(
       new DialMediaSinkServiceImpl(std::move(connector), sink_discovery_cb,
-                                   dial_sink_added_cb, task_runner),
+                                   task_runner),
       base::OnTaskRunnerDeleter(task_runner));
 }
 

chrome/browser/media/router/discovery/dial/dial_media_sink_service.h

@@ -29,6 +29,9 @@ using OnDialSinkAddedCallback =
 // Service to discover DIAL media sinks.  All public methods must be invoked on
 // the UI thread.  Delegates to DialMediaSinkServiceImpl by posting tasks to its
 // SequencedTaskRunner.
+// TODO(imcheng): Remove this class and moving the logic into a part
+// of DialMediaSinkServiceImpl that runs on the UI thread, and renaming
+// DialMediaSinkServiceImpl to DialMediaSinkService.
 class DialMediaSinkService {
  public:
   DialMediaSinkService();
@@ -37,19 +40,8 @@ class DialMediaSinkService {
   // Starts discovery of DIAL sinks. Can only be called once.
   // |sink_discovery_cb|: Callback to invoke on UI thread when the list of
   // discovered sinks has been updated.
-  // |dial_sink_added_cb|: Callback to invoke when a new DIAL sink has been
-  // discovered. The callback may be invoked on any thread, and may be invoked
-  // after |this| is destroyed. Can be null.
   // Marked virtual for tests.
-  virtual void Start(const OnSinksDiscoveredCallback& sink_discovery_cb,
-                     const OnDialSinkAddedCallback& dial_sink_added_cb);
-
-  // Initiates discovery immediately in response to a user gesture
-  // (i.e., opening the Media Router dialog). This method can only be called
-  // after |Start()|.
-  // TODO(imcheng): Rename to ManuallyInitiateDiscovery() or similar.
-  // Marked virtual for tests.
-  virtual void OnUserGesture();
+  virtual void Start(const OnSinksDiscoveredCallback& sink_discovery_cb);
 
   // Returns a raw pointer to |impl_|. This method is only valid to call after
   // |Start()| has been called. Always returns non-null.
@@ -61,8 +53,7 @@ class DialMediaSinkService {
  private:
   // Marked virtual for tests.
   virtual std::unique_ptr<DialMediaSinkServiceImpl, base::OnTaskRunnerDeleter>
-  CreateImpl(const OnSinksDiscoveredCallback& sink_discovery_cb,
-             const OnDialSinkAddedCallback& dial_sink_added_cb);
+  CreateImpl(const OnSinksDiscoveredCallback& sink_discovery_cb);
 
   void RunSinksDiscoveredCallback(
       const OnSinksDiscoveredCallback& sinks_discovered_cb,

chrome/browser/media/router/discovery/dial/dial_media_sink_service_impl.cc

@@ -52,11 +52,9 @@ SinkAppStatus GetSinkAppStatusFromResponse(const DialAppInfoResult& result) {
 DialMediaSinkServiceImpl::DialMediaSinkServiceImpl(
     std::unique_ptr<service_manager::Connector> connector,
     const OnSinksDiscoveredCallback& on_sinks_discovered_cb,
-    const OnDialSinkAddedCallback& dial_sink_added_cb,
     const scoped_refptr<base::SequencedTaskRunner>& task_runner)
     : MediaSinkServiceBase(on_sinks_discovered_cb),
       connector_(std::move(connector)),
-      dial_sink_added_cb_(dial_sink_added_cb),
       task_runner_(task_runner) {
   DETACH_FROM_SEQUENCE(sequence_checker_);
 }
@@ -86,7 +84,7 @@ void DialMediaSinkServiceImpl::Start() {
   app_discovery_service_ =
       std::make_unique<DialAppDiscoveryService>(connector_.get());
 
-  MediaSinkServiceBase::StartTimer();
+  StartTimer();
 
   dial_registry_ =
       test_dial_registry_ ? test_dial_registry_ : DialRegistry::GetInstance();
@@ -96,18 +94,6 @@ void DialMediaSinkServiceImpl::Start() {
 
 void DialMediaSinkServiceImpl::OnUserGesture() {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
-  // Re-sync sinks to CastMediaSinkService. It's possible that a DIAL-discovered
-  // sink was added to CastMediaSinkService earlier, but was removed due to
-  // flaky network. This gives CastMediaSinkService an opportunity to recover
-  // even if mDNS is not working for some reason.
-  DVLOG(2) << "OnUserGesture: re-syncing " << current_sinks_.size()
-           << " sinks to CastMediaSinkService";
-
-  if (dial_sink_added_cb_) {
-    for (const auto& sink_it : current_sinks_)
-      dial_sink_added_cb_.Run(sink_it.second);
-  }
-
   RescanAppInfo();
 }
 
@@ -125,8 +111,8 @@ DialMediaSinkServiceImpl::StartMonitoringAvailableSinksForApp(
         base::Unretained(this), app_name, callback_list.get()));
 
     // Start checking if |app_name| is available on existing sinks.
-    for (const auto& dial_sink_it : current_sinks_)
-      FetchAppInfoForSink(dial_sink_it.second, app_name);
+    for (const auto& sink : GetSinks())
+      FetchAppInfoForSink(sink.second, app_name);
   }
 
   return callback_list->Add(callback);
@@ -149,9 +135,35 @@ void DialMediaSinkServiceImpl::SetAppDiscoveryServiceForTest(
 }
 
 void DialMediaSinkServiceImpl::OnDiscoveryComplete() {
-  MediaSinkServiceBase::OnDiscoveryComplete();
+  std::vector<MediaSinkInternal> sinks_to_update;
+  std::vector<MediaSinkInternal> sinks_to_remove;
+  for (const auto& sink : GetSinks()) {
+    if (!base::ContainsKey(latest_sinks_, sink.first))
+      sinks_to_remove.push_back(sink.second);
+  }
+
+  for (const auto& latest_sink : latest_sinks_) {
+    // Sink is added or updated.
+    const MediaSinkInternal* sink = GetSinkById(latest_sink.first);
+    if (!sink || *sink != latest_sink.second)
+      sinks_to_update.push_back(latest_sink.second);
+  }
+
+  // Note: calling |AddOrUpdateSink()| or |RemoveSink()| here won't cause the
+  // discovery timer to fire again, since it is considered to be still running.
+  for (const auto& sink : sinks_to_update)
+    AddOrUpdateSink(sink);
+
+  for (const auto& sink : sinks_to_remove)
+    RemoveSink(sink);
+
+  latest_sinks_.clear();
+
+  // If discovered sinks are updated, then query results might have changed.
   for (const auto& query : sink_queries_)
     query.second->Notify(query.first);
+
+  MediaSinkServiceBase::OnDiscoveryComplete();
 }
 
 void DialMediaSinkServiceImpl::OnDialDeviceEvent(
@@ -160,13 +172,12 @@ void DialMediaSinkServiceImpl::OnDialDeviceEvent(
   DVLOG(2) << "DialMediaSinkServiceImpl::OnDialDeviceEvent found "
            << devices.size() << " devices";
 
-  current_sinks_.clear();
   current_devices_ = devices;
 
   description_service_->GetDeviceDescriptions(devices);
 
   // Makes sure the timer fires even if there is no device.
-  MediaSinkServiceBase::RestartTimer();
+  StartTimer();
 }
 
 void DialMediaSinkServiceImpl::OnDialError(DialRegistry::DialErrorCode type) {
@@ -185,7 +196,8 @@ void DialMediaSinkServiceImpl::OnDeviceDescriptionAvailable(
 
   std::string processed_uuid =
       MediaSinkInternal::ProcessDeviceUUID(description_data.unique_id);
-  std::string sink_id = base::StringPrintf("dial:<%s>", processed_uuid.c_str());
+  MediaSink::Id sink_id =
+      base::StringPrintf("dial:<%s>", processed_uuid.c_str());
   MediaSink sink(sink_id, description_data.friendly_name, SinkIconType::GENERIC,
                  MediaRouteProviderId::DIAL);
   DialSinkExtraData extra_data;
@@ -198,19 +210,14 @@ void DialMediaSinkServiceImpl::OnDeviceDescriptionAvailable(
   }
 
   MediaSinkInternal dial_sink(sink, extra_data);
-  current_sinks_.insert_or_assign(sink_id, dial_sink);
-  if (dial_sink_added_cb_)
-    dial_sink_added_cb_.Run(dial_sink);
+  latest_sinks_.insert_or_assign(sink_id, dial_sink);
+  StartTimer();
 
   if (!IsDiscoveryOnly(description_data.model_name)) {
     // Start checking if all registered apps are available on |dial_sink|.
     for (const auto& query : sink_queries_)
       FetchAppInfoForSink(dial_sink, query.first);
   }
-
-  // Start fetch timer again if device description comes back after
-  // |finish_timer_| fires.
-  MediaSinkServiceBase::RestartTimer();
 }
 
 void DialMediaSinkServiceImpl::OnDeviceDescriptionError(
@@ -244,7 +251,7 @@ void DialMediaSinkServiceImpl::OnAppInfoParseCompleted(
   // The sink might've been removed before the parse was complete. In that case
   // the callbacks won't be notified, but the app status will be saved for later
   // use.
-  if (!base::ContainsKey(current_sinks_, sink_id))
+  if (!GetSinkById(sink_id))
     return;
 
   auto query_it = sink_queries_.find(app_name);
@@ -274,17 +281,16 @@ void DialMediaSinkServiceImpl::FetchAppInfoForSink(
 }
 
 void DialMediaSinkServiceImpl::RescanAppInfo() {
-  for (const auto& dial_sink_it : current_sinks_) {
-    std::string model_name = dial_sink_it.second.dial_data().model_name;
+  for (const auto& sink : GetSinks()) {
+    std::string model_name = sink.second.dial_data().model_name;
     if (IsDiscoveryOnly(model_name)) {
       DVLOG(2) << "Model name does not support DIAL app availability: "
                << model_name;
       continue;
     }
 
-    for (const auto& query : sink_queries_) {
-      FetchAppInfoForSink(dial_sink_it.second, query.first);
-    }
+    for (const auto& query : sink_queries_)
+      FetchAppInfoForSink(sink.second, query.first);
   }
 }
 
@@ -306,7 +312,7 @@ void DialMediaSinkServiceImpl::SetAppStatus(const std::string& sink_id,
 
 void DialMediaSinkServiceImpl::RecordDeviceCounts() {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
-  metrics_.RecordDeviceCountsIfNeeded(current_sinks_.size(),
+  metrics_.RecordDeviceCountsIfNeeded(GetSinks().size(),
                                       current_devices_.size());
 }
 
@@ -324,10 +330,9 @@ std::vector<MediaSinkInternal> DialMediaSinkServiceImpl::GetAvailableSinks(
     const std::string& app_name) const {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   std::vector<MediaSinkInternal> sinks;
-  for (const auto& sink_it : current_sinks_) {
-    std::string sink_id = sink_it.second.sink().id();
-    if (GetAppStatus(sink_id, app_name) == SinkAppStatus::kAvailable)
-      sinks.push_back(sink_it.second);
+  for (const auto& sink : GetSinks()) {
+    if (GetAppStatus(sink.first, app_name) == SinkAppStatus::kAvailable)
+      sinks.push_back(sink.second);
   }
   return sinks;
 }

chrome/browser/media/router/discovery/dial/dial_media_sink_service_impl.h

@@ -39,6 +39,8 @@ class DialMediaSinkServiceImpl : public MediaSinkServiceBase,
   // Callbacks invoked when the list of available sinks for |app_name| changes.
   // The client can call |GetAvailableSinks()| to obtain the latest sink list.
   // |app_name|: app name, e.g. YouTube.
+  // TODO(imcheng): Move sink query logic into DialAppDiscoveryService and
+  // have it use MediaSinkServiceBase::Observer to observe sinks.
   using SinkQueryByAppFunc = void(const std::string& app_name);
   using SinkQueryByAppCallback = base::RepeatingCallback<SinkQueryByAppFunc>;
   using SinkQueryByAppCallbackList = base::CallbackList<SinkQueryByAppFunc>;
@@ -51,20 +53,18 @@ class DialMediaSinkServiceImpl : public MediaSinkServiceBase,
   // |connector|: connector to the ServiceManager suitable to use on
   // |task_runner|.
   // |on_sinks_discovered_cb|: Callback for MediaSinkServiceBase.
-  // |dial_sink_added_cb|: If not null, callback to invoke when a DIAL sink has
-  // been discovered.
   // Note that both callbacks are invoked on |task_runner|.
   // |task_runner|: The SequencedTaskRunner this class runs in.
   DialMediaSinkServiceImpl(
       std::unique_ptr<service_manager::Connector> connector,
       const OnSinksDiscoveredCallback& on_sinks_discovered_cb,
-      const OnDialSinkAddedCallback& dial_sink_added_cb,
       const scoped_refptr<base::SequencedTaskRunner>& task_runner);
   ~DialMediaSinkServiceImpl() override;
 
   virtual void Start();
 
-  void OnUserGesture();
+  // MediaSinkServiceBase implementation.
+  void OnUserGesture() override;
 
   // Returns the SequencedTaskRunner that should be used to invoke methods on
   // this instance. Can be invoked on any thread.
@@ -97,9 +97,6 @@ class DialMediaSinkServiceImpl : public MediaSinkServiceBase,
   void SetAppDiscoveryServiceForTest(
       std::unique_ptr<DialAppDiscoveryService> app_discovery_service);
 
-  // MediaSinkServiceBase implementation.
-  void OnDiscoveryComplete() override;
-
  private:
   friend class DialMediaSinkServiceImplTest;
   friend class MockDialMediaSinkServiceImpl;
@@ -111,8 +108,6 @@ class DialMediaSinkServiceImpl : public MediaSinkServiceBase,
                            OnDeviceDescriptionAvailable);
   FRIEND_TEST_ALL_PREFIXES(DialMediaSinkServiceImplTest,
                            OnDeviceDescriptionAvailableIPAddressChanged);
-  FRIEND_TEST_ALL_PREFIXES(DialMediaSinkServiceImplTest,
-                           OnDialSinkAddedCallback);
   FRIEND_TEST_ALL_PREFIXES(DialMediaSinkServiceImplTest,
                            StartStopMonitoringAvailableSinksForApp);
   FRIEND_TEST_ALL_PREFIXES(DialMediaSinkServiceImplTest,
@@ -171,6 +166,7 @@ class DialMediaSinkServiceImpl : public MediaSinkServiceBase,
       SinkQueryByAppCallbackList* callback_list);
 
   // MediaSinkServiceBase implementation.
+  void OnDiscoveryComplete() override;
   void RecordDeviceCounts() override;
 
   // Connector to ServiceManager for safe XML parsing requests.
@@ -182,8 +178,6 @@ class DialMediaSinkServiceImpl : public MediaSinkServiceBase,
   // Initialized in |Start()|.
   std::unique_ptr<DialAppDiscoveryService> app_discovery_service_;
 
-  OnDialSinkAddedCallback dial_sink_added_cb_;
-
   // Raw pointer to DialRegistry singleton.
   DialRegistry* dial_registry_ = nullptr;
 
@@ -193,6 +187,11 @@ class DialMediaSinkServiceImpl : public MediaSinkServiceBase,
   // Device data list from current round of discovery.
   DialRegistry::DeviceList current_devices_;
 
+  // Sinks that are added during the latest round of discovery. In
+  // |OnDiscoveryCompleted()| this will be merged into
+  // |MediaSinkServiceBase::sinks_| and then cleared.
+  base::flat_map<MediaSink::Id, MediaSinkInternal> latest_sinks_;
+
   // Map of app status, keyed by <sink id:app name>.
   base::flat_map<std::string, SinkAppStatus> app_statuses_;
 

chrome/browser/media/router/discovery/mdns/cast_media_sink_service.cc

@@ -32,7 +32,7 @@ CastMediaSinkService::~CastMediaSinkService() {
 
 void CastMediaSinkService::Start(
     const OnSinksDiscoveredCallback& sinks_discovered_cb,
-    CastMediaSinkServiceImpl::Observer* observer) {
+    MediaSinkServiceBase* dial_media_sink_service) {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   DCHECK(!impl_);
 
@@ -49,7 +49,7 @@ void CastMediaSinkService::Start(
                      base::BindRepeating(
                          &CastMediaSinkService::RunSinksDiscoveredCallback,
                          weak_ptr_factory_.GetWeakPtr(), sinks_discovered_cb)),
-                 observer);
+                 dial_media_sink_service);
   impl_->task_runner()->PostTask(
       FROM_HERE, base::BindOnce(&CastMediaSinkServiceImpl::Start,
                                 base::Unretained(impl_.get())));
@@ -62,7 +62,7 @@ void CastMediaSinkService::Start(
 std::unique_ptr<CastMediaSinkServiceImpl, base::OnTaskRunnerDeleter>
 CastMediaSinkService::CreateImpl(
     const OnSinksDiscoveredCallback& sinks_discovered_cb,
-    CastMediaSinkServiceImpl::Observer* observer) {
+    MediaSinkServiceBase* dial_media_sink_service) {
   cast_channel::CastSocketService* cast_socket_service =
       cast_channel::CastSocketService::GetInstance();
   scoped_refptr<base::SequencedTaskRunner> task_runner =
@@ -75,8 +75,8 @@ CastMediaSinkService::CreateImpl(
                           base::Unretained(this)));
   return std::unique_ptr<CastMediaSinkServiceImpl, base::OnTaskRunnerDeleter>(
       new CastMediaSinkServiceImpl(
-          sinks_discovered_cb, observer, cast_socket_service,
-          DiscoveryNetworkMonitor::GetInstance(),
+          sinks_discovered_cb, cast_socket_service,
+          DiscoveryNetworkMonitor::GetInstance(), dial_media_sink_service,
           GetCastAllowAllIPsPref(g_browser_process->local_state())),
       base::OnTaskRunnerDeleter(task_runner));
 }
@@ -109,7 +109,7 @@ void CastMediaSinkService::OnUserGesture() {
   DVLOG(2) << "OnUserGesture: open channel now for " << cast_sinks_.size()
            << " devices discovered in latest round of mDNS";
   impl_->task_runner()->PostTask(
-      FROM_HERE, base::BindOnce(&CastMediaSinkServiceImpl::AttemptConnection,
+      FROM_HERE, base::BindOnce(&CastMediaSinkServiceImpl::OpenChannelsNow,
                                 base::Unretained(impl_.get()), cast_sinks_));
 }
 
@@ -148,10 +148,6 @@ void CastMediaSinkService::OnDnsSdEvent(
                      CastMediaSinkServiceImpl::SinkSource::kMdns));
 }
 
-OnDialSinkAddedCallback CastMediaSinkService::GetDialSinkAddedCallback() {
-  return impl_->GetDialSinkAddedCallback();
-}
-
 void CastMediaSinkService::RunSinksDiscoveredCallback(
     const OnSinksDiscoveredCallback& sinks_discovered_cb,
     std::vector<MediaSinkInternal> sinks) {

chrome/browser/media/router/discovery/mdns/cast_media_sink_service.h

@@ -26,25 +26,23 @@ namespace media_router {
 // A service which can be used to start background discovery and resolution of
 // Cast devices.
 // This class is not thread safe. All methods must be invoked on the UI thread.
+// TODO(imcheng): Consider removing this class and moving the logic into a part
+// of CastMediaSinkServiceImpl that runs on the UI thread, and renaming
+// CastMediaSinkServiceImpl to CastMediaSinkService. Longer term, we
+// should look into eliminating dependencies on the UI thread.
 class CastMediaSinkService : public DnsSdRegistry::DnsSdObserver {
  public:
   CastMediaSinkService();
   ~CastMediaSinkService() override;
 
-  // Returns a callback to |impl_| when a DIAL sink is added (e.g., in order
-  // to perform dual discovery). The callback must be run on the same sequence
-  // as |impl_| and must not be run after |impl_| is destroyed.
-  // This method can only be called after |Start()| is called.
-  OnDialSinkAddedCallback GetDialSinkAddedCallback();
-
   // Starts Cast sink discovery. No-ops if already started.
   // |sink_discovery_cb|: Callback to invoke when the list of discovered sinks
   // has been updated.
-  // |observer|: Observer passed to |impl_|. Note that unlike the callback, the
-  // observer will be invoked on the sequence |impl_| runs on. Can be nullptr.
+  // |dial_media_sink_service|: Pointer to DIAL MediaSinkService for dual
+  // discovery.
   // Marked virtual for tests.
   virtual void Start(const OnSinksDiscoveredCallback& sinks_discovered_cb,
-                     CastMediaSinkServiceImpl::Observer* observer);
+                     MediaSinkServiceBase* dial_media_sink_service);
 
   // Initiates discovery immediately in response to a user gesture
   // (i.e., opening the Media Router dialog).
@@ -55,7 +53,9 @@ class CastMediaSinkService : public DnsSdRegistry::DnsSdObserver {
   // Marked virtual for tests.
   virtual std::unique_ptr<CastMediaSinkServiceImpl, base::OnTaskRunnerDeleter>
   CreateImpl(const OnSinksDiscoveredCallback& sinks_discovered_cb,
-             CastMediaSinkServiceImpl::Observer* observer);
+             MediaSinkServiceBase* dial_media_sink_service);
+
+  CastMediaSinkServiceImpl* impl() { return impl_.get(); }
 
   // Registers with DnsSdRegistry to listen for Cast devices. Note that this is
   // called on |Start()| on all platforms except for Windows. On Windows, this

chrome/browser/media/router/discovery/mdns/cast_media_sink_service_unittest.cc

@@ -12,6 +12,7 @@
 #include "chrome/browser/media/router/discovery/mdns/media_sink_util.h"
 #include "chrome/browser/media/router/test/mock_dns_sd_registry.h"
 #include "chrome/browser/media/router/test/test_helper.h"
+#include "chrome/common/media_router/test/test_helper.h"
 #include "components/cast_channel/cast_socket.h"
 #include "components/cast_channel/cast_socket_service.h"
 #include "components/cast_channel/cast_test_util.h"
@@ -59,13 +60,13 @@ class MockCastMediaSinkServiceImpl : public CastMediaSinkServiceImpl {
  public:
   MockCastMediaSinkServiceImpl(
       const OnSinksDiscoveredCallback& callback,
-      CastMediaSinkServiceImpl::Observer* observer,
       cast_channel::CastSocketService* cast_socket_service,
-      DiscoveryNetworkMonitor* network_monitor)
+      DiscoveryNetworkMonitor* network_monitor,
+      MediaSinkServiceBase* dial_media_sink_service)
       : CastMediaSinkServiceImpl(callback,
-                                 observer,
                                  cast_socket_service,
                                  network_monitor,
+                                 dial_media_sink_service,
                                  /* allow_all_ips */ false),
         sinks_discovered_cb_(callback) {}
   ~MockCastMediaSinkServiceImpl() override {}
@@ -94,12 +95,12 @@ class TestCastMediaSinkService : public CastMediaSinkService {
 
   std::unique_ptr<CastMediaSinkServiceImpl, base::OnTaskRunnerDeleter>
   CreateImpl(const OnSinksDiscoveredCallback& sinks_discovered_cb,
-             CastMediaSinkServiceImpl::Observer* observer) override {
+             MediaSinkServiceBase* dial_media_sink_service) override {
     auto mock_impl = std::unique_ptr<MockCastMediaSinkServiceImpl,
                                      base::OnTaskRunnerDeleter>(
-        new MockCastMediaSinkServiceImpl(sinks_discovered_cb, observer,
-                                         cast_socket_service_,
-                                         network_monitor_),
+        new MockCastMediaSinkServiceImpl(sinks_discovered_cb,
+                                         cast_socket_service_, network_monitor_,
+                                         dial_media_sink_service),
         base::OnTaskRunnerDeleter(cast_socket_service_->task_runner()));
     mock_impl_ = mock_impl.get();
     return mock_impl;
@@ -108,8 +109,8 @@ class TestCastMediaSinkService : public CastMediaSinkService {
   MockCastMediaSinkServiceImpl* mock_impl() { return mock_impl_; }
 
  private:
-  cast_channel::CastSocketService* const cast_socket_service_ = nullptr;
-  DiscoveryNetworkMonitor* const network_monitor_ = nullptr;
+  cast_channel::CastSocketService* const cast_socket_service_;
+  DiscoveryNetworkMonitor* const network_monitor_;
   MockCastMediaSinkServiceImpl* mock_impl_ = nullptr;
 };
 
@@ -130,7 +131,7 @@ class CastMediaSinkServiceTest : public ::testing::Test {
     EXPECT_CALL(test_dns_sd_registry_, RegisterDnsSdListener(_));
     media_sink_service_->SetDnsSdRegistryForTest(&test_dns_sd_registry_);
     media_sink_service_->Start(mock_sink_discovered_ui_cb_.Get(),
-                               /* observer */ nullptr);
+                               &dial_media_sink_service_);
     mock_impl_ = media_sink_service_->mock_impl();
     ASSERT_TRUE(mock_impl_);
     EXPECT_CALL(*mock_impl_, DoStart()).WillOnce(InvokeWithoutArgs([this]() {
@@ -153,6 +154,7 @@ class CastMediaSinkServiceTest : public ::testing::Test {
   std::unique_ptr<net::NetworkChangeNotifier> network_change_notifier_;
 
   base::MockCallback<OnSinksDiscoveredCallback> mock_sink_discovered_ui_cb_;
+  TestMediaSinkService dial_media_sink_service_;
   std::unique_ptr<cast_channel::MockCastSocketService>
       mock_cast_socket_service_;
 

chrome/browser/media/router/providers/cast/cast_app_discovery_service.cc

@@ -40,19 +40,25 @@ bool ShouldRefreshAppAvailability(
 CastAppDiscoveryService::CastAppDiscoveryService(
     cast_channel::CastMessageHandler* message_handler,
     cast_channel::CastSocketService* socket_service,
+    MediaSinkServiceBase* media_sink_service,
     const base::TickClock* clock)
     : message_handler_(message_handler),
       socket_service_(socket_service),
+      media_sink_service_(media_sink_service),
       clock_(clock),
       weak_ptr_factory_(this) {
   DETACH_FROM_SEQUENCE(sequence_checker_);
   DCHECK(message_handler_);
   DCHECK(socket_service_);
   DCHECK(clock_);
+  socket_service_->task_runner()->PostTask(
+      FROM_HERE,
+      base::BindOnce(&CastAppDiscoveryService::Init, base::Unretained(this)));
 }
 
 CastAppDiscoveryService::~CastAppDiscoveryService() {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  media_sink_service_->RemoveObserver(this);
 }
 
 CastAppDiscoveryService::Subscription
@@ -81,10 +87,11 @@ CastAppDiscoveryService::StartObservingMediaSinks(
     // it changed.
     base::flat_set<std::string> new_app_ids =
         availability_tracker_.RegisterSource(source);
+    const auto& sinks = media_sink_service_->GetSinks();
     for (const auto& app_id : new_app_ids) {
-      // Note: The following logic assumes |sinks_| will not change as it is
+      // Note: The following logic assumes |sinks| will not change as it is
       // being iterated.
-      for (const auto& sink : sinks_) {
+      for (const auto& sink : sinks) {
         int channel_id = sink.second.cast_data().cast_channel_id;
         cast_channel::CastSocket* socket =
             socket_service_->GetSocket(channel_id);
@@ -102,11 +109,13 @@ CastAppDiscoveryService::StartObservingMediaSinks(
 }
 
 void CastAppDiscoveryService::Refresh() {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   const auto app_ids = availability_tracker_.GetRegisteredApps();
   base::TimeTicks now = clock_->NowTicks();
-  // Note: The following logic assumes |sinks_| will not change as it is
+  const auto& sinks = media_sink_service_->GetSinks();
+  // Note: The following logic assumes |sinks| will not change as it is
   // being iterated.
-  for (const auto& sink : sinks_) {
+  for (const auto& sink : sinks) {
     for (const auto& app_id : app_ids) {
       if (ShouldRefreshAppAvailability(
               availability_tracker_.GetAvailability(sink.first, app_id), now)) {
@@ -136,12 +145,21 @@ void CastAppDiscoveryService::MaybeRemoveSinkQueryEntry(
   }
 }
 
+void CastAppDiscoveryService::Init() {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  media_sink_service_->AddObserver(this);
+}
+
 void CastAppDiscoveryService::OnSinkAddedOrUpdated(
-    const MediaSinkInternal& sink,
-    cast_channel::CastSocket* socket) {
+    const MediaSinkInternal& sink) {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+
+  cast_channel::CastSocket* socket =
+      socket_service_->GetSocket(sink.cast_data().cast_channel_id);
+  if (!socket)
+    return;
+
   const MediaSink::Id& sink_id = sink.sink().id();
-  sinks_.insert_or_assign(sink_id, sink);
   for (const std::string& app_id : availability_tracker_.GetRegisteredApps()) {
     auto availability = availability_tracker_.GetAvailability(sink_id, app_id);
     if (availability.first != cast_channel::GetAppAvailabilityResult::kUnknown)
@@ -154,7 +172,6 @@ void CastAppDiscoveryService::OnSinkAddedOrUpdated(
 void CastAppDiscoveryService::OnSinkRemoved(const MediaSinkInternal& sink) {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   const MediaSink::Id& sink_id = sink.sink().id();
-  sinks_.erase(sink_id);
   UpdateSinkQueries(availability_tracker_.RemoveResultsForSink(sink_id));
 }
 
@@ -176,7 +193,7 @@ void CastAppDiscoveryService::UpdateAppAvailability(
     cast_channel::GetAppAvailabilityResult availability) {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   RecordAppAvailabilityResult(availability, clock_->NowTicks() - start_time);
-  if (!base::ContainsKey(sinks_, sink_id))
+  if (!media_sink_service_->GetSinkById(sink_id))
     return;
 
   DVLOG(1) << "App " << app_id << " on sink " << sink_id << " is "
@@ -203,9 +220,9 @@ std::vector<MediaSinkInternal> CastAppDiscoveryService::GetSinksByIds(
     const base::flat_set<MediaSink::Id>& sink_ids) const {
   std::vector<MediaSinkInternal> sinks;
   for (const auto& sink_id : sink_ids) {
-    auto it = sinks_.find(sink_id);
-    if (it != sinks_.end())
-      sinks.push_back(it->second);
+    const MediaSinkInternal* sink = media_sink_service_->GetSinkById(sink_id);
+    if (sink)
+      sinks.push_back(*sink);
   }
   return sinks;
 }

chrome/browser/media/router/providers/cast/cast_app_discovery_service.h

@@ -16,8 +16,8 @@
 #include "base/memory/weak_ptr.h"
 #include "base/sequenced_task_runner.h"
 #include "base/time/time.h"
-#include "chrome/browser/media/router/discovery/mdns/cast_media_sink_service_impl.h"
 #include "chrome/browser/media/router/providers/cast/cast_app_availability_tracker.h"
+#include "chrome/common/media_router/discovery/media_sink_service_base.h"
 #include "chrome/common/media_router/media_sink.h"
 #include "chrome/common/media_router/media_source.h"
 #include "chrome/common/media_router/providers/cast/cast_media_source.h"
@@ -30,6 +30,7 @@ class TickClock;
 namespace cast_channel {
 class CastMessageHandler;
 class CastSocket;
+class CastSocketService;
 }  // namespace cast_channel
 
 namespace media_router {
@@ -38,7 +39,7 @@ namespace media_router {
 // updates, and issues app availability requests based on these signals. This
 // class may be created on any sequence. All other methods must be called on the
 // CastSocketService sequence.
-class CastAppDiscoveryService : public CastMediaSinkServiceImpl::Observer {
+class CastAppDiscoveryService : public MediaSinkServiceBase::Observer {
  public:
   using SinkQueryFunc = void(const MediaSource::Id& source_id,
                              const std::vector<MediaSinkInternal>& sinks);
@@ -48,6 +49,7 @@ class CastAppDiscoveryService : public CastMediaSinkServiceImpl::Observer {
 
   CastAppDiscoveryService(cast_channel::CastMessageHandler* message_handler,
                           cast_channel::CastSocketService* socket_service,
+                          MediaSinkServiceBase* media_sink_service,
                           const base::TickClock* clock);
   ~CastAppDiscoveryService() override;
 
@@ -67,9 +69,11 @@ class CastAppDiscoveryService : public CastMediaSinkServiceImpl::Observer {
  private:
   friend class CastAppDiscoveryServiceTest;
 
-  // CastMediaSinkServiceImpl::Observer
-  void OnSinkAddedOrUpdated(const MediaSinkInternal& sink,
-                            cast_channel::CastSocket* socket) override;
+  // Called on construction. Registers an observer with |media_sink_service_|.
+  void Init();
+
+  // MediaSinkServiceBase::Observer
+  void OnSinkAddedOrUpdated(const MediaSinkInternal& sink) override;
   void OnSinkRemoved(const MediaSinkInternal& sink) override;
 
   // Issues an app availability request for |app_id| to the sink given by
@@ -104,9 +108,9 @@ class CastAppDiscoveryService : public CastMediaSinkServiceImpl::Observer {
 
   cast_channel::CastMessageHandler* const message_handler_;
   cast_channel::CastSocketService* const socket_service_;
+  MediaSinkServiceBase* const media_sink_service_;
 
   CastAppAvailabilityTracker availability_tracker_;
-  base::flat_map<MediaSink::Id, MediaSinkInternal> sinks_;
 
   const base::TickClock* const clock_;
 

chrome/browser/media/router/providers/cast/cast_app_discovery_service_unittest.cc

@@ -8,7 +8,9 @@
 #include "base/test/simple_test_tick_clock.h"
 #include "base/test/test_simple_task_runner.h"
 #include "chrome/browser/media/router/test/test_helper.h"
+#include "chrome/common/media_router/discovery/media_sink_service_base.h"
 #include "chrome/common/media_router/providers/cast/cast_media_source.h"
+#include "chrome/common/media_router/test/test_helper.h"
 #include "components/cast_channel/cast_test_util.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -28,12 +30,14 @@ class CastAppDiscoveryServiceTest : public testing::Test {
         app_discovery_service_(
             std::make_unique<CastAppDiscoveryService>(&message_handler_,
                                                       &socket_service_,
+                                                      &media_sink_service_,
                                                       &clock_)),
         source_a_1_(*CastMediaSource::From("cast:AAAAAAAA?clientId=1")),
         source_a_2_(*CastMediaSource::From("cast:AAAAAAAA?clientId=2")),
         source_b_1_(*CastMediaSource::From("cast:BBBBBBBB?clientId=1")) {
     ON_CALL(socket_service_, GetSocket(_))
         .WillByDefault(testing::Return(&socket_));
+    task_runner_->RunPendingTasks();
   }
 
   ~CastAppDiscoveryServiceTest() override { task_runner_->RunPendingTasks(); }
@@ -43,11 +47,11 @@ class CastAppDiscoveryServiceTest : public testing::Test {
                     const std::vector<MediaSinkInternal>&));
 
   void AddOrUpdateSink(const MediaSinkInternal& sink) {
-    app_discovery_service_->OnSinkAddedOrUpdated(sink, &socket_);
+    media_sink_service_.AddOrUpdateSink(sink);
   }
 
   void RemoveSink(const MediaSinkInternal& sink) {
-    app_discovery_service_->OnSinkRemoved(sink);
+    media_sink_service_.RemoveSink(sink);
   }
 
   CastAppDiscoveryService::Subscription StartObservingMediaSinksInitially(
@@ -66,6 +70,7 @@ class CastAppDiscoveryServiceTest : public testing::Test {
   cast_channel::MockCastSocketService socket_service_;
   cast_channel::MockCastSocket socket_;
   cast_channel::MockCastMessageHandler message_handler_;
+  TestMediaSinkService media_sink_service_;
   std::unique_ptr<CastAppDiscoveryService> app_discovery_service_;
   CastMediaSource source_a_1_;
   CastMediaSource source_a_2_;

chrome/browser/media/router/providers/cast/dual_media_sink_service.cc

@@ -47,8 +47,7 @@ DualMediaSinkService::AddSinksDiscoveredCallback(
 
 void DualMediaSinkService::OnUserGesture() {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
-
-  dial_media_sink_service_->OnUserGesture();
+  // TODO(imcheng): Move this call into CastMediaRouteProvider.
   if (cast_media_sink_service_)
     cast_media_sink_service_->OnUserGesture();
 }
@@ -61,36 +60,34 @@ void DualMediaSinkService::StartMdnsDiscovery() {
 }
 
 DualMediaSinkService::DualMediaSinkService() {
-  OnDialSinkAddedCallback dial_sink_added_cb;
+  dial_media_sink_service_ = std::make_unique<DialMediaSinkService>();
+  dial_media_sink_service_->Start(
+      base::BindRepeating(&DualMediaSinkService::OnSinksDiscovered,
+                          base::Unretained(this), "dial"));
+
   if (CastDiscoveryEnabled()) {
+    cast_media_sink_service_ = std::make_unique<CastMediaSinkService>();
+    cast_media_sink_service_->Start(
+        base::BindRepeating(&DualMediaSinkService::OnSinksDiscovered,
+                            base::Unretained(this), "cast"),
+        dial_media_sink_service_->impl());
+
     if (CastMediaRouteProviderEnabled()) {
       cast_channel::CastSocketService* cast_socket_service =
           cast_channel::CastSocketService::GetInstance();
       cast_app_discovery_service_ = std::make_unique<CastAppDiscoveryService>(
           GetCastMessageHandler(), cast_socket_service,
+          cast_media_sink_service_->impl(),
           base::DefaultTickClock::GetInstance());
     }
-
-    cast_media_sink_service_ = std::make_unique<CastMediaSinkService>();
-    cast_media_sink_service_->Start(
-        base::BindRepeating(&DualMediaSinkService::OnSinksDiscovered,
-                            base::Unretained(this), "cast"),
-        cast_app_discovery_service_.get());
-    dial_sink_added_cb = cast_media_sink_service_->GetDialSinkAddedCallback();
   }
-
-  dial_media_sink_service_ = std::make_unique<DialMediaSinkService>();
-  dial_media_sink_service_->Start(
-      base::BindRepeating(&DualMediaSinkService::OnSinksDiscovered,
-                          base::Unretained(this), "dial"),
-      dial_sink_added_cb);
 }
 
 DualMediaSinkService::DualMediaSinkService(
     std::unique_ptr<CastMediaSinkService> cast_media_sink_service,
     std::unique_ptr<DialMediaSinkService> dial_media_sink_service)
-    : cast_media_sink_service_(std::move(cast_media_sink_service)),
-      dial_media_sink_service_(std::move(dial_media_sink_service)) {}
+    : dial_media_sink_service_(std::move(dial_media_sink_service)),
+      cast_media_sink_service_(std::move(cast_media_sink_service)) {}
 
 DualMediaSinkService::~DualMediaSinkService() = default;
 

chrome/browser/media/router/providers/cast/dual_media_sink_service.h

@@ -15,6 +15,8 @@
 #include "base/gtest_prod_util.h"
 #include "base/macros.h"
 #include "base/sequence_checker.h"
+#include "chrome/browser/media/router/discovery/mdns/cast_media_sink_service.h"
+#include "chrome/browser/media/router/discovery/mdns/cast_media_sink_service_impl.h"
 #include "chrome/browser/media/router/media_sinks_observer.h"
 #include "chrome/common/media_router/discovery/media_sink_internal.h"
 #include "chrome/common/media_router/media_source.h"
@@ -48,10 +50,6 @@ class DualMediaSinkService {
   static DualMediaSinkService* GetInstance();
   static void SetInstanceForTest(DualMediaSinkService* instance_for_test);
 
-  CastAppDiscoveryService* cast_app_discovery_service() {
-    return cast_app_discovery_service_.get();
-  }
-
   // Returns the current list of sinks, keyed by provider name.
   const base::flat_map<std::string, std::vector<MediaSinkInternal>>&
   current_sinks() {
@@ -61,6 +59,10 @@ class DualMediaSinkService {
   // Used by DialMediaRouteProvider only.
   DialMediaSinkServiceImpl* GetDialMediaSinkServiceImpl();
 
+  CastAppDiscoveryService* cast_app_discovery_service() {
+    return cast_app_discovery_service_.get();
+  }
+
   // Adds |callback| to be notified when the list of discovered sinks changes.
   // The caller is responsible for destroying the returned Subscription when it
   // no longer wishes to receive updates.
@@ -88,18 +90,20 @@ class DualMediaSinkService {
   friend class MediaRouterDesktopTestBase;
   FRIEND_TEST_ALL_PREFIXES(MediaRouterDesktopTest, ProvideSinks);
 
-  friend struct std::default_delete<DualMediaSinkService>;
-
   static DualMediaSinkService* instance_for_test_;
 
+  friend struct std::default_delete<DualMediaSinkService>;
+
   DualMediaSinkService();
 
   void OnSinksDiscovered(const std::string& provider_name,
                          std::vector<MediaSinkInternal> sinks);
 
+  // Note: Dual discovery logic assumes |dial_media_sink_service_| outlives
+  // |cast_media_sink_service_|.
+  std::unique_ptr<DialMediaSinkService> dial_media_sink_service_;
   std::unique_ptr<CastMediaSinkService> cast_media_sink_service_;
   std::unique_ptr<CastAppDiscoveryService> cast_app_discovery_service_;
-  std::unique_ptr<DialMediaSinkService> dial_media_sink_service_;
 
   OnSinksDiscoveredProviderCallbackList sinks_discovered_callbacks_;
   base::flat_map<std::string, std::vector<MediaSinkInternal>> current_sinks_;

chrome/browser/media/router/providers/cast/dual_media_sink_service_unittest.cc

@@ -46,7 +46,6 @@ class DualMediaSinkServiceTest : public testing::Test {
 
 TEST_F(DualMediaSinkServiceTest, OnUserGesture) {
   EXPECT_CALL(*cast_media_sink_service(), OnUserGesture());
-  EXPECT_CALL(*dial_media_sink_service(), OnUserGesture());
   dual_media_sink_service()->OnUserGesture();
 }
 

chrome/browser/media/router/providers/dial/dial_media_route_provider.cc

@@ -199,7 +199,7 @@ void DialMediaRouteProvider::EnableMdnsDiscovery() {
 }
 
 void DialMediaRouteProvider::UpdateMediaSinks(const std::string& media_source) {
-  NOTIMPLEMENTED();
+  media_sink_service_->OnUserGesture();
 }
 
 void DialMediaRouteProvider::SearchSinks(

chrome/browser/media/router/providers/dial/dial_media_route_provider_unittest.cc

@@ -25,7 +25,6 @@ class TestDialMediaSinkServiceImpl : public DialMediaSinkServiceImpl {
  public:
   TestDialMediaSinkServiceImpl()
       : DialMediaSinkServiceImpl(/* connector */ nullptr,
-                                 base::DoNothing(),
                                  base::DoNothing(),
                                  /* task_runner */ nullptr) {}
 

chrome/browser/media/router/test/test_helper.h

@@ -123,9 +123,7 @@ class MockDialMediaSinkService : public DialMediaSinkService {
   MockDialMediaSinkService();
   ~MockDialMediaSinkService() override;
 
-  MOCK_METHOD2(Start,
-               void(const OnSinksDiscoveredCallback&,
-                    const OnDialSinkAddedCallback&));
+  MOCK_METHOD1(Start, void(const OnSinksDiscoveredCallback&));
   MOCK_METHOD0(OnUserGesture, void());
 };
 
@@ -135,8 +133,7 @@ class MockCastMediaSinkService : public CastMediaSinkService {
   ~MockCastMediaSinkService() override;
 
   MOCK_METHOD2(Start,
-               void(const OnSinksDiscoveredCallback&,
-                    CastMediaSinkServiceImpl::Observer*));
+               void(const OnSinksDiscoveredCallback&, MediaSinkServiceBase*));
   MOCK_METHOD0(OnUserGesture, void());
   MOCK_METHOD0(StartMdnsDiscovery, void());
 };

chrome/common/media_router/BUILD.gn

@@ -50,3 +50,18 @@ static_library("router") {
     deps += [ "//components/cast_channel:cast_channel" ]
   }
 }
+
+static_library("test_support") {
+  testonly = true
+  deps = [
+    "//base:base",
+    "//chrome/test:test_support",
+  ]
+  public_deps = [
+    ":router",
+  ]
+  sources = [
+    "test/test_helper.cc",
+    "test/test_helper.h",
+  ]
+}

chrome/common/media_router/discovery/media_sink_service_base.cc

@@ -7,69 +7,105 @@
 #include <vector>
 
 namespace {
-// Time interval when media sink service sends sinks to MRP.
-const int kFetchCompleteTimeoutSecs = 3;
+// Timeout amount for |discovery_timer_|.
+const constexpr base::TimeDelta kDiscoveryTimeout =
+    base::TimeDelta::FromSeconds(3);
 }  // namespace
 
 namespace media_router {
 
 MediaSinkServiceBase::MediaSinkServiceBase(
     const OnSinksDiscoveredCallback& callback)
-    : on_sinks_discovered_cb_(callback) {
+    : discovery_timer_(std::make_unique<base::OneShotTimer>()),
+      on_sinks_discovered_cb_(callback) {
   DETACH_FROM_SEQUENCE(sequence_checker_);
-  DCHECK(!on_sinks_discovered_cb_.is_null());
 }
 
 MediaSinkServiceBase::~MediaSinkServiceBase() {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
 }
 
-void MediaSinkServiceBase::SetTimerForTest(std::unique_ptr<base::Timer> timer) {
-  DCHECK(!discovery_timer_);
-  discovery_timer_ = std::move(timer);
+void MediaSinkServiceBase::AddObserver(Observer* observer) {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  observers_.AddObserver(observer);
 }
 
-void MediaSinkServiceBase::OnDiscoveryComplete() {
+void MediaSinkServiceBase::RemoveObserver(Observer* observer) {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
-  if (current_sinks_ == mrp_sinks_) {
-    DVLOG(2) << "No update to sink list.";
-    return;
-  }
-
-  DVLOG(2) << "Send sinks to media router, [size]: " << current_sinks_.size();
+  observers_.RemoveObserver(observer);
+}
 
-  std::vector<MediaSinkInternal> sinks;
-  for (const auto& sink_it : current_sinks_)
-    sinks.push_back(sink_it.second);
+const base::flat_map<MediaSink::Id, MediaSinkInternal>&
+MediaSinkServiceBase::GetSinks() const {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  return sinks_;
+}
 
-  on_sinks_discovered_cb_.Run(std::move(sinks));
-  mrp_sinks_ = current_sinks_;
-  discovery_timer_->Stop();
-  RecordDeviceCounts();
+const MediaSinkInternal* MediaSinkServiceBase::GetSinkById(
+    const MediaSink::Id& sink_id) const {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  auto it = sinks_.find(sink_id);
+  return it != sinks_.end() ? &it->second : nullptr;
 }
 
-void MediaSinkServiceBase::StartTimer() {
+void MediaSinkServiceBase::AddOrUpdateSink(const MediaSinkInternal& sink) {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
-  // |discovery_timer_| is already set to a mock timer in unit tests only.
-  if (!discovery_timer_)
-    discovery_timer_.reset(new base::OneShotTimer());
+  sinks_.insert_or_assign(sink.sink().id(), sink);
+  for (auto& observer : observers_)
+    observer.OnSinkAddedOrUpdated(sink);
 
-  DoStart();
+  StartTimer();
 }
 
-void MediaSinkServiceBase::RestartTimer() {
+void MediaSinkServiceBase::RemoveSink(const MediaSinkInternal& sink) {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
-  if (!discovery_timer_->IsRunning())
-    DoStart();
+
+  // Make a copy of the sink to avoid potential use-after-free.
+  MediaSink::Id sink_id = sink.sink().id();
+  MediaSinkInternal sink_copy = sink;
+  if (!sinks_.erase(sink_id))
+    return;
+
+  for (auto& observer : observers_)
+    observer.OnSinkRemoved(sink_copy);
+
+  StartTimer();
+}
+
+void MediaSinkServiceBase::SetTimerForTest(std::unique_ptr<base::Timer> timer) {
+  discovery_timer_ = std::move(timer);
 }
 
-void MediaSinkServiceBase::DoStart() {
-  base::TimeDelta finish_delay =
-      base::TimeDelta::FromSeconds(kFetchCompleteTimeoutSecs);
+void MediaSinkServiceBase::StartTimer() {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  if (discovery_timer_->IsRunning())
+    return;
+
   discovery_timer_->Start(
-      FROM_HERE, finish_delay,
+      FROM_HERE, kDiscoveryTimeout,
       base::BindRepeating(&MediaSinkServiceBase::OnDiscoveryComplete,
                           base::Unretained(this)));
 }
 
+void MediaSinkServiceBase::OnDiscoveryComplete() {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  discovery_timer_->Stop();
+  RecordDeviceCounts();
+
+  // Only send discovered sinks back to MediaRouter if the list changed.
+  if (sinks_ == previous_sinks_) {
+    DVLOG(2) << "No update to sink list.";
+    return;
+  }
+
+  DVLOG(2) << "Send sinks to media router, [size]: " << sinks_.size();
+
+  std::vector<MediaSinkInternal> sinks;
+  for (const auto& sink_it : sinks_)
+    sinks.push_back(sink_it.second);
+
+  on_sinks_discovered_cb_.Run(std::move(sinks));
+  previous_sinks_ = sinks_;
+}
+
 }  // namespace media_router

chrome/common/media_router/discovery/media_sink_service_base.h

@@ -9,6 +9,7 @@
 
 #include "base/containers/flat_map.h"
 #include "base/gtest_prod_util.h"
+#include "base/observer_list.h"
 #include "base/sequence_checker.h"
 #include "base/timer/timer.h"
 #include "chrome/common/media_router/discovery/media_sink_internal.h"
@@ -16,39 +17,68 @@
 
 namespace media_router {
 
-// Base class for discovering MediaSinks and notifying the caller with updated
-// list. The class rate-limits notifications to |callback|, so that multiple
-// updates received in quick succession will be batched.
+// Base class for discovering MediaSinks. Responsible for bookkeeping of
+// current set of discovered sinks, and notifying observers when there are
+// updates.
+// In addition, this class maintains a "discovery timer", used for batching
+// updates in quick succession. The timer fires when it is assumed that
+// discovery has reached a relatively steady state. When the timer fires:
+// - The batched updated sink list will be sent back to the Media Router
+// extension via |callback|. This back-channel is necessary until all logic
+// dependent on MediaSinks are moved out of the extension.
+// - Subclasses may record discovered related metrics.
 // This class may be created on any thread, but all subsequent methods must be
 // invoked on the same thread.
 class MediaSinkServiceBase {
  public:
-  // |callback|: The callback to invoke when the list of MediaSinks has been
-  // updated.
+  // Listens for sink updates in MediaSinkServiceBase.
+  class Observer {
+   public:
+    virtual ~Observer() = default;
+
+    // Invoked when |sink| is added or updated.
+    virtual void OnSinkAddedOrUpdated(const MediaSinkInternal& sink) = 0;
+
+    // Invoked when |sink| is removed.
+    virtual void OnSinkRemoved(const MediaSinkInternal& sink) = 0;
+  };
+
+  // |callback|: Callback to invoke inform MediaRouter of discovered sinks
+  // updates.
   explicit MediaSinkServiceBase(const OnSinksDiscoveredCallback& callback);
   virtual ~MediaSinkServiceBase();
 
- protected:
-  void SetTimerForTest(std::unique_ptr<base::Timer> timer);
+  // Adds |observer| to observe |this| for sink updates.
+  // Caller is responsible for calling |RemoveObserver| before it is destroyed.
+  // Both methods are safe to call on any thread.
+  void AddObserver(Observer* observer);
+  void RemoveObserver(Observer* observer);
 
-  // Called when |discovery_timer_| expires.
-  virtual void OnDiscoveryComplete();
+  // Overridden by subclass to initiate action triggered by user gesture, e.g.
+  // start one-off round of discovery.
+  virtual void OnUserGesture() {}
 
-  // Overriden by subclass to report device counts.
-  virtual void RecordDeviceCounts() {}
+  // Adds or updates, or removes a sink.
+  // Notifies |observers_| that the sink has been added, updated, or removed.
+  // Also invokes |StartTimer()|.
+  void AddOrUpdateSink(const MediaSinkInternal& sink);
+  void RemoveSink(const MediaSinkInternal& sink);
 
-  // Creates |discovery_timer_| and starts it.
-  void StartTimer();
+  const base::flat_map<MediaSink::Id, MediaSinkInternal>& GetSinks() const;
+  const MediaSinkInternal* GetSinkById(const MediaSink::Id& sink_id) const;
 
-  // Restarts |discovery_timer| if it is stopped. Can only be called after
-  // |StartTimer()| is called.
-  void RestartTimer();
+  void SetTimerForTest(std::unique_ptr<base::Timer> timer);
 
-  // Sorted sinks from current round of discovery, keyed by sink ID. Subclasses
-  // may make modifications to this map while discovery is active. At the
-  // completion of the current round of discovery, |current_sinks_| will be
-  // copied over to |mrp_sinks_|, which will be used for |GetSink()| calls.
-  base::flat_map<std::string, MediaSinkInternal> current_sinks_;
+ protected:
+  // Called when |discovery_timer_| expires. Informs subclass to report device
+  // counts. Also informs Media Router of updated list of discovered sinks.
+  // May be overridden by subclass to perform additional operations, such as
+  // pruning old sinks.
+  virtual void OnDiscoveryComplete();
+
+  // Starts |discovery_timer_| to invoke |OnDiscoveryComplete()|. Subclasses
+  // may call this at the start of a round of discovery.
+  void StartTimer();
 
  private:
   friend class MediaSinkServiceBaseTest;
@@ -57,18 +87,34 @@ class MediaSinkServiceBase {
   FRIEND_TEST_ALL_PREFIXES(MediaSinkServiceBaseTest,
                            TestOnDiscoveryComplete_SameSinkDifferentOrders);
 
-  // Helper method to start |discovery_timer_|.
-  void DoStart();
+  // Overriden by subclass to report device counts.
+  virtual void RecordDeviceCounts() {}
 
-  OnSinksDiscoveredCallback on_sinks_discovered_cb_;
+  // The current set of discovered sinks keyed by MediaSink ID.
+  base::flat_map<MediaSink::Id, MediaSinkInternal> sinks_;
 
-  // Sorted sinks sent to Media Router Provider in last |OnDiscoveryComplete()|,
-  // keyed by sink ID.
-  base::flat_map<std::string, MediaSinkInternal> mrp_sinks_;
+  // Observers to notify when a sink is added, updated, or removed.
+  base::ObserverList<Observer> observers_;
 
-  // Timer for completing the current round of discovery.
+  // Timer for recording device counts after a sink list has changed. To ensure
+  // the metrics are recorded accurately, a small delay is introduced after a
+  // sink list change in order for the discovery process to reach a steady
+  // state before the metrics are recorded.
   std::unique_ptr<base::Timer> discovery_timer_;
 
+  // The following fields exist temporarily for sending back discovered sinks to
+  // the Media Router extension.
+  // TODO(https://crbug.com/809249): Remove once the extension no longer need
+  // the sinks.
+
+  // Callback to MediaRouter to provide sinks to the MR extension.
+  OnSinksDiscoveredCallback on_sinks_discovered_cb_;
+
+  // Sinks saved in the previous |OnDiscoveryComplete()| invocation. Checked
+  // against |sinks_| during |OnDiscoveryComplete()| before invoking
+  // |on_sinks_discovered_cb_|.
+  base::flat_map<MediaSink::Id, MediaSinkInternal> previous_sinks_;
+
   SEQUENCE_CHECKER(sequence_checker_);
   DISALLOW_COPY_AND_ASSIGN(MediaSinkServiceBase);
 };

chrome/common/media_router/discovery/media_sink_service_base_unittest.cc

@@ -3,9 +3,11 @@
 // found in the LICENSE file.
 
 #include "chrome/common/media_router/discovery/media_sink_service_base.h"
+
 #include "base/memory/ptr_util.h"
 #include "base/test/mock_callback.h"
 #include "base/timer/mock_timer.h"
+#include "chrome/common/media_router/test/test_helper.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
@@ -52,25 +54,16 @@ class MediaSinkServiceBaseTest : public ::testing::Test {
       : media_sink_service_(mock_sink_discovered_cb_.Get()) {}
   ~MediaSinkServiceBaseTest() override {}
 
-  void SetUp() override {
-    mock_timer_ =
-        new base::MockTimer(true /*retain_user_task*/, false /*is_repeating*/);
-    media_sink_service_.SetTimerForTest(base::WrapUnique(mock_timer_));
-  }
-
   void PopulateSinks(const std::vector<MediaSinkInternal>& old_sinks,
                      const std::vector<MediaSinkInternal>& new_sinks) {
-    media_sink_service_.mrp_sinks_.clear();
-    for (const auto& old_sink : old_sinks) {
-      std::string sink_id = old_sink.sink().id();
-      media_sink_service_.mrp_sinks_[sink_id] = old_sink;
-    }
-
-    media_sink_service_.current_sinks_.clear();
-    for (const auto& new_sink : new_sinks) {
-      std::string sink_id = new_sink.sink().id();
-      media_sink_service_.current_sinks_[sink_id] = new_sink;
-    }
+    media_sink_service_.previous_sinks_.clear();
+    for (const auto& old_sink : old_sinks)
+      media_sink_service_.previous_sinks_.emplace(old_sink.sink().id(),
+                                                  old_sink);
+
+    media_sink_service_.sinks_.clear();
+    for (const auto& new_sink : new_sinks)
+      media_sink_service_.sinks_.emplace(new_sink.sink().id(), new_sink);
   }
 
   void TestOnDiscoveryComplete(
@@ -83,9 +76,7 @@ class MediaSinkServiceBaseTest : public ::testing::Test {
 
  protected:
   base::MockCallback<OnSinksDiscoveredCallback> mock_sink_discovered_cb_;
-  base::MockTimer* mock_timer_;
-
-  MediaSinkServiceBase media_sink_service_;
+  TestMediaSinkService media_sink_service_;
 
   DISALLOW_COPY_AND_ASSIGN(MediaSinkServiceBaseTest);
 };

chrome/common/media_router/test/test_helper.cc

+// Copyright 2018 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 "chrome/common/media_router/test/test_helper.h"
+
+#include "base/bind_helpers.h"
+#include "base/memory/ptr_util.h"
+
+namespace media_router {
+
+#if !defined(OS_ANDROID)
+TestMediaSinkService::TestMediaSinkService()
+    : TestMediaSinkService(base::DoNothing()) {}
+
+TestMediaSinkService::TestMediaSinkService(
+    const OnSinksDiscoveredCallback& callback)
+    : MediaSinkServiceBase(callback),
+      timer_(new base::MockTimer(true /*retain_user_task*/,
+                                 false /*is_repeating*/)) {
+  SetTimerForTest(base::WrapUnique(timer_));
+}
+
+TestMediaSinkService::~TestMediaSinkService() = default;
+#endif  // !defined(OS_ANDROID)
+
+}  // namespace media_router

chrome/common/media_router/test/test_helper.h

+// Copyright 2018 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.
+
+#ifndef CHROME_COMMON_MEDIA_ROUTER_TEST_TEST_HELPER_H_
+#define CHROME_COMMON_MEDIA_ROUTER_TEST_TEST_HELPER_H_
+
+#include "base/macros.h"
+#include "base/timer/mock_timer.h"
+#include "build/build_config.h"
+#include "chrome/common/media_router/discovery/media_sink_service_base.h"
+
+namespace media_router {
+
+#if !defined(OS_ANDROID)
+class TestMediaSinkService : public MediaSinkServiceBase {
+ public:
+  TestMediaSinkService();
+  explicit TestMediaSinkService(const OnSinksDiscoveredCallback& callback);
+  ~TestMediaSinkService() override;
+
+  base::MockTimer* timer() { return timer_; }
+
+ private:
+  // Owned by MediaSinkService.
+  base::MockTimer* timer_;
+  DISALLOW_COPY_AND_ASSIGN(TestMediaSinkService);
+};
+#endif  // !defined(OS_ANDROID)
+
+}  // namespace media_router
+
+#endif  // CHROME_COMMON_MEDIA_ROUTER_TEST_TEST_HELPER_H_

chrome/test/BUILD.gn

@@ -2756,6 +2756,7 @@ test("unit_tests") {
     "//chrome:strings",
     "//chrome/browser/media/router:test_support",
     "//chrome/common:test_support",
+    "//chrome/common/media_router:test_support",
     "//components/autofill/content/renderer:test_support",
     "//components/browser_sync:test_support",
     "//components/component_updater:test_support",