Address comments

Signed-off-by: Arpit Bandejiya <abandeji@amazon.com>
opensearch-project · Sep 2, 2024 · 50c5d7b · 50c5d7b
1 parent 5d40f9f
commit 50c5d7b
Show file tree

Hide file tree

Showing 8 changed files with 49 additions and 11 deletions.
diff --git a/...nalClusterTest/java/org/opensearch/gateway/remote/RemoteClusterStateCleanupManagerIT.java b/...nalClusterTest/java/org/opensearch/gateway/remote/RemoteClusterStateCleanupManagerIT.java
@@ -254,9 +254,8 @@ private void verifyIndexRoutingFilesDeletion(
         assertNotNull(discoveryStats.getClusterStateStats());
         for (PersistedStateStats persistedStateStats : discoveryStats.getClusterStateStats().getPersistenceStats()) {
             Map<String, AtomicLong> extendedFields = persistedStateStats.getExtendedFields();
-            assertTrue(extendedFields.containsKey(RemotePersistenceStats.INDEX_ROUTING_FILES_CLEANUP_ATTEMPT_FAILED_COUNT));
-            long cleanupAttemptFailedCount = extendedFields.get(RemotePersistenceStats.INDEX_ROUTING_FILES_CLEANUP_ATTEMPT_FAILED_COUNT)
-                .get();
+            assertTrue(extendedFields.containsKey(RemoteUploadStats.INDEX_ROUTING_FILES_CLEANUP_ATTEMPT_FAILED_COUNT));
+            long cleanupAttemptFailedCount = extendedFields.get(RemoteUploadStats.INDEX_ROUTING_FILES_CLEANUP_ATTEMPT_FAILED_COUNT).get();
             assertEquals(0, cleanupAttemptFailedCount);
         }
     }

