1 files changed, 45 insertions, 28 deletions
diff --git a/core/src/main/java/org/elasticsearch/cluster/ClusterModule.java b/core/src/main/java/org/elasticsearch/cluster/ClusterModule.java
index 581844e28f..34cf7df6a2 100644
--- a/core/src/main/java/org/elasticsearch/cluster/ClusterModule.java
+++ b/core/src/main/java/org/elasticsearch/cluster/ClusterModule.java
@@ -56,16 +56,19 @@ import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.inject.AbstractModule;
 import org.elasticsearch.common.logging.ESLogger;
 import org.elasticsearch.common.logging.Loggers;
+import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Setting.Property;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.ExtensionPoint;
 import org.elasticsearch.gateway.GatewayAllocator;
+import org.elasticsearch.plugins.ClusterPlugin;
 import org.elasticsearch.tasks.TaskResultsService;
 
-import java.util.Arrays;
-import java.util.Collections;
+import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.function.Function;
 
 /**
@@ -77,48 +80,27 @@ public class ClusterModule extends AbstractModule {
     public static final String BALANCED_ALLOCATOR = "balanced"; // default
     public static final Setting<String> SHARDS_ALLOCATOR_TYPE_SETTING =
         new Setting<>("cluster.routing.allocation.type", BALANCED_ALLOCATOR, Function.identity(), Property.NodeScope);
-    public static final List<Class<? extends AllocationDecider>> DEFAULT_ALLOCATION_DECIDERS =
-        Collections.unmodifiableList(Arrays.asList(
-            MaxRetryAllocationDecider.class,
-            SameShardAllocationDecider.class,
-            FilterAllocationDecider.class,
-            ReplicaAfterPrimaryActiveAllocationDecider.class,
-            ThrottlingAllocationDecider.class,
-            RebalanceOnlyWhenActiveAllocationDecider.class,
-            ClusterRebalanceAllocationDecider.class,
-            ConcurrentRebalanceAllocationDecider.class,
-            EnableAllocationDecider.class,
-            AwarenessAllocationDecider.class,
-            ShardsLimitAllocationDecider.class,
-            NodeVersionAllocationDecider.class,
-            DiskThresholdDecider.class,
-            SnapshotInProgressAllocationDecider.class));
 
     private final Settings settings;
     private final ExtensionPoint.SelectedType<ShardsAllocator> shardsAllocators = new ExtensionPoint.SelectedType<>("shards_allocator", ShardsAllocator.class);
-    private final ExtensionPoint.ClassSet<AllocationDecider> allocationDeciders = new ExtensionPoint.ClassSet<>("allocation_decider", AllocationDecider.class, AllocationDeciders.class);
     private final ExtensionPoint.ClassSet<IndexTemplateFilter> indexTemplateFilters = new ExtensionPoint.ClassSet<>("index_template_filter", IndexTemplateFilter.class);
     private final ClusterService clusterService;
     private final IndexNameExpressionResolver indexNameExpressionResolver;
+    // pkg private for tests
+    final Collection<AllocationDecider> allocationDeciders;
 
     // pkg private so tests can mock
     Class<? extends ClusterInfoService> clusterInfoServiceImpl = InternalClusterInfoService.class;
 
-    public ClusterModule(Settings settings, ClusterService clusterService) {
+    public ClusterModule(Settings settings, ClusterService clusterService, List<ClusterPlugin> clusterPlugins) {
         this.settings = settings;
-        for (Class<? extends AllocationDecider> decider : ClusterModule.DEFAULT_ALLOCATION_DECIDERS) {
-            registerAllocationDecider(decider);
-        }
+        this.allocationDeciders = createAllocationDeciders(settings, clusterService.getClusterSettings(), clusterPlugins);
         registerShardsAllocator(ClusterModule.BALANCED_ALLOCATOR, BalancedShardsAllocator.class);
         registerShardsAllocator(ClusterModule.EVEN_SHARD_COUNT_ALLOCATOR, BalancedShardsAllocator.class);
         this.clusterService = clusterService;
         indexNameExpressionResolver = new IndexNameExpressionResolver(settings);
     }
 
-    public void registerAllocationDecider(Class<? extends AllocationDecider> allocationDecider) {
-        allocationDeciders.registerExtension(allocationDecider);
-    }
-
     public void registerShardsAllocator(String name, Class<? extends ShardsAllocator> clazz) {
         shardsAllocators.registerExtension(name, clazz);
     }
@@ -131,6 +113,41 @@ public class ClusterModule extends AbstractModule {
         return indexNameExpressionResolver;
     }
 
+    // TODO: this is public so allocation benchmark can access the default deciders...can we do that in another way?
+    /** Return a new {@link AllocationDecider} instance with builtin deciders as well as those from plugins. */
+    public static Collection<AllocationDecider> createAllocationDeciders(Settings settings, ClusterSettings clusterSettings,
+                                                                         List<ClusterPlugin> clusterPlugins) {
+        // collect deciders by class so that we can detect duplicates
+        Map<Class, AllocationDecider> deciders = new HashMap<>();
+        addAllocationDecider(deciders, new MaxRetryAllocationDecider(settings));
+        addAllocationDecider(deciders, new SameShardAllocationDecider(settings));
+        addAllocationDecider(deciders, new FilterAllocationDecider(settings, clusterSettings));
+        addAllocationDecider(deciders, new ReplicaAfterPrimaryActiveAllocationDecider(settings));
+        addAllocationDecider(deciders, new ThrottlingAllocationDecider(settings, clusterSettings));
+        addAllocationDecider(deciders, new RebalanceOnlyWhenActiveAllocationDecider(settings));
+        addAllocationDecider(deciders, new ClusterRebalanceAllocationDecider(settings, clusterSettings));
+        addAllocationDecider(deciders, new ConcurrentRebalanceAllocationDecider(settings, clusterSettings));
+        addAllocationDecider(deciders, new EnableAllocationDecider(settings, clusterSettings));
+        addAllocationDecider(deciders, new AwarenessAllocationDecider(settings, clusterSettings));
+        addAllocationDecider(deciders, new ShardsLimitAllocationDecider(settings, clusterSettings));
+        addAllocationDecider(deciders, new NodeVersionAllocationDecider(settings));
+        addAllocationDecider(deciders, new DiskThresholdDecider(settings, clusterSettings));
+        addAllocationDecider(deciders, new SnapshotInProgressAllocationDecider(settings, clusterSettings));
+
+        clusterPlugins.stream()
+            .flatMap(p -> p.createAllocationDeciders(settings, clusterSettings).stream())
+            .forEach(d -> addAllocationDecider(deciders, d));
+
+        return deciders.values();
+    }
+
+    /** Add the given allocation decider to the given deciders collection, erroring if the class name is already used. */
+    private static void addAllocationDecider(Map<Class, AllocationDecider> deciders, AllocationDecider decider) {
+        if (deciders.put(decider.getClass(), decider) != null) {
+            throw new IllegalArgumentException("Cannot specify allocation decider [" + decider.getClass().getName() + "] twice");
+        }
+    }
+
     @Override
     protected void configure() {
         // bind ShardsAllocator
@@ -139,7 +156,6 @@ public class ClusterModule extends AbstractModule {
             final ESLogger logger = Loggers.getLogger(getClass(), settings);
             logger.warn("{} allocator has been removed in 2.0 using {} instead", ClusterModule.EVEN_SHARD_COUNT_ALLOCATOR, ClusterModule.BALANCED_ALLOCATOR);
         }
-        allocationDeciders.bind(binder());
         indexTemplateFilters.bind(binder());
 
         bind(ClusterInfoService.class).to(clusterInfoServiceImpl).asEagerSingleton();
@@ -161,5 +177,6 @@ public class ClusterModule extends AbstractModule {
         bind(NodeMappingRefreshAction.class).asEagerSingleton();
         bind(MappingUpdatedAction.class).asEagerSingleton();
         bind(TaskResultsService.class).asEagerSingleton();
+        bind(AllocationDeciders.class).toInstance(new AllocationDeciders(settings, allocationDeciders));
     }
 }