opendistro-for-elasticsearch · kaituo · Apr 29, 2020 · Apr 28, 2020 · Apr 29, 2020
diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/ad/AnomalyDetectorPlugin.java b/src/main/java/com/amazon/opendistroforelasticsearch/ad/AnomalyDetectorPlugin.java
@@ -72,6 +72,7 @@
 import com.amazon.opendistroforelasticsearch.ad.transport.ThresholdResultTransportAction;
 import com.amazon.opendistroforelasticsearch.ad.transport.handler.AnomalyResultHandler;
 import com.amazon.opendistroforelasticsearch.ad.util.ClientUtil;
+import com.amazon.opendistroforelasticsearch.ad.util.ClusterStateUtils;
 import com.amazon.opendistroforelasticsearch.ad.util.ColdStartRunner;
 import com.amazon.opendistroforelasticsearch.ad.util.IndexUtils;
 import com.amazon.opendistroforelasticsearch.ad.util.Throttler;
@@ -91,6 +92,7 @@
 import org.elasticsearch.cluster.metadata.MetaData.Custom;
 import org.elasticsearch.cluster.node.DiscoveryNodes;
 import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.IndexScopedSettings;
@@ -121,6 +123,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Supplier;
@@ -145,6 +148,7 @@ public class AnomalyDetectorPlugin extends Plugin implements ActionPlugin, Scrip
     private ADStats adStats;
     private NamedXContentRegistry xContentRegistry;
     private ClientUtil clientUtil;
