HBASE-29238 ExportSnapshot support specify storage policy (#6874)

Signed-off-by: Duo Zhang <zhangduo@apache.org>
Signed-off-by: Wei-Chiu Chuang <weichiu@apache.org>
(cherry-picked from commit f6fdbbc)

Author: Reidddddd

hbase-mapreduce/src/main/java/org/apache/hadoop/hbase/snapshot/ExportSnapshot.java

@@ -26,9 +26,11 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
@@ -122,6 +124,7 @@ public class ExportSnapshot extends AbstractHBaseTool implements Tool {
     "snapshot.export.copy.references.threads";
   private static final int DEFAULT_COPY_MANIFEST_THREADS =
     Runtime.getRuntime().availableProcessors();
+  private static final String CONF_STORAGE_POLICY = "snapshot.export.storage.policy.family";
 
   static class Testing {
     static final String CONF_TEST_FAILURE = "test.snapshot.export.failure";
@@ -159,6 +162,8 @@ static final class Options {
       new Option(null, "bandwidth", true, "Limit bandwidth to this value in MB/second.");
     static final Option RESET_TTL =
       new Option(null, "reset-ttl", false, "Do not copy TTL for the snapshot");
+    static final Option STORAGE_POLICY = new Option(null, "storage-policy", true,
+      "Storage policy for export snapshot output directory, with format like: f=HOT&g=ALL_SDD");
   }
 
   // Export Map-Reduce Counters, to keep track of the progress
@@ -328,6 +333,13 @@ private void copyFile(final Context context, final SnapshotFileInfo inputInfo,
 
         // Ensure that the output folder is there and copy the file
         createOutputPath(outputPath.getParent());
+        String family = new Path(inputInfo.getHfile()).getParent().getName();
+        String familyStoragePolicy = generateFamilyStoragePolicyKey(family);
+        if (stringIsNotEmpty(context.getConfiguration().get(familyStoragePolicy))) {
+          String key = context.getConfiguration().get(familyStoragePolicy);
+          LOG.info("Setting storage policy {} for {}", key, outputPath.getParent());
+          outputFs.setStoragePolicy(outputPath.getParent(), key);
+        }
         FSDataOutputStream out = outputFs.create(outputPath, true);
 
         long stime = EnvironmentEdgeManager.currentTime();
@@ -426,7 +438,7 @@ private boolean preserveAttributes(final Path path, final FileStatus refStat) {
     }
 
     private boolean stringIsNotEmpty(final String str) {
-      return str != null && str.length() > 0;
+      return str != null && !str.isEmpty();
     }
 
     private long copyData(final Context context, final Path inputPath, final InputStream in,
@@ -907,8 +919,8 @@ public boolean nextKeyValue() {
    */
   private void runCopyJob(final Path inputRoot, final Path outputRoot, final String snapshotName,
     final Path snapshotDir, final boolean verifyChecksum, final String filesUser,
-    final String filesGroup, final int filesMode, final int mappers, final int bandwidthMB)
-    throws IOException, InterruptedException, ClassNotFoundException {
+    final String filesGroup, final int filesMode, final int mappers, final int bandwidthMB,
+    final String storagePolicy) throws IOException, InterruptedException, ClassNotFoundException {
     Configuration conf = getConf();
     if (filesGroup != null) conf.set(CONF_FILES_GROUP, filesGroup);
     if (filesUser != null) conf.set(CONF_FILES_USER, filesUser);
@@ -923,6 +935,11 @@ private void runCopyJob(final Path inputRoot, final Path outputRoot, final Strin
     conf.setInt(CONF_BANDWIDTH_MB, bandwidthMB);
     conf.set(CONF_SNAPSHOT_NAME, snapshotName);
     conf.set(CONF_SNAPSHOT_DIR, snapshotDir.toString());
+    if (storagePolicy != null) {
+      for (Map.Entry<String, String> entry : storagePolicyPerFamily(storagePolicy).entrySet()) {
+        conf.set(generateFamilyStoragePolicyKey(entry.getKey()), entry.getValue());
+      }
+    }
 
     String jobname = conf.get(CONF_MR_JOB_NAME, "ExportSnapshot-" + snapshotName);
     Job job = new Job(conf);
@@ -1009,6 +1026,23 @@ private void setPermissionParallel(final FileSystem outputFs, final short filesM
     }, conf);
   }
 
+  private Map<String, String> storagePolicyPerFamily(String storagePolicy) {
+    Map<String, String> familyStoragePolicy = new HashMap<>();
+    for (String familyConf : storagePolicy.split("&")) {
+      String[] familySplit = familyConf.split("=");
+      if (familySplit.length != 2) {
+        continue;
+      }
+      // family is key, storage policy is value
+      familyStoragePolicy.put(familySplit[0], familySplit[1]);
+    }
+    return familyStoragePolicy;
+  }
+
+  private static String generateFamilyStoragePolicyKey(String family) {
+    return CONF_STORAGE_POLICY + "." + family;
+  }
+
   private boolean verifyTarget = true;
   private boolean verifySource = true;
   private boolean verifyChecksum = true;
@@ -1023,6 +1057,7 @@ private void setPermissionParallel(final FileSystem outputFs, final short filesM
   private int filesMode = 0;
   private int mappers = 0;
   private boolean resetTtl = false;
+  private String storagePolicy = null;
 
   @Override
   protected void processOptions(CommandLine cmd) {
@@ -1045,6 +1080,9 @@ protected void processOptions(CommandLine cmd) {
     verifyTarget = !cmd.hasOption(Options.NO_TARGET_VERIFY.getLongOpt());
     verifySource = !cmd.hasOption(Options.NO_SOURCE_VERIFY.getLongOpt());
     resetTtl = cmd.hasOption(Options.RESET_TTL.getLongOpt());
+    if (cmd.hasOption(Options.STORAGE_POLICY.getLongOpt())) {
+      storagePolicy = cmd.getOptionValue(Options.STORAGE_POLICY.getLongOpt());
+    }
   }
 
   /**
@@ -1213,7 +1251,7 @@ public int doWork() throws IOException {
     // by the HFileArchiver, since they have no references.
     try {
       runCopyJob(inputRoot, outputRoot, snapshotName, snapshotDir, verifyChecksum, filesUser,
-        filesGroup, filesMode, mappers, bandwidthMB);
+        filesGroup, filesMode, mappers, bandwidthMB, storagePolicy);
 
       LOG.info("Finalize the Snapshot Export");
       if (!skipTmp) {