Use RestoreClusterStateListener with generic type that extends ActionResponse

Signed-off-by: Sachin Kale <kalsac@amazon.com>

Author: Sachin Kale

server/src/main/java/org/opensearch/action/admin/cluster/remotestore/restore/RemoteStoreRestoreClusterStateListener.java

@@ -9,6 +9,7 @@
 package org.opensearch.action.admin.cluster.remotestore.restore;
 
 import org.opensearch.action.ActionListener;
+import org.opensearch.action.admin.cluster.snapshots.restore.RestoreClusterStateListener;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeAction;
 import org.opensearch.cluster.ClusterState;
@@ -78,7 +79,7 @@ protected ClusterBlockException checkBlock(RestoreRemoteStoreRequest request, Cl
     }
 
     @Override
-    protected void masterOperation(
+    protected void clusterManagerOperation(
         final RestoreRemoteStoreRequest request,
         final ClusterState state,
         final ActionListener<RestoreRemoteStoreResponse> listener
@@ -87,10 +88,11 @@ protected void masterOperation(
             request,
             ActionListener.delegateFailure(listener, (delegatedListener, restoreCompletionResponse) -> {
                 if (restoreCompletionResponse.getRestoreInfo() == null && request.waitForCompletion()) {
-                    RemoteStoreRestoreClusterStateListener.createAndRegisterListener(
+                    RestoreClusterStateListener.createAndRegisterListener(
                         clusterService,
                         restoreCompletionResponse,
-                        delegatedListener
+                        delegatedListener,
+                        RestoreRemoteStoreResponse::new
                     );
                 } else {
                     delegatedListener.onResponse(new RestoreRemoteStoreResponse(restoreCompletionResponse.getRestoreInfo()));

server/src/main/java/org/opensearch/action/admin/cluster/remotestore/restore/TransportRestoreRemoteStoreAction.java

@@ -35,6 +35,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.action.ActionListener;
+import org.opensearch.action.ActionResponse;
 import org.opensearch.cluster.ClusterChangedEvent;
 import org.opensearch.cluster.ClusterStateListener;
 import org.opensearch.cluster.RestoreInProgress;
@@ -44,29 +45,36 @@
 import org.opensearch.snapshots.RestoreInfo;
 import org.opensearch.snapshots.RestoreService;
 
+import java.util.function.Function;
+
 import static org.opensearch.snapshots.RestoreService.restoreInProgress;
 
 /**
  * Transport listener for cluster state updates
  *
  * @opensearch.internal
  */
-public class RestoreClusterStateListener implements ClusterStateListener {
+public class RestoreClusterStateListener<T extends ActionResponse> implements ClusterStateListener {
 
     private static final Logger logger = LogManager.getLogger(RestoreClusterStateListener.class);
 
     private final ClusterService clusterService;
     private final String uuid;
-    private final ActionListener<RestoreSnapshotResponse> listener;
+    private final String restoreIdentifier;
+    private final ActionListener<T> listener;
+    private final Function<RestoreInfo, T> actionResponseFactory;
 
     private RestoreClusterStateListener(
         ClusterService clusterService,
         RestoreService.RestoreCompletionResponse response,
-        ActionListener<RestoreSnapshotResponse> listener
+        ActionListener<T> listener,
+        Function<RestoreInfo, T> actionResponseFactory
     ) {
         this.clusterService = clusterService;
         this.uuid = response.getUuid();
+        this.restoreIdentifier = response.getSnapshot() != null ? response.getSnapshot().getSnapshotId().getName() : "remote_store";
         this.listener = listener;
+        this.actionResponseFactory = actionResponseFactory;
     }
 
     @Override
@@ -78,35 +86,36 @@ public void clusterChanged(ClusterChangedEvent changedEvent) {
             // on the current cluster-manager and as such it might miss some intermediary cluster states due to batching.
             // Clean up listener in that case and acknowledge completion of restore operation to client.
             clusterService.removeListener(this);
-            listener.onResponse(new RestoreSnapshotResponse((RestoreInfo) null));
+            listener.onResponse(actionResponseFactory.apply(null));
         } else if (newEntry == null) {
             clusterService.removeListener(this);
             ImmutableOpenMap<ShardId, RestoreInProgress.ShardRestoreStatus> shards = prevEntry.shards();
-            assert prevEntry.state().completed() : "expected completed snapshot state but was " + prevEntry.state();
+            assert prevEntry.state().completed() : "expected completed snapshot/remote store restore state but was " + prevEntry.state();
             assert RestoreService.completed(shards) : "expected all restore entries to be completed";
             RestoreInfo ri = new RestoreInfo(
-                prevEntry.snapshot().getSnapshotId().getName(),
+                restoreIdentifier,
                 prevEntry.indices(),
                 shards.size(),
                 shards.size() - RestoreService.failedShards(shards)
             );
-            RestoreSnapshotResponse response = new RestoreSnapshotResponse(ri);
-            logger.debug("restore of [{}] completed", prevEntry.snapshot().getSnapshotId());
+            T response = actionResponseFactory.apply(ri);
+            logger.debug("restore of [{}] completed", restoreIdentifier);
             listener.onResponse(response);
         } else {
-            // restore not completed yet, wait for next cluster state update
+            logger.debug("restore not completed yet, wait for next cluster state update");
         }
     }
 
     /**
      * Creates a cluster state listener and registers it with the cluster service. The listener passed as a
      * parameter will be called when the restore is complete.
      */
-    public static void createAndRegisterListener(
+    public static <T extends ActionResponse> void createAndRegisterListener(
         ClusterService clusterService,
         RestoreService.RestoreCompletionResponse response,
-        ActionListener<RestoreSnapshotResponse> listener
+        ActionListener<T> listener,
+        Function<RestoreInfo, T> actionResponseFactory
     ) {
-        clusterService.addListener(new RestoreClusterStateListener(clusterService, response, listener));
+        clusterService.addListener(new RestoreClusterStateListener<T>(clusterService, response, listener, actionResponseFactory));
     }
 }

server/src/main/java/org/opensearch/action/admin/cluster/snapshots/restore/RestoreClusterStateListener.java

@@ -109,7 +109,12 @@ protected void clusterManagerOperation(
     ) {
         restoreService.restoreSnapshot(request, ActionListener.delegateFailure(listener, (delegatedListener, restoreCompletionResponse) -> {
             if (restoreCompletionResponse.getRestoreInfo() == null && request.waitForCompletion()) {
-                RestoreClusterStateListener.createAndRegisterListener(clusterService, restoreCompletionResponse, delegatedListener);
+                RestoreClusterStateListener.createAndRegisterListener(
+                    clusterService,
+                    restoreCompletionResponse,
+                    delegatedListener,
+                    RestoreSnapshotResponse::new
+                );
             } else {
                 delegatedListener.onResponse(new RestoreSnapshotResponse(restoreCompletionResponse.getRestoreInfo()));
             }