From 8ba2782810f3dda60e0c3656f0db71d77f295b3d Mon Sep 17 00:00:00 2001
From: Anshu Agarwal <anshukag@amazon.com>
Date: Tue, 2 Jan 2024 10:04:28 +0530
Subject: [PATCH] Fix for deserialization bug in weighted round robin metadata

Signed-off-by: Anshu Agarwal <anshukag@amazon.com>
---
 .../metadata/WeightedRoutingMetadata.java     |  7 ++--
 .../WeightedRoutingMetadataTests.java         | 41 ++++++++++++++++---
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/server/src/main/java/org/opensearch/cluster/metadata/WeightedRoutingMetadata.java b/server/src/main/java/org/opensearch/cluster/metadata/WeightedRoutingMetadata.java
index bc24dd22f5c6e..b303c3a2034d5 100644
--- a/server/src/main/java/org/opensearch/cluster/metadata/WeightedRoutingMetadata.java
+++ b/server/src/main/java/org/opensearch/cluster/metadata/WeightedRoutingMetadata.java
@@ -26,6 +26,7 @@
 import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Objects;
 
 /**
  * Contains metadata for weighted routing
@@ -99,7 +100,7 @@ public static NamedDiff<Metadata.Custom> readDiffFrom(StreamInput in) throws IOE
     public static WeightedRoutingMetadata fromXContent(XContentParser parser) throws IOException {
         String attrKey = null;
         Double attrValue;
-        String attributeName = null;
+        String attributeName = "";
         Map<String, Double> weights = new HashMap<>();
         WeightedRouting weightedRouting;
         XContentParser.Token token;
@@ -162,12 +163,12 @@ public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         WeightedRoutingMetadata that = (WeightedRoutingMetadata) o;
-        return weightedRouting.equals(that.weightedRouting);
+        return weightedRouting.equals(that.weightedRouting) && version == that.version;
     }
 
     @Override
     public int hashCode() {
-        return weightedRouting.hashCode();
+        return Objects.hash(weightedRouting.hashCode(), version);
     }
 
     @Override
diff --git a/server/src/test/java/org/opensearch/cluster/metadata/WeightedRoutingMetadataTests.java b/server/src/test/java/org/opensearch/cluster/metadata/WeightedRoutingMetadataTests.java
index b78d1b56364eb..e19bde5d53d8a 100644
--- a/server/src/test/java/org/opensearch/cluster/metadata/WeightedRoutingMetadataTests.java
+++ b/server/src/test/java/org/opensearch/cluster/metadata/WeightedRoutingMetadataTests.java
@@ -8,29 +8,60 @@
 
 package org.opensearch.cluster.metadata;
 
+import org.opensearch.cluster.ClusterModule;
+import org.opensearch.cluster.Diff;
 import org.opensearch.cluster.routing.WeightedRouting;
+import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
+import org.opensearch.core.common.io.stream.Writeable;
 import org.opensearch.core.xcontent.XContentParser;
-import org.opensearch.test.AbstractXContentTestCase;
+import org.opensearch.test.AbstractDiffableSerializationTestCase;
 
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.Map;
 
-public class WeightedRoutingMetadataTests extends AbstractXContentTestCase<WeightedRoutingMetadata> {
+public class WeightedRoutingMetadataTests extends AbstractDiffableSerializationTestCase<Metadata.Custom> {
+
+    @Override
+    protected Writeable.Reader<Metadata.Custom> instanceReader() {
+        return WeightedRoutingMetadata::new;
+    }
+
     @Override
     protected WeightedRoutingMetadata createTestInstance() {
+        String attributeName = "zone";
         Map<String, Double> weights = Map.of("a", 1.0, "b", 1.0, "c", 0.0);
-        WeightedRouting weightedRouting = new WeightedRouting("zone", weights);
+        if (randomBoolean()) {
+            weights = new HashMap<>();
+            attributeName = "";
+        }
+        WeightedRouting weightedRouting = new WeightedRouting(attributeName, weights);
         WeightedRoutingMetadata weightedRoutingMetadata = new WeightedRoutingMetadata(weightedRouting, -1);
+
         return weightedRoutingMetadata;
     }
 
+    @Override
+    protected NamedWriteableRegistry getNamedWriteableRegistry() {
+        return new NamedWriteableRegistry(ClusterModule.getNamedWriteables());
+    }
+
     @Override
     protected WeightedRoutingMetadata doParseInstance(XContentParser parser) throws IOException {
         return WeightedRoutingMetadata.fromXContent(parser);
     }
 
     @Override
-    protected boolean supportsUnknownFields() {
-        return false;
+    protected Metadata.Custom makeTestChanges(Metadata.Custom testInstance) {
+
+        WeightedRouting weightedRouting = new WeightedRouting("", new HashMap<>());
+        WeightedRoutingMetadata weightedRoutingMetadata = new WeightedRoutingMetadata(weightedRouting, -1);
+        return weightedRoutingMetadata;
     }
+
+    @Override
+    protected Writeable.Reader<Diff<Metadata.Custom>> diffReader() {
+        return WeightedRoutingMetadata::readDiffFrom;
+    }
+
 }