Change ServiceManager::SetLoaderForUrl() to SetLoaderForScheme()

BUG=None
R=viettrungluu@chromium.org, viettrungluu

Review URL: https://codereview.chromium.org/214513009

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

3 files changed, 100 insertions, 26 deletions

diff --git a/mojo/service_manager/service_manager.cc b/mojo/service_manager/service_manager.cc
index fe8575d..8b80825 100644
--- a/mojo/service_manager/service_manager.cc
+++ b/mojo/service_manager/service_manager.cc
@@ -32,6 +32,7 @@ class ServiceManager::ServiceFactory : public Shell, public ErrorHandler {
                                                 url,
                                                 pipe.handle_to_self.Pass());
   }
+
   virtual ~ServiceFactory() {}
 
   void ConnectToClient(ScopedMessagePipeHandle handle) {
@@ -73,7 +74,7 @@ ServiceManager::ServiceManager() : default_loader_(NULL) {
 }
 
 ServiceManager::~ServiceManager() {
-  for (ServiceFactoryMap::iterator it = url_to_service_factory_.begin();
+  for (URLToServiceFactoryMap::iterator it = url_to_service_factory_.begin();
        it != url_to_service_factory_.end(); ++it) {
     delete it->second;
   }
@@ -88,22 +89,9 @@ ServiceManager* ServiceManager::GetInstance() {
   return &instance.Get();
 }
 
-void ServiceManager::SetLoaderForURL(ServiceLoader* loader, const GURL& gurl) {
-  DCHECK(url_to_loader_.find(gurl) == url_to_loader_.end());
-  url_to_loader_[gurl] = loader;
-}
-
-ServiceLoader* ServiceManager::GetLoaderForURL(const GURL& gurl) {
-  LoaderMap::const_iterator it = url_to_loader_.find(gurl);
-  if (it != url_to_loader_.end())
-    return it->second;
-  DCHECK(default_loader_);
-  return default_loader_;
-}
-
 void ServiceManager::Connect(const GURL& url,
                              ScopedMessagePipeHandle client_handle) {
-  ServiceFactoryMap::const_iterator service_it =
+  URLToServiceFactoryMap::const_iterator service_it =
       url_to_service_factory_.find(url);
   ServiceFactory* service_factory;
   if (service_it != url_to_service_factory_.end()) {
@@ -115,9 +103,32 @@ void ServiceManager::Connect(const GURL& url,
   service_factory->ConnectToClient(client_handle.Pass());
 }
 
+void ServiceManager::SetLoaderForURL(ServiceLoader* loader, const GURL& url) {
+  DCHECK(url_to_loader_.find(url) == url_to_loader_.end());
+  url_to_loader_[url] = loader;
+}
+
+void ServiceManager::SetLoaderForScheme(ServiceLoader* loader,
+                                        const std::string& scheme) {
+  DCHECK(scheme_to_loader_.find(scheme) == scheme_to_loader_.end());
+  scheme_to_loader_[scheme] = loader;
+}
+
+ServiceLoader* ServiceManager::GetLoaderForURL(const GURL& url) {
+  URLToLoaderMap::const_iterator url_it = url_to_loader_.find(url);
+  if (url_it != url_to_loader_.end())
+    return url_it->second;
+  SchemeToLoaderMap::const_iterator scheme_it =
+      scheme_to_loader_.find(url.scheme());
+  if (scheme_it != scheme_to_loader_.end())
+    return scheme_it->second;
+  DCHECK(default_loader_);
+  return default_loader_;
+}
+
 void ServiceManager::OnServiceFactoryError(ServiceFactory* service_factory) {
   const GURL url = service_factory->url();
-  ServiceFactoryMap::iterator it = url_to_service_factory_.find(url);
+  URLToServiceFactoryMap::iterator it = url_to_service_factory_.find(url);
   DCHECK(it != url_to_service_factory_.end());
   delete it->second;
   url_to_service_factory_.erase(it);
diff --git a/mojo/service_manager/service_manager.h b/mojo/service_manager/service_manager.h
index cf4c385..513c792 100644
--- a/mojo/service_manager/service_manager.h
+++ b/mojo/service_manager/service_manager.h
@@ -45,28 +45,40 @@ class MOJO_SERVICE_MANAGER_EXPORT ServiceManager {
   // Returns a shared instance, creating it if necessary.
   static ServiceManager* GetInstance();
 
-  // Sets the default Loader to be used if not overridden by SetLoaderForURL().
+  // Loads a service if necessary and establishes a new client connection.
+  void Connect(const GURL& url, ScopedMessagePipeHandle client_handle);
+
+  // Sets the default Loader to be used if not overridden by
+  // SetLoaderForURL() or SetLoaderForScheme().
   // Does not take ownership of |loader|.
   void set_default_loader(ServiceLoader* loader) { default_loader_ = loader; }
   // Sets a Loader to be used for a specific url.
   // Does not take ownership of |loader|.
-  void SetLoaderForURL(ServiceLoader* loader, const GURL& gurl);
-  // Returns the Loader to use for a url (using default if not overridden.)
-  ServiceLoader* GetLoaderForURL(const GURL& gurl);
-  // Loads a service if necessary and establishes a new client connection.
-  void Connect(const GURL& url, ScopedMessagePipeHandle client_handle);
+  void SetLoaderForURL(ServiceLoader* loader, const GURL& url);
+  // Sets a Loader to be used for a specific url scheme.
+  // Does not take ownership of |loader|.
+  void SetLoaderForScheme(ServiceLoader* loader, const std::string& scheme);
 
  private:
   class ServiceFactory;
+  typedef std::map<std::string, ServiceLoader*> SchemeToLoaderMap;
+  typedef std::map<GURL, ServiceLoader*> URLToLoaderMap;
+  typedef std::map<GURL, ServiceFactory*> URLToServiceFactoryMap;
+
+  // Returns the Loader to use for a url (using default if not overridden.)
+  // The preference is to use a loader that's been specified for an url first,
+  // then one that's been specified for a scheme, then the default.
+  ServiceLoader* GetLoaderForURL(const GURL& url);
 
   // Removes a ServiceFactory when it no longer has any connections.
   void OnServiceFactoryError(ServiceFactory* service_factory);
 
+  // Loader management.
+  URLToLoaderMap url_to_loader_;
+  SchemeToLoaderMap scheme_to_loader_;
   ServiceLoader* default_loader_;
-  typedef std::map<GURL, ServiceFactory*> ServiceFactoryMap;
-  ServiceFactoryMap url_to_service_factory_;
-  typedef std::map<GURL, ServiceLoader*> LoaderMap;
-  LoaderMap url_to_loader_;
+
+  URLToServiceFactoryMap url_to_service_factory_;
   DISALLOW_COPY_AND_ASSIGN(ServiceManager);
 };
 
diff --git a/mojo/service_manager/service_manager_unittest.cc b/mojo/service_manager/service_manager_unittest.cc
index 26273e7..4d467da 100644
--- a/mojo/service_manager/service_manager_unittest.cc
+++ b/mojo/service_manager/service_manager_unittest.cc
@@ -124,6 +124,24 @@ class ServiceManagerTest : public testing::Test, public ServiceLoader {
   DISALLOW_COPY_AND_ASSIGN(ServiceManagerTest);
 };
 
+class TestServiceLoader : public ServiceLoader {
+ public:
+  TestServiceLoader() : num_loads_(0) {}
+  int num_loads() const { return num_loads_; }
+
+ private:
+  virtual void LoadService(ServiceManager* manager,
+                           const GURL& url,
+                           ScopedShellHandle service_handle) OVERRIDE {
+    ++num_loads_;
+  }
+  virtual void OnServiceError(ServiceManager* manager, const GURL& url)
+      OVERRIDE {}
+
+  int num_loads_;
+  DISALLOW_COPY_AND_ASSIGN(TestServiceLoader);
+};
+
 TEST_F(ServiceManagerTest, Basic) {
   test_client_->Test("test");
   loop_.Run();
@@ -140,4 +158,37 @@ TEST_F(ServiceManagerTest, ClientError) {
   EXPECT_EQ(0, context_.num_impls);
   EXPECT_FALSE(HasFactoryForTestURL());
 }
+
+// Confirm that both urls and schemes can have their loaders explicitly set.
+TEST_F(ServiceManagerTest, SetLoaders) {
+  ServiceManager sm;
+  TestServiceLoader default_loader;
+  TestServiceLoader url_loader;
+  TestServiceLoader scheme_loader;
+  sm.set_default_loader(&default_loader);
+  sm.SetLoaderForURL(&url_loader, GURL("test:test1"));
+  sm.SetLoaderForScheme(&scheme_loader, "test");
+
+  // test::test1 should go to url_loader.
+  InterfacePipe<TestService, AnyInterface> pipe1;
+  sm.Connect(GURL("test:test1"), pipe1.handle_to_peer.Pass());
+  EXPECT_EQ(1, url_loader.num_loads());
+  EXPECT_EQ(0, scheme_loader.num_loads());
+  EXPECT_EQ(0, default_loader.num_loads());
+
+  // test::test2 should go to scheme loader.
+  InterfacePipe<TestService, AnyInterface> pipe2;
+  sm.Connect(GURL("test:test2"), pipe2.handle_to_peer.Pass());
+  EXPECT_EQ(1, url_loader.num_loads());
+  EXPECT_EQ(1, scheme_loader.num_loads());
+  EXPECT_EQ(0, default_loader.num_loads());
+
+  // http::test1 should go to default loader.
+  InterfacePipe<TestService, AnyInterface> pipe3;
+  sm.Connect(GURL("http:test1"), pipe3.handle_to_peer.Pass());
+  EXPECT_EQ(1, url_loader.num_loads());
+  EXPECT_EQ(1, scheme_loader.num_loads());
+  EXPECT_EQ(1, default_loader.num_loads());
+}
+
 }  // namespace mojo