HDDS-13063. Support splitting json output to multiple valid json files in a directory

hadoop-ozone/cli-debug/src/main/java/org/apache/hadoop/ozone/debug/replicas/ReplicasVerify.java

@@ -17,12 +17,15 @@
 
 package org.apache.hadoop.ozone.debug.replicas;
 
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.apache.hadoop.ozone.conf.OzoneServiceConfig.DEFAULT_SHUTDOWN_HOOK_PRIORITY;
 
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import java.io.File;
 import java.io.IOException;
 import java.io.PrintStream;
+import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.Iterator;
@@ -89,6 +92,18 @@ public class ReplicasVerify extends Handler {
       defaultValue = "1000000")
   private long containerCacheSize;
 
+  @CommandLine.Option(names = {"--out", "-o"},
+      description = "Output directory to dump verification output.The directory is created if it does not exist, " +
+          "and files are named using the directory's name as the base, e.g. <dirName>.0, <dirName>.1, " +
+          "when splitting with --max-records-per-file (otherwise a single <dirName> file is written).")
+  private String outputDir;
+
+  @CommandLine.Option(names = {"--max-records-per-file"},
+      description = "Maximum number of keys to write per output file. When greater than zero (and --out is set), " +
+          "output is split into multiple valid JSON files named <out>.0, <out>.1, Requires --out.",
+      defaultValue = "0")
+  private long recordsPerFile;
+
   private List<ReplicaVerifier> replicaVerifiers;
 
   private static final String DURATION_FORMAT = "HH:mm:ss,SSS";
@@ -197,8 +212,68 @@ void findCandidateKeys(OzoneClient ozoneClient, OzoneAddress address) throws IOE
         checkVolume(ozoneClient, it.next(), keysArray, allKeysPassed);
       }
     }
-    root.put("pass", allKeysPassed.get());
-    System.out.println(JsonUtils.toJsonStringWithDefaultPrettyPrinter(root));
+    if (outputDir == null) {
+      root.put("pass", allKeysPassed.get());
+      System.out.println(JsonUtils.toJsonStringWithDefaultPrettyPrinter(root));
+    } else {
+      writeOutputToFiles(root, keysArray, allKeysPassed.get());
+    }
+  }
+
+  /**
+   * Writes verification output to file(s) instead of stdout.
+   * When recordsPerFile is greater than zero, the keys are split into multiple valid JSON files.
+   */
+  private void writeOutputToFiles(ObjectNode root, ArrayNode keysArray, boolean allKeysPassed) throws IOException {
+    String outputPrefix = resolveOutputPrefix();
+
+    if (recordsPerFile <= 0) {
+      root.put("pass", allKeysPassed);
+      writeJsonToFile(root, outputPrefix);
+      return;
+    }
+
+    int suffix = 0;
+    ObjectNode chunkNode = null;
+    ArrayNode chunkKeys = null;
+    for (int i = 0; i < keysArray.size(); i++) {
+      if (chunkNode == null) {
+        chunkNode = JsonUtils.createObjectNode(null);
+        chunkKeys = chunkNode.putArray("keys");
+      }
+      chunkKeys.add(keysArray.get(i));
+      if (chunkKeys.size() >= recordsPerFile) {
+        writeJsonToFile(chunkNode, outputPrefix + "." + suffix++);
+        chunkNode = null;
+      }
+    }
+    if (chunkNode != null) {
+      writeJsonToFile(chunkNode, outputPrefix + "." + suffix++);
+    }
+  }
+
+  /**
+   * Resolves the file name prefix used for output files from --out. The value of --out is always treated as a directory
+   * it is created when missing, and files are written inside it using the directory's own name as the base.
+   */
+  private String resolveOutputPrefix() throws IOException {
+    File outputDirectory = new File(outputDir);
+    if (outputDirectory.exists()) {
+      if (!outputDirectory.isDirectory()) {
+        throw new IOException("Output path already exists and is not a directory: " +
+            outputDirectory.getAbsolutePath());
+      }
+    } else if (!outputDirectory.mkdirs()) {
+      throw new IOException("An exception occurred while creating the directory. Directory: " +
+          outputDirectory.getAbsolutePath());
+    }
+    return new File(outputDirectory, outputDirectory.getName()).getPath();
+  }
+
+  private void writeJsonToFile(ObjectNode node, String targetFileName) throws IOException {
+    try (PrintWriter writer = new PrintWriter(targetFileName, UTF_8.name())) {
+      writer.println(JsonUtils.toJsonStringWithDefaultPrettyPrinter(node));
+    }
   }
 
   void checkVolume(OzoneClient ozoneClient, OzoneVolume volume, ArrayNode keysArray, AtomicBoolean allKeysPassed)

hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/shell/TestOzoneDebugReplicasVerify.java

@@ -24,10 +24,14 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.fail;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import java.io.File;
 import java.io.IOException;
+import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Iterator;
@@ -57,6 +61,7 @@
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
@@ -71,6 +76,7 @@ public abstract class TestOzoneDebugReplicasVerify implements NonHATests.TestCas
   private static final Logger LOG = LoggerFactory.getLogger(TestOzoneDebugReplicasVerify.class);
   private static final String CHUNKS_DIR_NAME = "chunks";
   private static final String BLOCK_FILE_EXTENSION = ".block";
+  private static final ObjectMapper MAPPER = new ObjectMapper();
 
   private OzoneDebug ozoneDebugShell;
   private String ozoneAddress;
@@ -261,4 +267,66 @@ void testChecksumsWithEmptyBlockFile() {
         .contains("Unexpected read size")
         .doesNotContain("Checksum mismatch");
   }
+
+  @Test
+  void testSplitOutputToNewDirectory(@TempDir Path tempDir) throws IOException {
+    int maxRecordsPerFile = 2;
+    int expectedKeyFiles = (int) Math.ceil(keyInfoMap.size() / (maxRecordsPerFile * 1.0));
+    // Directory does not exist yet: it should be created and files written inside as <dirName>.0, <dirName>.1
+    String dirName = "verify-replica";
+    File outDir = new File(tempDir.toFile(), dirName);
+
+    runVerifyToDirectory(outDir.getAbsolutePath(), maxRecordsPerFile);
+
+    assertSplitFilesInDirectory(outDir, dirName, expectedKeyFiles, maxRecordsPerFile);
+  }
+
+  @Test
+  void testSplitOutputToExistingDirectory(@TempDir Path tempDir) throws IOException {
+    int maxRecordsPerFile = 2;
+    int expectedKeyFiles = (int) Math.ceil(keyInfoMap.size() / (maxRecordsPerFile * 1.0));
+    // Directory already exists: files are written inside it using the directory's name as the base.
+    String dirName = "verify-output";
+    File outDir = new File(tempDir.toFile(), dirName);
+    assertTrue(outDir.mkdirs(), "Failed to create output directory: " + outDir.getAbsolutePath());
+
+    runVerifyToDirectory(outDir.getAbsolutePath(), maxRecordsPerFile);
+
+    assertSplitFilesInDirectory(outDir, dirName, expectedKeyFiles, maxRecordsPerFile);
+  }
+
+  private void runVerifyToDirectory(String outputDir, int maxRecordsPerFile) {
+    List<String> parameters = new ArrayList<>();
+    parameters.add(0, getSetConfStringFromConf(ScmConfigKeys.OZONE_SCM_CLIENT_ADDRESS_KEY));
+    parameters.add(0, getSetConfStringFromConf(OMConfigKeys.OZONE_OM_ADDRESS_KEY));
+    parameters.add("replicas");
+    parameters.add("verify");
+    parameters.add("--checksums");
+    parameters.add("--all-results");
+    parameters.add("--out");
+    parameters.add(outputDir);
+    parameters.add("--max-records-per-file");
+    parameters.add(String.valueOf(maxRecordsPerFile));
+    parameters.add(ozoneAddress);
+
+    int exitCode = ozoneDebugShell.execute(parameters.toArray(new String[0]));
+    assertEquals(0, exitCode, err.get());
+  }
+
+  private void assertSplitFilesInDirectory(File outDir, String baseName, int expectedKeyFiles, int maxRecordsPerFile)
+      throws IOException {
+    assertTrue(outDir.isDirectory(), "Output directory should be created: " + outDir.getAbsolutePath());
+    int keysInFiles = 0;
+    for (int i = 0; i < expectedKeyFiles; i++) {
+      File keyFile = new File(outDir, baseName + "." + i);
+      assertTrue(keyFile.isFile(), "Expected key file: " + keyFile.getAbsolutePath());
+      JsonNode jsonNode = MAPPER.readTree(keyFile);
+      assertNotNull(jsonNode, "Output file must be valid JSON: " + keyFile.getAbsolutePath());
+      JsonNode keys = jsonNode.get("keys");
+      assertNotNull(keys, "Each split file must contain a 'keys' array");
+      assertThat(keys.size()).isLessThanOrEqualTo(maxRecordsPerFile);
+      keysInFiles += keys.size();
+    }
+    assertEquals(keyInfoMap.size(), keysInFiles, "All keys should be written across the split files");
+  }
 }