Apply review comments

linkedin · Oct 12, 2024 · 3032ab4 · 3032ab4
1 parent ef7c87c
commit 3032ab4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/storage/StorageService.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/storage/StorageService.java
@@ -397,23 +397,22 @@ public synchronized void checkWhetherStoragePartitionsShouldBeKeptOrNot(SafeHeli
           idealStatePartitionIds.add(RocksDBUtils.parsePartitionIdFromPartitionDbName(partitionDbName));
         });
 
-        Map<String, Map<String, String>> recordMapFields = idealState.getRecord().getMapFields();
-        for (String partitionDbName: recordMapFields.keySet()) {
+        Map<String, Map<String, String>> mapFields = idealState.getRecord().getMapFields();
+        for (Map.Entry<String, Map<String, String>> entry: mapFields.entrySet()) {
           boolean keepPartition = false;
+          String partitionDbName = entry.getKey();
           int partitionId = RocksDBUtils.parsePartitionIdFromPartitionDbName(partitionDbName);
           if (!storageEnginePartitionIds.contains(partitionId)) {
             continue;
           }
-          for (String hostName: recordMapFields.get(partitionDbName).keySet()) {
+          for (String hostName: entry.getValue().keySet()) {
             if (hostName.equals(listenerHostName)) {
               keepPartition = true;
               break;
             }
           }
           if (!keepPartition) {
             storageEngine.dropPartition(partitionId);
-            partitionSet.remove(partitionDbName);
-            recordMapFields.remove(partitionDbName);
           }
         }
       }

diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/storage/StorageServiceTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/storage/StorageServiceTest.java
@@ -155,7 +155,6 @@ public void testCheckWhetherStoragePartitionsShouldBeKeptOrNot() throws NoSuchFi
     when(abstractStorageEngine.getStoreVersionName()).thenReturn(resourceName);
     abstractStorageEngine.addStoragePartition(0);
     abstractStorageEngine.addStoragePartition(1);
-    abstractStorageEngine.addStoragePartition(2);
 
     String clusterName = "test_cluster";
     VeniceConfigLoader mockVeniceConfigLoader = mock(VeniceConfigLoader.class);
@@ -191,7 +190,7 @@ public void testCheckWhetherStoragePartitionsShouldBeKeptOrNot() throws NoSuchFi
     when(idealState.getRecord()).thenReturn(record);
     when(manager.getInstanceName()).thenReturn("lor1-app56586.prod.linkedin.com_1690");
 
-    Set<Integer> partitionSet = new HashSet<>(Arrays.asList(0, 1, 2));
+    Set<Integer> partitionSet = new HashSet<>(Arrays.asList(0, 1));
     when(abstractStorageEngine.getPartitionIds()).thenReturn(partitionSet);
     doAnswer(new Answer() {
       @Override
@@ -221,6 +220,6 @@ public Object answer(InvocationOnMock invocation) throws Throwable {
 
     doCallRealMethod().when(mockStorageService).checkWhetherStoragePartitionsShouldBeKeptOrNot(manager);
     mockStorageService.checkWhetherStoragePartitionsShouldBeKeptOrNot(manager);
-    Assert.assertFalse(idealState.getRecord().getMapFields().containsKey("test_store_v1_0"));
+    Assert.assertFalse(abstractStorageEngine.containsPartition(0));
   }
 }