HDFS-15362. FileWithSnapshotFeature#updateQuotaAndCollectBlocks should collect all distinct blocks. Contributed by hemanthboyina.

hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/server/namenode/snapshot/FileWithSnapshotFeature.java

@@ -17,9 +17,10 @@
  */
 package org.apache.hadoop.hdfs.server.namenode.snapshot;
 
-import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.fs.StorageType;
@@ -156,7 +157,8 @@ public void updateQuotaAndCollectBlocks(INode.ReclaimContext reclaimContext,
     QuotaCounts oldCounts;
     if (removed.snapshotINode != null) {
       oldCounts = new QuotaCounts.Builder().build();
-      List<BlockInfo> allBlocks = new ArrayList<BlockInfo>();
+      // collect all distinct blocks
+      Set<BlockInfo> allBlocks = new HashSet<BlockInfo>();
       if (file.getBlocks() != null) {
         allBlocks.addAll(Arrays.asList(file.getBlocks()));
       }

hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/namenode/snapshot/TestFileWithSnapshotFeature.java

@@ -30,11 +30,13 @@
 import org.apache.hadoop.test.Whitebox;
 import org.junit.Assert;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 import java.util.ArrayList;
 
 import static org.apache.hadoop.fs.StorageType.DISK;
 import static org.apache.hadoop.fs.StorageType.SSD;
+import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.anyByte;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -94,4 +96,75 @@ public void testUpdateQuotaAndCollectBlocks() {
     Assert.assertEquals(-BLOCK_SIZE, counts.getTypeSpaces().get(SSD));
   }
 
+  /**
+   * Test update quota with same blocks.
+   */
+  @Test
+  public void testUpdateQuotaDistinctBlocks() {
+    BlockStoragePolicySuite bsps = mock(BlockStoragePolicySuite.class);
+    BlockStoragePolicy bsp = mock(BlockStoragePolicy.class);
+    BlockInfo[] blocks = new BlockInfo[] {
+        new BlockInfoContiguous(new Block(1, BLOCK_SIZE, 1), REPL_3) };
+
+    INodeFile file = mock(INodeFile.class);
+    when(file.getBlocks()).thenReturn(blocks);
+    when(file.getStoragePolicyID()).thenReturn((byte) 1);
+    when(file.getPreferredBlockReplication()).thenReturn((short) 3);
+
+    when(bsps.getPolicy(anyByte())).thenReturn(bsp);
+    INode.BlocksMapUpdateInfo collectedBlocks =
+        mock(INode.BlocksMapUpdateInfo.class);
+    ArrayList<INode> removedINodes = new ArrayList<>();
+    INode.ReclaimContext ctx =
+        new INode.ReclaimContext(bsps, collectedBlocks, removedINodes, null);
+    QuotaCounts counts = ctx.quotaDelta().getCountsCopy();
+    INodeFile snapshotINode = mock(INodeFile.class);
+
+    // add same blocks in file diff
+    FileDiff diff1 = new FileDiff(0, snapshotINode, null, 0);
+    FileDiff diff = Mockito.spy(diff1);
+    Mockito.doReturn(blocks).when(diff).getBlocks();
+
+    // removed file diff
+    FileDiff removed = new FileDiff(0, snapshotINode, null, 0);
+
+    // remaining file diffs
+    FileDiffList diffs = new FileDiffList();
+    diffs.addFirst(diff);
+    FileWithSnapshotFeature sf = new FileWithSnapshotFeature(diffs);
+
+    // update quota and collect same blocks in file and file diff
+    when(file.getFileWithSnapshotFeature()).thenReturn(sf);
+    sf.updateQuotaAndCollectBlocks(ctx, file, removed);
+    counts = ctx.quotaDelta().getCountsCopy();
+    assertEquals(0, counts.getStorageSpace());
+
+    // different blocks in file and file's diff and in removed diff
+    BlockInfo[] blocks1 = new BlockInfo[] {
+        new BlockInfoContiguous(new Block(2, BLOCK_SIZE, 1), REPL_3) };
+    Mockito.doReturn(blocks1).when(diff).getBlocks();
+    // remaining file diffs
+    FileDiffList diffs1 = new FileDiffList();
+    diffs1.addFirst(diff);
+    FileWithSnapshotFeature sf1 = new FileWithSnapshotFeature(diffs1);
+    when(file.getFileWithSnapshotFeature()).thenReturn(sf1);
+    BlockInfo[] removedBlocks = new BlockInfo[] {
+        new BlockInfoContiguous(new Block(3, BLOCK_SIZE, 1), REPL_3) };
+    FileDiff removed1 = new FileDiff(0, snapshotINode, null, 1024);
+    removed1.setBlocks(removedBlocks);
+    INode.ReclaimContext ctx1 =
+        new INode.ReclaimContext(bsps, collectedBlocks, removedINodes, null);
+    sf1.updateQuotaAndCollectBlocks(ctx1, file, removed1);
+    counts = ctx1.quotaDelta().getCountsCopy();
+    assertEquals(3072, counts.getStorageSpace());
+
+    // same blocks in file and removed diff
+    removed1 = new FileDiff(0, snapshotINode, null, 1024);
+    removed1.setBlocks(blocks);
+    INode.ReclaimContext ctx2 =
+        new INode.ReclaimContext(bsps, collectedBlocks, removedINodes, null);
+    sf1.updateQuotaAndCollectBlocks(ctx2, file, removed1);
+    counts = ctx2.quotaDelta().getCountsCopy();
+    assertEquals(0, counts.getStorageSpace());
+  }
 }