Make write/read from io backword compatible

Signed-off-by: Siddhant Deshmukh <[email protected]>
opensearch-project · Sep 4, 2024 · 357d2cd · 357d2cd
1 parent 07ce827
commit 357d2cd
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 11 deletions.
diff --git a/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java b/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java
@@ -117,7 +117,7 @@ public QueryInsightsListener(
 
         clusterService.getClusterSettings()
             .addSettingsUpdateConsumer(
-                    TOP_N_QUERIES_MAX_GROUPS_EXCLUDING_N,
+                TOP_N_QUERIES_MAX_GROUPS_EXCLUDING_N,
                 v -> this.queryInsightsService.setMaximumGroups(v),
                 v -> this.queryInsightsService.validateMaximumGroups(v)
             );

diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java b/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java
@@ -235,7 +235,12 @@ public void setMaximumGroups(final int maxGroups) {
     public void validateMaximumGroups(final int maxGroups) {
         if (maxGroups < 0 || maxGroups > QueryInsightsSettings.MAX_GROUPS_EXCLUDING_TOPN_LIMIT) {
             throw new IllegalArgumentException(
-                "Max groups setting" + " should be between 0 and " + QueryInsightsSettings.MAX_GROUPS_EXCLUDING_TOPN_LIMIT + ", was (" + maxGroups + ")"
+                "Max groups setting"
+                    + " should be between 0 and "
+                    + QueryInsightsSettings.MAX_GROUPS_EXCLUDING_TOPN_LIMIT
+                    + ", was ("
+                    + maxGroups
+                    + ")"
             );
         }
     }

diff --git a/src/main/java/org/opensearch/plugin/insights/rules/model/MetricType.java b/src/main/java/org/opensearch/plugin/insights/rules/model/MetricType.java
@@ -96,4 +96,21 @@ public int compare(final Number a, final Number b) {
         }
         return -1;
     }
+
+    /**
+     * Parse a value with the correct type based on MetricType
+     *
+     * @param o the generic object to parse
+     * @return {@link Number}
+     */
+    Number parseValue(final Object o) {
+        switch (this) {
+            case LATENCY:
+            case CPU:
+            case MEMORY:
+                return (Long) o;
+            default:
+                return (Number) o;
+        }
+    }
 }
diff --git a/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java b/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java
@@ -9,9 +9,11 @@
 package org.opensearch.plugin.insights.rules.model;
 
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Objects;
+import org.opensearch.Version;
 import org.opensearch.core.common.Strings;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
@@ -41,11 +43,22 @@ public class SearchQueryRecord implements ToXContentObject, Writeable {
      */
     public SearchQueryRecord(final StreamInput in) throws IOException, ClassCastException {
         this.timestamp = in.readLong();
-        measurements = new LinkedHashMap<>();
-        in.readOrderedMap(MetricType::readFromStream, Measurement::readFromStream)
-            .forEach(((metricType, measurement) -> measurements.put(metricType, measurement)));
+        if (in.getVersion().onOrAfter(Version.V_2_17_0)) {
+            measurements = new LinkedHashMap<>();
+            in.readOrderedMap(MetricType::readFromStream, Measurement::readFromStream)
+                .forEach(((metricType, measurement) -> measurements.put(metricType, measurement)));
+            this.groupingId = null;
+        } else {
+            measurements = new HashMap<>();
+            in.readMap(MetricType::readFromStream, StreamInput::readGenericValue).forEach((metricType, o) -> {
+                try {
+                    measurements.put(metricType, new Measurement(metricType.parseValue(o)));
+                } catch (ClassCastException e) {
+                    throw new ClassCastException("Error parsing value for metric type: " + metricType);
+                }
+            });
+        }
         this.attributes = Attribute.readAttributeMap(in);
-        this.groupingId = null;
     }
 
     /**
@@ -155,11 +168,15 @@ public XContentBuilder toXContent(final XContentBuilder builder, final ToXConten
     @Override
     public void writeTo(final StreamOutput out) throws IOException {
         out.writeLong(timestamp);
-        out.writeMap(
-            measurements,
-            (stream, metricType) -> MetricType.writeTo(out, metricType),
-            (stream, measurement) -> measurement.writeTo(out)
-        );
+        if (out.getVersion().onOrAfter(Version.V_2_17_0)) {
+            out.writeMap(
+                measurements,
+                (stream, metricType) -> MetricType.writeTo(out, metricType),
+                (stream, measurement) -> measurement.writeTo(out)
+            );
+        } else {
+            out.writeMap(measurements, (stream, metricType) -> MetricType.writeTo(out, metricType), StreamOutput::writeGenericValue);
+        }
         out.writeMap(
             attributes,
             (stream, attribute) -> Attribute.writeTo(out, attribute),