HDDS-12707. Recon - In-memory extraction of om tarball from network input stream

hadoop-ozone/recon/src/main/java/org/apache/hadoop/ozone/recon/spi/impl/OzoneManagerServiceProviderImpl.java

@@ -42,27 +42,39 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Iterators;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
+import java.io.ByteArrayInputStream;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.channels.Channels;
+import java.nio.channels.FileChannel;
+import java.nio.channels.ReadableByteChannel;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.nio.file.StandardOpenOption;
 import java.nio.file.attribute.PosixFilePermission;
 import java.nio.file.attribute.PosixFilePermissions;
-import java.util.Arrays;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
+import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.ForkJoinWorkerThread;
+import java.util.concurrent.Future;
+import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 import javax.inject.Inject;
 import javax.inject.Singleton;
-import org.apache.commons.io.FileUtils;
+import org.apache.commons.compress.archivers.tar.TarArchiveEntry;
+import org.apache.commons.compress.archivers.tar.TarArchiveInputStream;
 import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.apache.hadoop.hdds.conf.OzoneConfiguration;
 import org.apache.hadoop.hdds.protocol.proto.HddsProtos;
@@ -113,6 +125,8 @@ public class OzoneManagerServiceProviderImpl
       LoggerFactory.getLogger(OzoneManagerServiceProviderImpl.class);
   private URLConnectionFactory connectionFactory;
 
+  private int omDBTarProcessorThreadCount; // Number of parallel workers
+
   private File omSnapshotDBParentDir = null;
   private File reconDbDir = null;
   private String omDBSnapshotUrl;
@@ -134,6 +148,8 @@ public class OzoneManagerServiceProviderImpl
   private ThreadFactory threadFactory;
   private ReconContext reconContext;
   private ReconTaskStatusUpdaterManager taskStatusUpdaterManager;
+  private ForkJoinPool.ForkJoinWorkerThreadFactory fetchOMDBThreadFactory;
+  private ForkJoinPool omDbTarExecutor;
 
   /**
    * OM Snapshot related task names.
@@ -220,6 +236,13 @@ public OzoneManagerServiceProviderImpl(
     this.threadFactory =
         new ThreadFactoryBuilder().setNameFormat(threadNamePrefix + "SyncOM-%d")
             .build();
+    this.omDBTarProcessorThreadCount = Runtime.getRuntime().availableProcessors();
+    // Create a custom ForkJoinWorkerThreadFactory to name threads
+    this.fetchOMDBThreadFactory = pool -> {
+      ForkJoinWorkerThread worker = ForkJoinPool.defaultForkJoinWorkerThreadFactory.newThread(pool);
+      worker.setName("FetchOMDBTar-" + threadNamePrefix + "-" + worker.getPoolIndex());
+      return worker;
+    };
     this.reconContext = reconContext;
     this.taskStatusUpdaterManager = taskStatusUpdaterManager;
     this.omDBLagThreshold = configuration.getLong(RECON_OM_DELTA_UPDATE_LAG_THRESHOLD,
@@ -235,6 +258,7 @@ public OMMetadataManager getOMMetadataManagerInstance() {
   public void start() {
     LOG.info("Starting Ozone Manager Service Provider.");
     scheduler = Executors.newScheduledThreadPool(1, threadFactory);
+    omDbTarExecutor = new ForkJoinPool(omDBTarProcessorThreadCount, fetchOMDBThreadFactory, null, false);
     try {
       omMetadataManager.start(configuration);
     } catch (IOException ioEx) {
@@ -357,6 +381,7 @@ public void stop() throws Exception {
     scheduler.shutdownNow();
     metrics.unRegister();
     connectionFactory.destroy();
+    omDbTarExecutor.shutdownNow();
   }
 
   /**
@@ -393,50 +418,176 @@ private boolean isOmSpnegoEnabled() {
    * @return DBCheckpoint instance.
    */
   @VisibleForTesting
