public void updateQuotaAndCollectBlocks(INode.ReclaimContext reclaimContext,
      INodeFile file, FileDiff removed) {
    byte storagePolicyID = file.getStoragePolicyID();
    BlockStoragePolicy bsp = null;
    if (storagePolicyID != HdfsConstants.BLOCK_STORAGE_POLICY_ID_UNSPECIFIED) {
      bsp = reclaimContext.storagePolicySuite().
          getPolicy(file.getStoragePolicyID());
    }

    QuotaCounts oldCounts;
    if (removed.snapshotINode != null) {
      oldCounts = new QuotaCounts.Builder().build();
      // collect all distinct blocks
      Set<BlockInfo> allBlocks = new HashSet<BlockInfo>();
      if (file.getBlocks() != null) {
        allBlocks.addAll(Arrays.asList(file.getBlocks()));
      }
      if (removed.getBlocks() != null) {
        allBlocks.addAll(Arrays.asList(removed.getBlocks()));
      }
      for (FileDiff diff : diffs) {
        BlockInfo[] diffBlocks = diff.getBlocks();
        if (diffBlocks != null) {
          allBlocks.addAll(Arrays.asList(diffBlocks));
        }
      }
      for (BlockInfo b: allBlocks) {
        short replication = b.getReplication();
        long blockSize = b.isComplete() ? b.getNumBytes() : file
            .getPreferredBlockSize();

        oldCounts.addStorageSpace(blockSize * replication);

        if (bsp != null) {
          List<StorageType> oldTypeChosen = bsp.chooseStorageTypes(replication);
          for (StorageType t : oldTypeChosen) {
            if (t.supportTypeQuota()) {
              oldCounts.addTypeSpace(t, blockSize);
            }
          }
        }
      }

      AclFeature aclFeature = removed.getSnapshotINode().getAclFeature();
      if (aclFeature != null) {
        AclStorage.removeAclFeature(aclFeature);
      }
    } else {
      oldCounts = file.storagespaceConsumed(null);
    }

    getDiffs().combineAndCollectSnapshotBlocks(reclaimContext, file, removed);
    if (file.getBlocks() != null) {
      short replInDiff = getMaxBlockRepInDiffs(removed);
      short repl = (short) Math.max(file.getPreferredBlockReplication(),
                                    replInDiff);
      for (BlockInfo b : file.getBlocks()) {
        if (repl != b.getReplication()) {
          reclaimContext.collectedBlocks().addUpdateReplicationFactor(b, repl);
        }
      }
    }
    QuotaCounts current = file.storagespaceConsumed(bsp);
    reclaimContext.quotaDelta().add(oldCounts.subtract(current));
  }
@Test
  public void testUpdateQuotaAndCollectBlocks() {
    FileDiffList diffs = new FileDiffList();
    FileWithSnapshotFeature sf = new FileWithSnapshotFeature(diffs);
    FileDiff diff = mock(FileDiff.class);
    BlockStoragePolicySuite bsps = mock(BlockStoragePolicySuite.class);
    BlockStoragePolicy bsp = mock(BlockStoragePolicy.class);
    BlockInfo[] blocks = new BlockInfo[] {
        new BlockInfoContiguous(new Block(1, BLOCK_SIZE, 1), REPL_1)
    };
    BlockManager bm = mock(BlockManager.class);

    // No snapshot
    INodeFile file = mock(INodeFile.class);
    when(file.getFileWithSnapshotFeature()).thenReturn(sf);
    when(file.getBlocks()).thenReturn(blocks);
    when(file.getStoragePolicyID()).thenReturn((byte) 1);
    Whitebox.setInternalState(file, "header", (long) REPL_1 << 48);
    when(file.getPreferredBlockReplication()).thenReturn(REPL_1);

    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);
    sf.updateQuotaAndCollectBlocks(ctx, file, diff);
    QuotaCounts counts = ctx.quotaDelta().getCountsCopy();
    Assert.assertEquals(0, counts.getStorageSpace());
    Assert.assertTrue(counts.getTypeSpaces().allLessOrEqual(0));

    // INode only exists in the snapshot
    INodeFile snapshotINode = mock(INodeFile.class);
    Whitebox.setInternalState(snapshotINode, "header", (long) REPL_3 << 48);
    Whitebox.setInternalState(diff, "snapshotINode", snapshotINode);
    when(diff.getSnapshotINode()).thenReturn(snapshotINode);

    when(bsp.chooseStorageTypes(REPL_1))
        .thenReturn(Lists.newArrayList(SSD));
    when(bsp.chooseStorageTypes(REPL_3))
        .thenReturn(Lists.newArrayList(DISK));
    blocks[0].setReplication(REPL_3);
    sf.updateQuotaAndCollectBlocks(ctx, file, diff);
    counts = ctx.quotaDelta().getCountsCopy();
    Assert.assertEquals((REPL_3 - REPL_1) * BLOCK_SIZE,
                        counts.getStorageSpace());
    Assert.assertEquals(BLOCK_SIZE, counts.getTypeSpaces().get(DISK));
    Assert.assertEquals(-BLOCK_SIZE, counts.getTypeSpaces().get(SSD));
  }