Expose zen ElectMasterService as a Discovery extension point

Some users need to override ElectMasterService from zen to add functionality
inside their infrastructure. This commit allows to extend it.

core/src/main/java/org/elasticsearch/discovery/DiscoveryModule.java

@@ -41,15 +41,18 @@ import java.util.Map;
 public class DiscoveryModule extends AbstractModule {
 
     public static final String DISCOVERY_TYPE_KEY = "discovery.type";
+    public static final String ZEN_MASTER_SERVICE_TYPE_KEY = "discovery.zen.masterservice.type";
 
     private final Settings settings;
     private final List<Class<? extends UnicastHostsProvider>> unicastHostProviders = Lists.newArrayList();
     private final Map<String, Class<? extends Discovery>> discoveryTypes = new HashMap<>();
+    private final Map<String, Class<? extends ElectMasterService>> masterServiceType = new HashMap<>();
 
     public DiscoveryModule(Settings settings) {
         this.settings = settings;
         addDiscoveryType("local", LocalDiscovery.class);
         addDiscoveryType("zen", ZenDiscovery.class);
+        addElectMasterService("zen", ElectMasterService.class);
     }
 
     /**
@@ -63,9 +66,22 @@ public class DiscoveryModule extends AbstractModule {
      * Adds a custom Discovery type.
      */
     public void addDiscoveryType(String type, Class<? extends Discovery> clazz) {
+        if (discoveryTypes.containsKey(type)) {
+            throw new IllegalArgumentException("discovery type [" + type + "] is already registered");
+        }
         discoveryTypes.put(type, clazz);
     }
 
+    /**
+     * Adds a custom zen master service type.
+     */
+    public void addElectMasterService(String type, Class<? extends ElectMasterService> masterService) {
+        if (masterServiceType.containsKey(type)) {
+            throw new IllegalArgumentException("master service type [" + type + "] is already registered");
+        }
+        this.masterServiceType.put(type, masterService);
+    }
+
     @Override
     protected void configure() {
         String defaultType = DiscoveryNode.localNode(settings) ? "local" : "zen";
@@ -76,7 +92,16 @@ public class DiscoveryModule extends AbstractModule {
         }
 
         if (discoveryType.equals("local") == false) {
-            bind(ElectMasterService.class).asEagerSingleton();
+            String masterServiceTypeKey = settings.get(ZEN_MASTER_SERVICE_TYPE_KEY, "zen");
+            final Class<? extends ElectMasterService> masterService = masterServiceType.get(masterServiceTypeKey);
+            if (masterService == null) {
+                throw new IllegalArgumentException("Unknown master service type [" + masterServiceTypeKey + "]");
+            }
+            if (masterService == ElectMasterService.class) {
+                bind(ElectMasterService.class).asEagerSingleton();
+            } else {
+                bind(ElectMasterService.class).to(masterService).asEagerSingleton();
+            }
             bind(ZenPingService.class).asEagerSingleton();
             Multibinder<UnicastHostsProvider> unicastHostsProviderMultibinder = Multibinder.newSetBinder(binder(), UnicastHostsProvider.class);
             for (Class<? extends UnicastHostsProvider> unicastHostProvider : unicastHostProviders) {

core/src/test/java/org/elasticsearch/common/inject/ModuleTestCase.java

@@ -1,4 +1,3 @@
-
 /*
  * Licensed to Elasticsearch under one or more contributor
  * license agreements. See the NOTICE file distributed with
@@ -68,7 +67,7 @@ public abstract class ModuleTestCase extends ESTestCase {
             fail("Expected exception from configuring module. Found these bindings:\n" + s);
         } catch (IllegalArgumentException e) {
             for (String msg : msgs) {
-                assertTrue(e.getMessage().contains(msg));
+                assertTrue(e.getMessage() + " didn't contain: " + msg, e.getMessage().contains(msg));
             }
         }
     }

core/src/test/java/org/elasticsearch/discovery/DiscoveryModuleTests.java

@@ -0,0 +1,149 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.discovery;
+
+import org.elasticsearch.Version;
+import org.elasticsearch.cluster.ClusterChangedEvent;
+import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.cluster.routing.RoutingService;
+import org.elasticsearch.common.Nullable;
+import org.elasticsearch.common.component.Lifecycle;
+import org.elasticsearch.common.component.LifecycleListener;
+import org.elasticsearch.common.inject.ModuleTestCase;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.discovery.local.LocalDiscovery;
+import org.elasticsearch.discovery.zen.ZenDiscovery;
+import org.elasticsearch.discovery.zen.elect.ElectMasterService;
+import org.elasticsearch.node.service.NodeService;
+
+/**
+ */
+public class DiscoveryModuleTests extends ModuleTestCase {
+
+    public static class DummyMasterElectionService extends ElectMasterService {
+
+        public DummyMasterElectionService(Settings settings, Version version) {
+            super(settings, version);
+        }
+    }
+
+
+    public void testRegisterMasterElectionService() {
+        Settings settings = Settings.builder().put("node.local", false).
+                put(DiscoveryModule.ZEN_MASTER_SERVICE_TYPE_KEY, "custom").build();
+        DiscoveryModule module = new DiscoveryModule(settings);
+        module.addElectMasterService("custom", DummyMasterElectionService.class);
+        assertBinding(module, ElectMasterService.class, DummyMasterElectionService.class);
+        assertBinding(module, Discovery.class, ZenDiscovery.class);
+    }
+
+    public void testLoadUnregisteredMasterElectionService() {
+        Settings settings = Settings.builder().put("node.local", false).
+                put(DiscoveryModule.ZEN_MASTER_SERVICE_TYPE_KEY, "foobar").build();
+        DiscoveryModule module = new DiscoveryModule(settings);
+        module.addElectMasterService("custom", DummyMasterElectionService.class);
+        assertBindingFailure(module, "Unknown master service type [foobar]");
+    }
+
+    public void testRegisterDefaults() {
+        boolean local = randomBoolean();
+        Settings settings = Settings.builder().put("node.local", local).build();
+        DiscoveryModule module = new DiscoveryModule(settings);
+        assertBinding(module, Discovery.class, local ? LocalDiscovery.class : ZenDiscovery.class);
+    }
+
+    public void testRegisterDiscovery() {
+        boolean local = randomBoolean();
+        Settings settings = Settings.builder().put("node.local", local).
+                put(DiscoveryModule.DISCOVERY_TYPE_KEY, "custom").build();
+        DiscoveryModule module = new DiscoveryModule(settings);
+        module.addDiscoveryType("custom", DummyDisco.class);
+        assertBinding(module, Discovery.class, DummyDisco.class);
+    }
+
+
+    public static class DummyDisco implements Discovery {
+
+
+        @Override
+        public DiscoveryNode localNode() {
+            return null;
+        }
+
+        @Override
+        public void addListener(InitialStateDiscoveryListener listener) {
+
+        }
+
+        @Override
+        public void removeListener(InitialStateDiscoveryListener listener) {
+
+        }
+
+        @Override
+        public String nodeDescription() {
+            return null;
+        }
+
+        @Override
+        public void setNodeService(@Nullable NodeService nodeService) {
+
+        }
+
+        @Override
+        public void setRoutingService(RoutingService routingService) {
+
+        }
+
+        @Override
+        public void publish(ClusterChangedEvent clusterChangedEvent, AckListener ackListener) {
+
+        }
+
+        @Override
+        public Lifecycle.State lifecycleState() {
+            return null;
+        }
+
+        @Override
+        public void addLifecycleListener(LifecycleListener listener) {
+
+        }
+
+        @Override
+        public void removeLifecycleListener(LifecycleListener listener) {
+
+        }
+
+        @Override
+        public Discovery start() {
+            return null;
+        }
+
+        @Override
+        public Discovery stop() {
+            return null;
+        }
+
+        @Override
+        public void close() {
+
+        }
+    }
+}