Move common functionalities to BaseShardResponse

Signed-off-by: Aman Khare <[email protected]>

Author: Aman Khare

server/src/internalClusterTest/java/org/opensearch/gateway/RecoveryFromGatewayIT.java

@@ -812,15 +812,15 @@ public void testShardFetchCorruptedShardsUsingBatchAction() throws Exception {
             .get(discoveryNodes[0].getId())
             .getNodeGatewayStartedShardsBatch()
             .get(shardId);
-        assertNotNull(nodeGatewayStartedShards.storeException());
+        assertNotNull(nodeGatewayStartedShards.getException());
         assertNotNull(nodeGatewayStartedShards.allocationId());
         assertTrue(nodeGatewayStartedShards.primary());
     }
 
     private void assertNodeGatewayStartedShardsHappyCase(
         TransportNodesListGatewayStartedShardsBatch.NodeGatewayStartedShard nodeGatewayStartedShards
     ) {
-        assertNull(nodeGatewayStartedShards.storeException());
+        assertNull(nodeGatewayStartedShards.getException());
         assertNotNull(nodeGatewayStartedShards.allocationId());
         assertTrue(nodeGatewayStartedShards.primary());
     }

server/src/main/java/org/opensearch/gateway/BaseShardResponse.java

@@ -21,16 +21,34 @@
  * @opensearch.internal
  */
 public abstract class BaseShardResponse extends TransportResponse {
-    public BaseShardResponse() {}
+
+    private Exception storeException;
+
+    public BaseShardResponse(Exception storeException) {
+        this.storeException = storeException;
+    }
 
     public abstract boolean isEmpty();
 
-    public abstract Exception getException();
+    public Exception getException() {
+        return storeException;
+    }
 
     public BaseShardResponse(StreamInput in) throws IOException {
-        super(in);
+        if (in.readBoolean()) {
+            storeException = in.readException();
+        } else {
+            storeException = null;
+        }
     }
 
     @Override
-    public void writeTo(StreamOutput out) throws IOException {}
+    public void writeTo(StreamOutput out) throws IOException {
+        if (storeException != null) {
+            out.writeBoolean(true);
+            out.writeException(storeException);
+        } else {
+            out.writeBoolean(false);
+        }
+    }
 }

server/src/main/java/org/opensearch/gateway/TransportNodesListGatewayStartedShards.java

@@ -169,7 +169,7 @@ protected NodeGatewayStartedShards nodeOperation(NodeRequest request) {
                 shardInfo.allocationId(),
                 shardInfo.primary(),
                 shardInfo.replicationCheckpoint(),
-                shardInfo.storeException()
+                shardInfo.getException()
             );
         } catch (Exception e) {
             throw new OpenSearchException("failed to load started shards", e);

server/src/main/java/org/opensearch/gateway/TransportNodesListGatewayStartedShardsBatch.java

@@ -260,20 +260,20 @@ public void writeTo(StreamOutput out) throws IOException {
      *
      * @opensearch.internal
      */
-    public static class NodeGatewayStartedShard {
+    public static class NodeGatewayStartedShard extends BaseShardResponse {
         private final String allocationId;
         private final boolean primary;
-        private final Exception storeException;
         private final ReplicationCheckpoint replicationCheckpoint;
 
+        @Override
+        public boolean isEmpty() {
+            return allocationId == null && primary == false && getException() == null && replicationCheckpoint == null;
+        }
+
         public NodeGatewayStartedShard(StreamInput in) throws IOException {
+            super(in);
             allocationId = in.readOptionalString();
             primary = in.readBoolean();
-            if (in.readBoolean()) {
-                storeException = in.readException();
-            } else {
-                storeException = null;
-            }
             if (in.readBoolean()) {
                 replicationCheckpoint = new ReplicationCheckpoint(in);
             } else {
@@ -291,10 +291,10 @@ public NodeGatewayStartedShard(
             ReplicationCheckpoint replicationCheckpoint,
             Exception storeException
         ) {
+            super(storeException);
             this.allocationId = allocationId;
             this.primary = primary;
             this.replicationCheckpoint = replicationCheckpoint;
-            this.storeException = storeException;
         }
 
         public String allocationId() {
@@ -309,19 +309,10 @@ public ReplicationCheckpoint replicationCheckpoint() {
             return this.replicationCheckpoint;
         }
 
-        public Exception storeException() {
-            return this.storeException;
-        }
-
         public void writeTo(StreamOutput out) throws IOException {
+            super.writeTo(out);
             out.writeOptionalString(allocationId);
             out.writeBoolean(primary);
-            if (storeException != null) {
-                out.writeBoolean(true);
-                out.writeException(storeException);
-            } else {
-                out.writeBoolean(false);
-            }
             if (replicationCheckpoint != null) {
                 out.writeBoolean(true);
                 replicationCheckpoint.writeTo(out);
@@ -343,15 +334,15 @@ public boolean equals(Object o) {
 
             return primary == that.primary
                 && Objects.equals(allocationId, that.allocationId)
-                && Objects.equals(storeException, that.storeException)
+                && Objects.equals(getException(), that.getException())
                 && Objects.equals(replicationCheckpoint, that.replicationCheckpoint);
         }
 
         @Override
         public int hashCode() {
             int result = (allocationId != null ? allocationId.hashCode() : 0);
             result = 31 * result + (primary ? 1 : 0);
-            result = 31 * result + (storeException != null ? storeException.hashCode() : 0);
+            result = 31 * result + (getException() != null ? getException().hashCode() : 0);
             result = 31 * result + (replicationCheckpoint != null ? replicationCheckpoint.hashCode() : 0);
             return result;
         }
@@ -360,8 +351,8 @@ public int hashCode() {
         public String toString() {
             StringBuilder buf = new StringBuilder();
             buf.append("NodeGatewayStartedShards[").append("allocationId=").append(allocationId).append(",primary=").append(primary);
-            if (storeException != null) {
-                buf.append(",storeException=").append(storeException);
+            if (getException() != null) {
+                buf.append(",storeException=").append(getException());
             }
             if (replicationCheckpoint != null) {
                 buf.append(",ReplicationCheckpoint=").append(replicationCheckpoint.toString());
@@ -387,13 +378,28 @@ public Map<ShardId, NodeGatewayStartedShard> getNodeGatewayStartedShardsBatch()
 
         public NodeGatewayStartedShardsBatch(StreamInput in) throws IOException {
             super(in);
-            this.nodeGatewayStartedShardsBatch = in.readMap(ShardId::new, NodeGatewayStartedShard::new);
+            this.nodeGatewayStartedShardsBatch = in.readMap(ShardId::new,
+                i -> {
+                    if (i.readBoolean()) {
+                        return new NodeGatewayStartedShard(i);
+                    } else {
+                        return null;
+                    }
+                });
         }
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             super.writeTo(out);
-            out.writeMap(nodeGatewayStartedShardsBatch, (o, k) -> k.writeTo(o), (o, v) -> v.writeTo(o));
+            out.writeMap(nodeGatewayStartedShardsBatch, (o, k) -> k.writeTo(o),
+                (o, v) -> {
+                    if (v != null) {
+                        o.writeBoolean(true);
+                        v.writeTo(o);
+                    } else {
+                        o.writeBoolean(false);
+                    }
+                });
         }
 
         public NodeGatewayStartedShardsBatch(DiscoveryNode node, Map<ShardId, NodeGatewayStartedShard> nodeGatewayStartedShardsBatch) {