HDDS-7566. Refactor TestRocksDBCheckpointDiffer tests.

hadoop-hdds/rocksdb-checkpoint-differ/src/test/java/org/apache/ozone/rocksdiff/TestCompactionDag.java

@@ -0,0 +1,256 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.ozone.rocksdiff;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import com.google.common.graph.GraphBuilder;
+import com.google.common.graph.MutableGraph;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+/**
+ * Test for CompactionDag.
+ */
+public class TestCompactionDag {
+
+  private static final List<List<String>> SST_FILES_BY_LEVEL = Arrays.asList(
+      Arrays.asList("000015", "000013", "000011", "000009"),
+      Arrays.asList("000018", "000016", "000017", "000026", "000024", "000022",
+          "000020"),
+      Arrays.asList("000027", "000030", "000028", "000029", "000031", "000039",
+          "000037", "000035", "000033"),
+      Arrays.asList("000040", "000044", "000042", "000043", "000045", "000041",
+          "000046", "000054", "000052", "000050", "000048"),
+      Arrays.asList("000059", "000055", "000056", "000060", "000057", "000058")
+  );
+
+  private static final List<List<CompactionNode>> COMPACTION_NODES_BY_LEVEL =
+      SST_FILES_BY_LEVEL.stream()
+          .map(sstFiles ->
+              sstFiles.stream()
+                  .map(
+                      sstFile -> new CompactionNode(sstFile,
+                          1000L,
+                          null, null, null
+                      ))
+                  .collect(Collectors.toList()))
+          .collect(Collectors.toList());
+
+  /**
+   * Creates a backward compaction DAG from a list of level nodes.
+   * It assumes that at each level files get compacted to the half of number
+   * of files at the next level.
+   * e.g. if level-1 has 7 files and level-2 has 9 files, so first 4 files
+   * at level-2 are from compaction of level-1 and rests are new.
+   */
+  private static MutableGraph<CompactionNode> createBackwardDagFromLevelNodes(
+      int fromLevel,
+      int toLevel
+  ) {
+    MutableGraph<CompactionNode> dag  = GraphBuilder.directed().build();
+
+    if (fromLevel == toLevel) {
+      COMPACTION_NODES_BY_LEVEL.get(fromLevel).forEach(dag::addNode);
+      return dag;
+    }
+
+    for (int level = fromLevel; level < toLevel; level++) {
+      List<CompactionNode> currentLevel = COMPACTION_NODES_BY_LEVEL.get(level);
+      List<CompactionNode> nextLevel = COMPACTION_NODES_BY_LEVEL.get(level + 1);
+
+      for (CompactionNode compactionNode : currentLevel) {
+        for (int j = 0; j < nextLevel.size(); j++) {
+          dag.addNode(compactionNode);
+          dag.addNode(nextLevel.get(j));
+
+          int child = nextLevel.size();
+          if (level < COMPACTION_NODES_BY_LEVEL.size() - 2) {
+            child /= 2;
+          }
+
+          if (j < child) {
+            dag.putEdge(compactionNode, nextLevel.get(j));
+          }
+        }
+      }
+    }
+
+    return dag;
+  }
+
+  /**
+   * Creates a forward compaction DAG from a list of level nodes.
+   * It assumes that at each level first half of the files are from the
+   * compaction of the previous level.
+   * e.g. if level-1 has 7 files and level-2 has 9 files, so first 4 files
+   * at level-2 are from compaction of level-1 and rests are new.
+   */
+  private static MutableGraph<CompactionNode> createForwardDagFromLevelNodes(
+      int fromLevel,
+      int toLevel
+  ) {
+    MutableGraph<CompactionNode> dag  = GraphBuilder.directed().build();
+
+    if (fromLevel == toLevel) {
+      COMPACTION_NODES_BY_LEVEL.get(fromLevel).forEach(dag::addNode);
+      return dag;
+    }
+
+    dag  = GraphBuilder.directed().build();
+    for (int level = fromLevel; level > toLevel; level--) {
+      List<CompactionNode> currentLevel = COMPACTION_NODES_BY_LEVEL.get(level);
+      List<CompactionNode> nextLevel = COMPACTION_NODES_BY_LEVEL.get(level - 1);
+
+      for (int i = 0; i < currentLevel.size(); i++) {
+        for (CompactionNode compactionNode : nextLevel) {
+          dag.addNode(currentLevel.get(i));
+          dag.addNode(compactionNode);
+
+          int parent = currentLevel.size();
+          if (level < COMPACTION_NODES_BY_LEVEL.size() - 1) {
+            parent /= 2;
+          }
+
+          if (i < parent) {
+            dag.putEdge(currentLevel.get(i), compactionNode);
+          }
+        }
+      }
+    }
+
+    return dag;
+  }
+
+  /**
+   * Test cases for pruneBackwardDag.
+   */
+  private static Stream<Arguments> pruneBackwardDagScenarios() {
+    Set<String> level0Files = new HashSet<>(SST_FILES_BY_LEVEL.get(0));
+    Set<String> level1Files = new HashSet<>(SST_FILES_BY_LEVEL.get(1));
+    Set<String> level2Files = new HashSet<>(SST_FILES_BY_LEVEL.get(2));
+    Set<String> level3Files = new HashSet<>(SST_FILES_BY_LEVEL.get(3));
+
+    level1Files.addAll(level0Files);
+    level2Files.addAll(level1Files);
+    level3Files.addAll(level2Files);
+
+    return Stream.of(
+        Arguments.of("Remove level 0 from backward DAG",
+            createBackwardDagFromLevelNodes(0, 4),
+            new HashSet<>(COMPACTION_NODES_BY_LEVEL.get(0)),
+            createBackwardDagFromLevelNodes(1, 4),
+            level0Files
+        ),
+        Arguments.of("Remove level 1 from backward DAG",
+            createBackwardDagFromLevelNodes(0, 4),
+            new HashSet<>(COMPACTION_NODES_BY_LEVEL.get(1)),
+            createBackwardDagFromLevelNodes(2, 4),
+            level1Files
+        ),
+        Arguments.of("Remove level 2 from backward DAG",
+            createBackwardDagFromLevelNodes(0, 4),
+            new HashSet<>(COMPACTION_NODES_BY_LEVEL.get(2)),
+            createBackwardDagFromLevelNodes(3, 4),
+            level2Files
+        ),
+        Arguments.of("Remove level 3 from backward DAG",
+            createBackwardDagFromLevelNodes(0, 4),
+            new HashSet<>(COMPACTION_NODES_BY_LEVEL.get(3)),
+            createBackwardDagFromLevelNodes(4, 4),
+            level3Files
+        )
+    );
+  }
+
+  @ParameterizedTest(name = "{0}")
+  @MethodSource("pruneBackwardDagScenarios")
+  public void testPruneBackwardDag(String description,
+      MutableGraph<CompactionNode> originalDag,
+      Set<CompactionNode> levelToBeRemoved,
+      MutableGraph<CompactionNode> expectedDag,
+      Set<String> expectedFileNodesRemoved) {
+    CompactionDag compactionDag = new CompactionDag();
+    Set<String> actualFileNodesRemoved =
+        compactionDag.pruneBackwardDag(originalDag, levelToBeRemoved);
+    assertEquals(expectedDag, originalDag);
+    assertEquals(actualFileNodesRemoved, expectedFileNodesRemoved);
+  }
+
+  /**
+   * Test cases for pruneBackwardDag.
+   */
+  private static Stream<Arguments> pruneForwardDagScenarios() {
+    Set<String> level0Files = new HashSet<>(SST_FILES_BY_LEVEL.get(0));
+    Set<String> level1Files = new HashSet<>(SST_FILES_BY_LEVEL.get(1));
+    Set<String> level2Files = new HashSet<>(SST_FILES_BY_LEVEL.get(2));
+    Set<String> level3Files = new HashSet<>(SST_FILES_BY_LEVEL.get(3));
+
+    level1Files.addAll(level0Files);
+    level2Files.addAll(level1Files);
+    level3Files.addAll(level2Files);
+
+    return Stream.of(
+        Arguments.of("Remove level 0 from forward DAG",
+            createForwardDagFromLevelNodes(4, 0),
+            new HashSet<>(COMPACTION_NODES_BY_LEVEL.get(0)),
+            createForwardDagFromLevelNodes(4, 1),
+            level0Files
+        ),
+        Arguments.of("Remove level 1 from forward DAG",
+            createForwardDagFromLevelNodes(4, 0),
+            new HashSet<>(COMPACTION_NODES_BY_LEVEL.get(1)),
+            createForwardDagFromLevelNodes(4, 2),
+            level1Files
+        ),
+        Arguments.of("Remove level 2 from forward DAG",
+            createForwardDagFromLevelNodes(4, 0),
+            new HashSet<>(COMPACTION_NODES_BY_LEVEL.get(2)),
+            createForwardDagFromLevelNodes(4, 3),
+            level2Files
+        ),
+        Arguments.of("Remove level 3 from forward DAG",
+            createForwardDagFromLevelNodes(4, 0),
+            new HashSet<>(COMPACTION_NODES_BY_LEVEL.get(3)),
+            createForwardDagFromLevelNodes(4, 4),
+            level3Files
+        )
+    );
+  }
+
+  @ParameterizedTest(name = "{0}")
+  @MethodSource("pruneForwardDagScenarios")
+  public void testPruneForwardDag(String description,
+      MutableGraph<CompactionNode> originalDag,
+      Set<CompactionNode> levelToBeRemoved,
+      MutableGraph<CompactionNode> expectedDag,
+      Set<String> expectedFileNodesRemoved) {
+    CompactionDag compactionDag = new CompactionDag();
+    Set<String> actualFileNodesRemoved =
+        compactionDag.pruneForwardDag(originalDag, levelToBeRemoved);
+    assertEquals(expectedDag, originalDag);
+    assertEquals(actualFileNodesRemoved, expectedFileNodesRemoved);
+  }
+}