-  DBCheckpoint getOzoneManagerDBSnapshot() {
-    String snapshotFileName = RECON_OM_SNAPSHOT_DB + "_" +
-        System.currentTimeMillis();
-    File targetFile = new File(omSnapshotDBParentDir, snapshotFileName +
-        ".tar");
+  public DBCheckpoint getOzoneManagerDBSnapshot() {
+    String snapshotFileName = RECON_OM_SNAPSHOT_DB + "_" + System.currentTimeMillis();
+    Path untarredDbDir = Paths.get(omSnapshotDBParentDir.getAbsolutePath(), snapshotFileName);
+
+    BlockingQueue<TarEntryData> queue = new LinkedBlockingQueue<>(omDBTarProcessorThreadCount * 10);
+
+    List<Future<Void>> futures = new ArrayList<>();
+
     try {
       SecurityUtil.doAsLoginUser(() -> {
         try (InputStream inputStream = reconUtils.makeHttpCall(
-            connectionFactory, getOzoneManagerSnapshotUrl(),
-            isOmSpnegoEnabled()).getInputStream()) {
-          FileUtils.copyInputStreamToFile(inputStream, targetFile);
+            connectionFactory, getOzoneManagerSnapshotUrl(), isOmSpnegoEnabled()).getInputStream();
+             TarArchiveInputStream tarInput = new TarArchiveInputStream(inputStream)) {
+
+          // Start worker threads (Consumers)
+          for (int i = 0; i < omDBTarProcessorThreadCount; i++) {
+            futures.add(omDbTarExecutor.submit(() -> {
+              processQueue(queue, untarredDbDir);
+              return null;
+            }));
+          }
+
+          // Producer: Read TAR entries & enqueue them
+          TarArchiveEntry entry;
+          while ((entry = tarInput.getNextTarEntry()) != null) {
+            if (entry.isDirectory()) {
+              File dir = new File(untarredDbDir.toFile(), entry.getName());
+              if (!dir.exists() && !dir.mkdirs()) {
+                throw new IOException("Failed to create directory: " + dir);
+              }
+            } else {
+              long expectedSize = entry.getSize();
+              byte[] fileData = new byte[(int) expectedSize];
+              int bytesReadTotal = 0;
+
+              while (bytesReadTotal < expectedSize) {
+                int bytesRead = tarInput.read(fileData, bytesReadTotal, (int) (expectedSize - bytesReadTotal));
+                if (bytesRead == -1) {
+                  throw new IOException("Unexpected end of stream for entry: " + entry.getName());
+                }
+                bytesReadTotal += bytesRead;
+              }
+
+              queue.put(new TarEntryData(entry.getName(), new ByteArrayInputStream(fileData), expectedSize));
+            }
+          }
+
+          // Signal consumers that no more data will come
+          for (int i = 0; i < omDBTarProcessorThreadCount; i++) {
+            queue.put(TarEntryData.EOF);
+          }
+
+        } catch (IOException | InterruptedException e) {
+          throw new RuntimeException("Error while extracting OM DB Snapshot TAR.", e);
         }
         return null;
       });
-      // Untar the checkpoint file.
-      Path untarredDbDir = Paths.get(omSnapshotDBParentDir.getAbsolutePath(), snapshotFileName);
-      reconUtils.untarCheckpointFile(targetFile, untarredDbDir);
 
-      // Validate the presence of required SST files
-      File[] sstFiles = untarredDbDir.toFile().listFiles((dir, name) -> name.endsWith(".sst"));
-      if (sstFiles == null || sstFiles.length == 0) {
-        LOG.warn("No SST files found in the OM snapshot directory: {}", untarredDbDir);
+      // Wait for all worker threads to finish
+      for (Future<Void> future : futures) {
+        future.get();
       }
 
-      List<String> sstFileNames = Arrays.stream(sstFiles)
-          .map(File::getName)
-          .collect(Collectors.toList());
-      LOG.debug("Valid SST files found: {}", sstFileNames);
+      // Validate extracted files
+      File[] sstFiles = untarredDbDir.toFile().listFiles((dir, name) -> name.endsWith(".sst"));
+      if (sstFiles != null && sstFiles.length > 0) {
+        LOG.warn("Number of SST files found in the OM snapshot directory: {} - {}", untarredDbDir, sstFiles.length);
+      }
 
-      // Currently, OM DB type is not configurable. Hence, defaulting to
-      // RocksDB.
-      reconContext.updateHealthStatus(new AtomicBoolean(true));
-      reconContext.getErrors().remove(ReconContext.ErrorCode.GET_OM_DB_SNAPSHOT_FAILED);
       return new RocksDBCheckpoint(untarredDbDir);
-    } catch (IOException e) {
-      LOG.error("Unable to obtain Ozone Manager DB Snapshot. ", e);
-      reconContext.updateHealthStatus(new AtomicBoolean(false));
-      reconContext.updateErrors(ReconContext.ErrorCode.GET_OM_DB_SNAPSHOT_FAILED);
+
+    } catch (IOException | InterruptedException | ExecutionException e) {
+      LOG.error("Unable to obtain Ozone Manager DB Snapshot.", e);
     } finally {
-      FileUtils.deleteQuietly(targetFile);
+      omDbTarExecutor.shutdown();
+      try {
+        if (!omDbTarExecutor.awaitTermination(10, TimeUnit.SECONDS)) {
+          LOG.warn("Forcing shutdown of OM DB Tar Executor...");
+          omDbTarExecutor.shutdownNow();
+        }
+      } catch (InterruptedException e) {
+        LOG.warn("Interrupted during shutdown. Forcing shutdown...");
+        omDbTarExecutor.shutdownNow();
+        Thread.currentThread().interrupt();
+      }
     }
+
     return null;
   }
 
+  /**
+   * Worker method that processes the queue and writes files in parallel.
+   */
+  private void processQueue(BlockingQueue<TarEntryData> queue, Path untarredDbDir) {
+    try {
+      while (true) {
+        TarEntryData entryData = queue.take(); // Blocking take
+        if (entryData == TarEntryData.EOF) {
+          break; // Exit thread when EOF is received
+        }
+
+        File outputFile = new File(untarredDbDir.toFile(), entryData.fileName);
+        Path parentDir = outputFile.toPath().getParent();
+
+        // Fix: Check if parentDir is null before calling createDirectories
+        if (parentDir != null && !Files.exists(parentDir)) {
+          Files.createDirectories(parentDir);
+        }
+
+        // Use new writeFile method that takes InputStream
+        if (entryData.inputStream != null) {
+          try (InputStream inputStream = entryData.inputStream) {
+            writeFile(untarredDbDir, entryData.fileName, inputStream, entryData.fileSize);
+          }
+        } else {
+          LOG.warn("Skipping file {} as inputStream is null", entryData.fileName);
+        }
+      }
+    } catch (IOException e) {
+      LOG.error("Error writing extracted file.", e);
+      throw new RuntimeException(e);
+    } catch (InterruptedException e) {
+      LOG.warn("Worker thread interrupted. Exiting...");
+      Thread.currentThread().interrupt();
+    }
+  }
+
+  /**
+   * Writes the extracted file to disk.
+   */
+  private void writeFile(Path untarredDbDir, String entryName, InputStream inputStream, long fileSize)
+      throws IOException {
+    File file = new File(untarredDbDir.toFile(), entryName);
+
+    // Ensure parent directories exist
+    File parentDir = file.getParentFile();
+    if (parentDir != null && !parentDir.exists() && !parentDir.mkdirs()) {
+      throw new IOException("Failed to create parent directories for: " + file);
+    }
+
+    // Write using FileChannel with Buffered Input
+    try (FileChannel outChannel = FileChannel.open(file.toPath(), StandardOpenOption.CREATE, StandardOpenOption.WRITE);
+         ReadableByteChannel inChannel = Channels.newChannel(inputStream)) {
+      long transferred = 0;
+      while (transferred < fileSize) {
+        long bytes = outChannel.transferFrom(inChannel, transferred, Math.min(8192, fileSize - transferred));
+        if (bytes <= 0) {
+          break; // Avoid infinite loops
+        }
+        transferred += bytes;
+      }
+    }
+  }
+
+  /**
+   * Wrapper class for TAR entry data.
+   */
+  private static class TarEntryData {
+    static final TarEntryData EOF = new TarEntryData(null, null, 0);
+    private final String fileName;
+    private final InputStream inputStream;
+    private final long fileSize;
+
+    TarEntryData(String fileName, InputStream inputStream, long fileSize) {
+      this.fileName = fileName;
+      this.inputStream = inputStream;
+      this.fileSize = fileSize;
+    }
+  }
+
   /**
    * Update Local OM DB with new OM DB snapshot.
    * @throws IOException