From 3afe59ff69828306868eb9931c3fb74b3b1269e7 Mon Sep 17 00:00:00 2001
From: David Pilato <david@pilato.fr>
Date: Wed, 14 Jan 2015 16:59:42 +0100
Subject: [PATCH] Move AzureUnicastHostsProvider to AzureDiscoveryModule

Related to elasticsearch/elasticsearch#9099

Closes #48
---
 .../discovery/azure/AzureDiscovery.java       | 30 ++-----------------
 .../discovery/azure/AzureDiscoveryModule.java |  4 +++
 2 files changed, 6 insertions(+), 28 deletions(-)

diff --git a/src/main/java/org/elasticsearch/discovery/azure/AzureDiscovery.java b/src/main/java/org/elasticsearch/discovery/azure/AzureDiscovery.java
index f7383d830c3..9d3391b0874 100755
--- a/src/main/java/org/elasticsearch/discovery/azure/AzureDiscovery.java
+++ b/src/main/java/org/elasticsearch/discovery/azure/AzureDiscovery.java
@@ -19,22 +19,16 @@
 
 package org.elasticsearch.discovery.azure;
 
-import org.elasticsearch.Version;
-import org.elasticsearch.cloud.azure.AzureComputeService;
 import org.elasticsearch.cluster.ClusterName;
 import org.elasticsearch.cluster.ClusterService;
 import org.elasticsearch.cluster.node.DiscoveryNodeService;
 import org.elasticsearch.cluster.settings.DynamicSettings;
-import org.elasticsearch.common.collect.ImmutableList;
 import org.elasticsearch.common.inject.Inject;
-import org.elasticsearch.common.network.NetworkService;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.discovery.DiscoverySettings;
 import org.elasticsearch.discovery.zen.ZenDiscovery;
 import org.elasticsearch.discovery.zen.elect.ElectMasterService;
-import org.elasticsearch.discovery.zen.ping.ZenPing;
 import org.elasticsearch.discovery.zen.ping.ZenPingService;
-import org.elasticsearch.discovery.zen.ping.unicast.UnicastZenPing;
 import org.elasticsearch.node.settings.NodeSettingsService;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
@@ -49,29 +43,9 @@ public class AzureDiscovery extends ZenDiscovery {
     @Inject
     public AzureDiscovery(Settings settings, ClusterName clusterName, ThreadPool threadPool, TransportService transportService,
                           ClusterService clusterService, NodeSettingsService nodeSettingsService, ZenPingService pingService,
-                          DiscoveryNodeService discoveryNodeService, AzureComputeService azureService, NetworkService networkService,
-                          DiscoverySettings discoverySettings, ElectMasterService electMasterService, DynamicSettings dynamicSettings,
-                          Version version) {
+                          DiscoveryNodeService discoveryNodeService,
+                          DiscoverySettings discoverySettings, ElectMasterService electMasterService, DynamicSettings dynamicSettings) {
         super(settings, clusterName, threadPool, transportService, clusterService, nodeSettingsService,
                 discoveryNodeService, pingService, electMasterService, discoverySettings, dynamicSettings);
-        if (settings.getAsBoolean("cloud.enabled", true)) {
-            ImmutableList<? extends ZenPing> zenPings = pingService.zenPings();
-            UnicastZenPing unicastZenPing = null;
-            for (ZenPing zenPing : zenPings) {
-                if (zenPing instanceof UnicastZenPing) {
-                    unicastZenPing = (UnicastZenPing) zenPing;
-                    break;
-                }
-            }
-
-            if (unicastZenPing != null) {
-                // update the unicast zen ping to add cloud hosts provider
-                // and, while we are at it, use only it and not the multicast for example
-                unicastZenPing.addHostsProvider(new AzureUnicastHostsProvider(settings, azureService, transportService, networkService, version));
-                pingService.zenPings(ImmutableList.of(unicastZenPing));
-            } else {
-                logger.warn("failed to apply azure unicast discovery, no unicast ping found");
-            }
-        }
     }
 }
diff --git a/src/main/java/org/elasticsearch/discovery/azure/AzureDiscoveryModule.java b/src/main/java/org/elasticsearch/discovery/azure/AzureDiscoveryModule.java
index 490ee61719f..d66fa0a1138 100644
--- a/src/main/java/org/elasticsearch/discovery/azure/AzureDiscoveryModule.java
+++ b/src/main/java/org/elasticsearch/discovery/azure/AzureDiscoveryModule.java
@@ -40,7 +40,11 @@ public class AzureDiscoveryModule extends ZenDiscoveryModule {
         super();
         this.logger = Loggers.getLogger(getClass(), settings);
         this.settings = settings;
+        if (AzureModule.isDiscoveryReady(settings, logger)) {
+            addUnicastHostProvider(AzureUnicastHostsProvider.class);
+        }
     }
+
     @Override
     protected void bindDiscovery() {
         if (AzureModule.isDiscoveryReady(settings, logger)) {