+    private ClusterStateUtils clusterStateUtils;
 
     static {
         SpecialPermission.check();
@@ -201,7 +205,11 @@ public List<RestHandler> getRestHandlers(
             clusterService,
             anomalyDetectorRunner
         );
-        RestStatsAnomalyDetectorAction statsAnomalyDetectorAction = new RestStatsAnomalyDetectorAction(restController, adStats);
+        RestStatsAnomalyDetectorAction statsAnomalyDetectorAction = new RestStatsAnomalyDetectorAction(
+            restController,
+            adStats,
+            this.clusterStateUtils
+        );
         RestAnomalyDetectorJobAction anomalyDetectorJobAction = new RestAnomalyDetectorJobAction(
             settings,
             restController,
@@ -259,8 +267,11 @@ public Collection<Object> createComponents(
         RandomCutForestSerDe rcfSerde = new RandomCutForestSerDe();
         CheckpointDao checkpoint = new CheckpointDao(client, clientUtil, CommonName.CHECKPOINT_INDEX_NAME);
 
+        HashMap<String, String> ignoredAttributes = new HashMap<>();
+        ignoredAttributes.put(CommonName.BOX_TYPE_KEY, CommonName.WARM_BOX_TYPE);
+        this.clusterStateUtils = new ClusterStateUtils(clusterService, ignoredAttributes);
         ModelManager modelManager = new ModelManager(
-            clusterService,
+            clusterStateUtils,
             jvmService,
             rcfSerde,
             checkpoint,
@@ -284,7 +295,7 @@ public Collection<Object> createComponents(
             AnomalyDetectorSettings.SHINGLE_SIZE
         );
 
-        HashRing hashRing = new HashRing(clusterService, clock, settings);
+        HashRing hashRing = new HashRing(clusterStateUtils, clock, settings);
         ADStateManager stateManager = new ADStateManager(
             client,
             xContentRegistry,
@@ -353,11 +364,12 @@ public Collection<Object> createComponents(
                 clock,
                 stateManager,
                 runner,
-                new ADClusterEventListener(clusterService, hashRing, modelManager),
+                new ADClusterEventListener(clusterService, hashRing, modelManager, clusterStateUtils),
                 deleteUtil,
                 adCircuitBreakerService,
                 adStats,
-                new MasterEventListener(clusterService, threadPool, deleteUtil, client, clock, clientUtil)
+                new MasterEventListener(clusterService, threadPool, deleteUtil, client, clock, clientUtil, clusterStateUtils),
+                clusterStateUtils
             );
     }
 

diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/ad/cluster/ADClusterEventListener.java b/src/main/java/com/amazon/opendistroforelasticsearch/ad/cluster/ADClusterEventListener.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License").
  * You may not use this file except in compliance with the License.
@@ -20,6 +20,8 @@
 import java.util.concurrent.Semaphore;
 
 import com.amazon.opendistroforelasticsearch.ad.ml.ModelManager;
+import com.amazon.opendistroforelasticsearch.ad.util.ClusterStateUtils;
+
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.cluster.ClusterChangedEvent;
@@ -32,7 +34,7 @@
 
 public class ADClusterEventListener implements ClusterStateListener {
     private static final Logger LOG = LogManager.getLogger(ADClusterEventListener.class);
-    static final String MASTER_NOT_APPLIED_MSG = "AD does not use master nodes";
+    static final String NODE_NOT_APPLIED_MSG = "AD does not use master or ultrawarm nodes";
     static final String NOT_RECOVERED_MSG = "CLuster is not recovered yet.";
     static final String IN_PROGRESS_MSG = "Cluster state change in progress, return.";
     static final String REMOVE_MODEL_MSG = "Remove model";
@@ -43,21 +45,28 @@ public class ADClusterEventListener implements ClusterStateListener {
     private HashRing hashRing;
     private ModelManager modelManager;
     private final ClusterService clusterService;
+    private final ClusterStateUtils clusterStateUtils;
 
     @Inject
-    public ADClusterEventListener(ClusterService clusterService, HashRing hashRing, ModelManager modelManager) {
+    public ADClusterEventListener(
+        ClusterService clusterService,
+        HashRing hashRing,
+        ModelManager modelManager,
+        ClusterStateUtils clusterStateUtils
+    ) {
         this.clusterService = clusterService;
         this.clusterService.addListener(this);
         this.hashRing = hashRing;
         this.modelManager = modelManager;
         this.inProgress = new Semaphore(1);
+        this.clusterStateUtils = clusterStateUtils;
     }
 
     @Override
     public void clusterChanged(ClusterChangedEvent event) {
 
-        if (!event.state().nodes().getLocalNode().isDataNode()) {
-            LOG.debug(MASTER_NOT_APPLIED_MSG);
+        if (clusterStateUtils.isIgnoredNode(event.state().nodes().getLocalNode())) {
+            LOG.debug(NODE_NOT_APPLIED_MSG);
             return;
         }
 
@@ -79,7 +88,7 @@ public void clusterChanged(ClusterChangedEvent event) {
             // Check whether it was a data node that was removed
             boolean dataNodeRemoved = false;
             for (DiscoveryNode removedNode : delta.removedNodes()) {
-                if (removedNode.isDataNode()) {
+                if (!clusterStateUtils.isIgnoredNode(removedNode)) {
                     LOG.info(NODE_REMOVED_MSG + " {}", removedNode.getId());
                     dataNodeRemoved = true;
                     break;
@@ -89,7 +98,7 @@ public void clusterChanged(ClusterChangedEvent event) {
             // Check whether it was a data node that was added
             boolean dataNodeAdded = false;
             for (DiscoveryNode addedNode : delta.addedNodes()) {
-                if (addedNode.isDataNode()) {
+                if (!clusterStateUtils.isIgnoredNode(addedNode)) {
                     LOG.info(NODE_ADDED_MSG + " {}", addedNode.getId());
                     dataNodeAdded = true;
                     break;

diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/ad/cluster/HashRing.java b/src/main/java/com/amazon/opendistroforelasticsearch/ad/cluster/HashRing.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License").
  * You may not use this file except in compliance with the License.
@@ -28,10 +28,10 @@
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.routing.Murmur3HashFunction;
-import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.TimeValue;
 
+import com.amazon.opendistroforelasticsearch.ad.util.ClusterStateUtils;
 import com.carrotsearch.hppc.cursors.ObjectCursor;
 
 public class HashRing {
@@ -43,7 +43,7 @@ public class HashRing {
     static final String COOLDOWN_MSG = "Hash ring doesn't respond to cluster state change within the cooldown period.";
 
     private final int VIRTUAL_NODE_COUNT = 100;
-    private final ClusterService clusterService;
+    private final ClusterStateUtils clusterStateUtils;
     private TreeMap<Integer, DiscoveryNode> circle;
     private Semaphore inProgress;
     // the UTC epoch milliseconds of the most recent successful update
@@ -52,9 +52,9 @@ public class HashRing {
     private final Clock clock;
     private AtomicBoolean membershipChangeRequied;
 
-    public HashRing(ClusterService clusterService, Clock clock, Settings settings) {
+    public HashRing(ClusterStateUtils clusterStateUtils, Clock clock, Settings settings) {
         this.circle = new TreeMap<Integer, DiscoveryNode>();
-        this.clusterService = clusterService;
+        this.clusterStateUtils = clusterStateUtils;
         this.inProgress = new Semaphore(1);
         this.clock = clock;
         this.coolDownPeriod = COOLDOWN_MINUTES.get(settings);
@@ -93,7 +93,7 @@ public boolean build() {
         TreeMap<Integer, DiscoveryNode> newCircle = new TreeMap<>();
 
         try {
-            for (ObjectCursor<DiscoveryNode> cursor : clusterService.state().nodes().getDataNodes().values()) {
+            for (ObjectCursor<DiscoveryNode> cursor : clusterStateUtils.getEligibleDataNodes().values()) {
                 DiscoveryNode curNode = cursor.value;
                 for (int i = 0; i < VIRTUAL_NODE_COUNT; i++) {
                     newCircle.put(Murmur3HashFunction.hash(curNode.getId() + i), curNode);

diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/ad/cluster/HourlyCron.java b/src/main/java/com/amazon/opendistroforelasticsearch/ad/cluster/HourlyCron.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License").
  * You may not use this file except in compliance with the License.
@@ -21,27 +21,27 @@
 import org.elasticsearch.action.FailedNodeException;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.cluster.node.DiscoveryNode;
-import org.elasticsearch.cluster.service.ClusterService;
 
 import com.amazon.opendistroforelasticsearch.ad.transport.CronAction;
 import com.amazon.opendistroforelasticsearch.ad.transport.CronRequest;
+import com.amazon.opendistroforelasticsearch.ad.util.ClusterStateUtils;
 
 public class HourlyCron implements Runnable {
     private static final Logger LOG = LogManager.getLogger(HourlyCron.class);
     static final String SUCCEEDS_LOG_MSG = "Hourly maintenance succeeds";
     static final String NODE_EXCEPTION_LOG_MSG = "Hourly maintenance of node has exception";
     static final String EXCEPTION_LOG_MSG = "Hourly maintenance has exception.";
-    private ClusterService clusterService;
+    private ClusterStateUtils clusterStateUtils;
     private Client client;
 
-    public HourlyCron(ClusterService clusterService, Client client) {
-        this.clusterService = clusterService;
+    public HourlyCron(Client client, ClusterStateUtils clusterStateUtils) {
+        this.clusterStateUtils = clusterStateUtils;
         this.client = client;
     }
 
     @Override
     public void run() {
-        DiscoveryNode[] dataNodes = clusterService.state().nodes().getDataNodes().values().toArray(DiscoveryNode.class);
+        DiscoveryNode[] dataNodes = clusterStateUtils.getEligibleDataNodes().values().toArray(DiscoveryNode.class);
 
         // we also add the cancel query function here based on query text from the negative cache.
 

diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/ad/cluster/MasterEventListener.java b/src/main/java/com/amazon/opendistroforelasticsearch/ad/cluster/MasterEventListener.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License").
  * You may not use this file except in compliance with the License.
@@ -26,6 +26,7 @@
 
 import com.amazon.opendistroforelasticsearch.ad.settings.AnomalyDetectorSettings;
 import com.amazon.opendistroforelasticsearch.ad.util.ClientUtil;
+import com.amazon.opendistroforelasticsearch.ad.util.ClusterStateUtils;
 
 import org.elasticsearch.threadpool.ThreadPool;
 
@@ -39,14 +40,16 @@ public class MasterEventListener implements LocalNodeMasterListener {
     private Client client;
     private Clock clock;
     private ClientUtil clientUtil;
+    private ClusterStateUtils clusterStateUtils;
 
     public MasterEventListener(
         ClusterService clusterService,
         ThreadPool threadPool,
         DeleteDetector deleteUtil,
         Client client,
         Clock clock,
-        ClientUtil clientUtil
+        ClientUtil clientUtil,
+        ClusterStateUtils clusterStateUtils
     ) {
         this.clusterService = clusterService;
         this.threadPool = threadPool;
@@ -55,13 +58,14 @@ public MasterEventListener(
         this.clusterService.addLocalNodeMasterListener(this);
         this.clock = clock;
         this.clientUtil = clientUtil;
+        this.clusterStateUtils = clusterStateUtils;
     }
 
     @Override
     public void onMaster() {
         if (hourlyCron == null) {
             hourlyCron = threadPool
-                .scheduleWithFixedDelay(new HourlyCron(clusterService, client), TimeValue.timeValueHours(1), executorName());
+                .scheduleWithFixedDelay(new HourlyCron(client, clusterStateUtils), TimeValue.timeValueHours(1), executorName());
             clusterService.addLifecycleListener(new LifecycleListener() {
                 @Override
                 public void beforeStop() {

diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/ad/constant/CommonName.java b/src/main/java/com/amazon/opendistroforelasticsearch/ad/constant/CommonName.java
@@ -31,4 +31,14 @@ public class CommonName {
     // Anomaly Detector name for X-Opaque-Id header
     // ======================================
     public static final String ANOMALY_DETECTOR = "[Anomaly Detector]";
+
+    // ======================================
+    // Ultrawarm node attributes
+    // ======================================
+
+    // warm node
+    public static String WARM_BOX_TYPE = "warm";
+
+    // box type
+    public static final String BOX_TYPE_KEY = "box_type";
 }
diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/ad/ml/ModelManager.java b/src/main/java/com/amazon/opendistroforelasticsearch/ad/ml/ModelManager.java
@@ -42,14 +42,14 @@
 import com.google.gson.Gson;
 
 import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.monitor.jvm.JvmService;
 
 import com.amazon.opendistroforelasticsearch.ad.common.exception.LimitExceededException;
 import com.amazon.opendistroforelasticsearch.ad.common.exception.ResourceNotFoundException;
 import com.amazon.opendistroforelasticsearch.ad.constant.CommonErrorMessages;
 import com.amazon.opendistroforelasticsearch.ad.ml.rcf.CombinedRcfResult;
 import com.amazon.opendistroforelasticsearch.ad.model.AnomalyDetector;
+import com.amazon.opendistroforelasticsearch.ad.util.ClusterStateUtils;
 import com.amazon.randomcutforest.RandomCutForest;
 import com.amazon.randomcutforest.serialize.RandomCutForestSerDe;
 
@@ -103,7 +103,7 @@ public String getName() {
     private final Duration checkpointInterval;
 
     // dependencies
-    private final ClusterService clusterService;
+    private final ClusterStateUtils clusterStateUtils;
     private final JvmService jvmService;
     private final RandomCutForestSerDe rcfSerde;
     private final CheckpointDao checkpointDao;
@@ -119,7 +119,7 @@ public String getName() {
     /**
      * Constructor.
      *
-     * @param clusterService cluster info
+     * @param clusterStateUtils cluster info
      * @param jvmService jvm info
      * @param rcfSerde RCF model serialization
      * @param checkpointDao model checkpoint storage
@@ -140,9 +140,10 @@ public String getName() {
      * @param minPreviewSize minimum number of data points for preview
      * @param modelTtl time to live for hosted models
      * @param checkpointInterval interval between checkpoints
+     * @param shingleSize size of a shingle
      */
     public ModelManager(
-        ClusterService clusterService,
+        ClusterStateUtils clusterStateUtils,
         JvmService jvmService,
         RandomCutForestSerDe rcfSerde,
         CheckpointDao checkpointDao,
@@ -166,7 +167,7 @@ public ModelManager(
         int shingleSize
     ) {
 
-        this.clusterService = clusterService;
+        this.clusterStateUtils = clusterStateUtils;
         this.jvmService = jvmService;
         this.rcfSerde = rcfSerde;
         this.checkpointDao = checkpointDao;
@@ -259,7 +260,7 @@ public Entry<Integer, Integer> getPartitionedForestSizes(RandomCutForest forest,
         int numPartitions = (int) Math.ceil((double) totalSize / (double) partitionSize);
         int forestSize = (int) Math.ceil((double) forest.getNumberOfTrees() / (double) numPartitions);
 
-        int numNodes = clusterService.state().nodes().getDataNodes().size();
+        int numNodes = clusterStateUtils.getEligibleDataNodes().size();
         if (numPartitions > numNodes) {
             // partition by cluster size
             partitionSize = (long) Math.ceil((double) totalSize / (double) numNodes);