diff --git a/server/src/main/java/org/opensearch/cluster/coordination/PublicationTransportHandler.java b/server/src/main/java/org/opensearch/cluster/coordination/PublicationTransportHandler.java
@@ -262,6 +262,7 @@ PublishWithJoinResponse handleIncomingRemotePublishRequest(RemotePublishRequest
             }
 
             if (applyFullState == true) {
+                remoteClusterStateService.fullDownloadState();
                 logger.debug(
                     () -> new ParameterizedMessage(
                         "Downloading full cluster state for term {}, version {}, stateUUID {}",
@@ -281,6 +282,7 @@ PublishWithJoinResponse handleIncomingRemotePublishRequest(RemotePublishRequest
                 lastSeenClusterState.set(clusterState);
                 return response;
             } else {
+                remoteClusterStateService.diffDownloadState();
                 logger.debug(
                     () -> new ParameterizedMessage(
                         "Downloading diff cluster state for term {}, version {}, previousUUID {}, current UUID {}",

diff --git a/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateCleanupManager.java b/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateCleanupManager.java
@@ -81,7 +81,7 @@ public RemoteClusterStateCleanupManager(
         RemoteRoutingTableService remoteRoutingTableService
     ) {
         this.remoteClusterStateService = remoteClusterStateService;
-        this.remoteStateStats = remoteClusterStateService.getStats();
+        this.remoteStateStats = remoteClusterStateService.getRemoteStateStats();
         ClusterSettings clusterSettings = clusterService.getClusterSettings();
         this.clusterApplierService = clusterService.getClusterApplierService();
         this.staleFileCleanupInterval = clusterSettings.get(REMOTE_CLUSTER_STATE_CLEANUP_INTERVAL_SETTING);

diff --git a/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateService.java b/server/src/main/java/org/opensearch/gateway/remote/RemoteClusterStateService.java
@@ -1662,6 +1662,18 @@ public void readMetadataFailed() {
         remoteStateStats.stateDownloadFailed();
     }
 
+    public void fullDownloadState() {
+        remoteStateStats.fullDownloadState();
+    }
+
+    public void diffDownloadState() {
+        remoteStateStats.diffDownloadState();
+    }
+
+    public RemotePersistenceStats getRemoteStateStats() {
+        return remoteStateStats;
+    }
+
     public PersistedStateStats getUploadStats() {
         return remoteStateStats.getUploadStats();
     }

diff --git a/server/src/main/java/org/opensearch/gateway/remote/RemoteDownloadStats.java b/server/src/main/java/org/opensearch/gateway/remote/RemoteDownloadStats.java
@@ -10,10 +10,27 @@
 
 import org.opensearch.cluster.coordination.PersistedStateStats;
 
+import java.util.concurrent.atomic.AtomicLong;
+
 public class RemoteDownloadStats extends PersistedStateStats {
     static final String REMOTE_DOWNLOAD = "remote_download";
+    static final String DIFF_DOWNLOAD = "diff_download";
+    private AtomicLong diffDownloadCount = new AtomicLong(0);
+    static final String FULL_DOWNLOAD = "full_download";
+    private AtomicLong fullDownloadCount = new AtomicLong(0);
 
     public RemoteDownloadStats() {
         super(REMOTE_DOWNLOAD);
+        addToExtendedFields(DIFF_DOWNLOAD, diffDownloadCount);
+        addToExtendedFields(FULL_DOWNLOAD, fullDownloadCount);
+    }
+
+    public void diffDownloadState() {
+        diffDownloadCount.incrementAndGet();
     }
+
+    public void fullDownloadState() {
+        fullDownloadCount.incrementAndGet();
+    }
+
 }
diff --git a/server/src/main/java/org/opensearch/gateway/remote/RemotePersistenceStats.java b/server/src/main/java/org/opensearch/gateway/remote/RemotePersistenceStats.java
@@ -81,4 +81,12 @@ public PersistedStateStats getDownloadStats() {
         return remoteDownloadStats;
     }
 
+    public void diffDownloadState() {
+        remoteDownloadStats.diffDownloadState();
+    }
+
+    public void fullDownloadState() {
+        remoteDownloadStats.fullDownloadState();
+    }
+
 }
diff --git a/...er/src/test/java/org/opensearch/gateway/remote/RemoteClusterStateCleanupManagerTests.java b/...er/src/test/java/org/opensearch/gateway/remote/RemoteClusterStateCleanupManagerTests.java
@@ -144,7 +144,7 @@ public void setup() {
         remoteManifestManager = mock(RemoteManifestManager.class);
         remoteClusterStateService = mock(RemoteClusterStateService.class);
         when(remoteClusterStateService.getRemoteManifestManager()).thenReturn(remoteManifestManager);
-        when(remoteClusterStateService.getStats()).thenReturn(new RemotePersistenceStats());
+        when(remoteClusterStateService.getRemoteStateStats()).thenReturn(new RemotePersistenceStats());
         when(remoteClusterStateService.getThreadpool()).thenReturn(threadPool);
         when(remoteClusterStateService.getBlobStore()).thenReturn(blobStore);
         when(remoteClusterStateService.getBlobStoreRepository()).thenReturn(blobStoreRepository);

diff --git a/server/src/test/java/org/opensearch/gateway/remote/RemoteClusterStateServiceTests.java b/server/src/test/java/org/opensearch/gateway/remote/RemoteClusterStateServiceTests.java
@@ -575,7 +575,7 @@ public void testWriteFullMetadataInParallelFailureForIndexMetadata() throws IOEx
             RemoteStateTransferException.class,
             () -> remoteClusterStateService.writeFullMetadata(clusterState, randomAlphaOfLength(10), MANIFEST_CURRENT_CODEC_VERSION)
         );
-        assertEquals(0, remoteClusterStateService.getStats().getSuccessCount());
+        assertEquals(0, remoteClusterStateService.getUploadStats().getSuccessCount());
     }
 
     public void testFailWriteIncrementalMetadataNonClusterManagerNode() throws IOException {
@@ -587,7 +587,7 @@ public void testFailWriteIncrementalMetadataNonClusterManagerNode() throws IOExc
             null
         );
         Assert.assertThat(manifestDetails, nullValue());
-        assertEquals(0, remoteClusterStateService.getStats().getSuccessCount());
+        assertEquals(0, remoteClusterStateService.getUploadStats().getSuccessCount());
     }
 
     public void testFailWriteIncrementalMetadataWhenTermChanged() {
@@ -2600,10 +2600,10 @@ public void testRemoteStateStats() throws IOException {
             MANIFEST_CURRENT_CODEC_VERSION
         ).getClusterMetadataManifest();
 
-        assertTrue(remoteClusterStateService.getStats() != null);
-        assertEquals(1, remoteClusterStateService.getStats().getSuccessCount());
-        assertEquals(0, remoteClusterStateService.getStats().getCleanupAttemptFailedCount());
-        assertEquals(0, remoteClusterStateService.getStats().getFailedCount());
+        assertTrue(remoteClusterStateService.getUploadStats() != null);
+        assertEquals(1, remoteClusterStateService.getUploadStats().getSuccessCount());
+        assertEquals(0, remoteClusterStateService.getRemoteStateStats().getCleanupAttemptFailedCount());
+        assertEquals(0, remoteClusterStateService.getUploadStats().getFailedCount());
     }
 
     public void testRemoteRoutingTableNotInitializedWhenDisabled() {