fix shingle memory calculation

src/main/java/com/amazon/opendistroforelasticsearch/ad/task/ADTaskCacheManager.java

@@ -27,6 +27,7 @@
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 
+import com.amazon.opendistroforelasticsearch.ad.model.AnomalyDetector;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.settings.Settings;
 
@@ -146,7 +147,7 @@ public boolean isThresholdModelTrained(String taskId) {
      * @param taskId task id
      * @param trained threshold model trained or not
      */
-    public void setThresholdModelTrained(String taskId, boolean trained) {
+    protected void setThresholdModelTrained(String taskId, boolean trained) {
         ADBatchTaskCache taskCache = getBatchTaskCache(taskId);
         taskCache.setThresholdModelTrained(trained);
         if (trained) {
@@ -209,8 +210,9 @@ private ADBatchTaskCache getBatchTaskCache(String taskId) {
      * @return how many bytes will consume
      */
     private long calculateADTaskCacheSize(ADTask adTask) {
-        return memoryTracker.estimateModelSize(adTask.getDetector(), NUM_TREES) + trainingDataMemorySize(THRESHOLD_MODEL_TRAINING_SIZE)
-            + shingleMemorySize(adTask.getDetector().getShingleSize());
+        AnomalyDetector detector = adTask.getDetector();
+        return memoryTracker.estimateModelSize(detector, NUM_TREES) + trainingDataMemorySize(THRESHOLD_MODEL_TRAINING_SIZE)
+            + shingleMemorySize(detector.getShingleSize(), detector.getEnabledFeatureIds().size());
     }
 
     /**
@@ -297,15 +299,18 @@ public long trainingDataMemorySize(int size) {
 
     /**
      * Estimate max memory usage of shingle data.
-     * Based on the test, one shingle data point consumes about 96 bytes.
-     * The shingle data is stored in {@link java.util.Deque}
+     * One feature aggregated data point(double) consumes 8 bytes.
+     * The shingle data is stored in {@link java.util.Deque}. From testing,
+     * other parts except feature data consume 80 bytes.
+     *
      * Check {@link ADBatchTaskCache#getShingle()}
      *
      * @param shingleSize shingle data point count
+     * @param enabledFeatureSize enabled feature count
      * @return how many bytes will consume
      */
-    public long shingleMemorySize(int shingleSize) {
-        return 96 * shingleSize;
+    public long shingleMemorySize(int shingleSize, int enabledFeatureSize) {
+        return (80 + 8 * enabledFeatureSize) * shingleSize;
     }
 
 }