HubSpot Backport: HBASE-29432: Provide mechanism to plug in rack or host locality logic into ExportSnapshot (not yet merged upstream)

Author: charlesconnell

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

@@ -24,17 +24,19 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.List;
-import java.util.HashSet;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.function.BiConsumer;
+import java.util.stream.Collectors;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FSDataInputStream;
 import org.apache.hadoop.fs.FSDataOutputStream;
@@ -72,12 +74,15 @@
 import org.apache.hadoop.mapreduce.TaskAttemptContext;
 import org.apache.hadoop.mapreduce.lib.output.NullOutputFormat;
 import org.apache.hadoop.mapreduce.security.TokenCache;
+import org.apache.hadoop.util.ReflectionUtils;
 import org.apache.hadoop.util.StringUtils;
 import org.apache.hadoop.util.Tool;
 import org.apache.yetus.audience.InterfaceAudience;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import org.apache.hbase.thirdparty.com.google.common.collect.ImmutableList;
+import org.apache.hbase.thirdparty.com.google.common.collect.ImmutableSet;
 import org.apache.hbase.thirdparty.org.apache.commons.cli.CommandLine;
 import org.apache.hbase.thirdparty.org.apache.commons.cli.Option;
 
@@ -115,6 +120,10 @@ public class ExportSnapshot extends AbstractHBaseTool implements Tool {
   private static final String CONF_MAP_GROUP = "snapshot.export.default.map.group";
   private static final String CONF_BANDWIDTH_MB = "snapshot.export.map.bandwidth.mb";
   private static final String CONF_MR_JOB_NAME = "mapreduce.job.name";
+  private static final String CONF_INPUT_FILE_GROUPER_CLASS =
+    "snapshot.export.input.file.grouper.class";
+  private static final String CONF_INPUT_FILE_LOCATION_RESOLVER_CLASS =
+    "snapshot.export.input.file.location.resolver.class";
   protected static final String CONF_SKIP_TMP = "snapshot.export.skip.tmp";
   private static final String CONF_COPY_MANIFEST_THREADS =
     "snapshot.export.copy.references.threads";
@@ -152,11 +161,19 @@ static final class Options {
     static final Option CHMOD =
       new Option(null, "chmod", true, "Change the permission of the files to the specified one.");
     static final Option MAPPERS = new Option(null, "mappers", true,
-      "Number of mappers to use during the copy (mapreduce.job.maps).");
+      "Number of mappers to use during the copy (mapreduce.job.maps). "
+        + "If you provide a --custom-file-grouper, "
+        + "then --mappers is interpreted as the number of mappers per group.");
     static final Option BANDWIDTH =
       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 CUSTOM_FILE_GROUPER = new Option(null, "custom-file-grouper", true,
+      "Fully qualified class name of an implementation of ExportSnapshot.CustomFileGrouper. "
+        + "See JavaDoc on that class for more information.");
+    static final Option FILE_LOCATION_RESOLVER = new Option(null, "file-location-resolver", true,
+      "Fully qualified class name of an implementation of ExportSnapshot.FileLocationResolver. "
+        + "See JavaDoc on that class for more information.");
   }
 
   // Export Map-Reduce Counters, to keep track of the progress
@@ -179,6 +196,54 @@ public enum ChecksumComparison {
     INCOMPATIBLE, // checksum comparison is not compatible.
   }
 
+  /**
+   * If desired, you may implement a CustomFileGrouper in order to influence how ExportSnapshot
+   * chooses which input files go into the MapReduce job's {@link InputSplit}s. Your implementation
+   * must return a data structure that contains each input file exactly once. Files that appear in
+   * separate entries in the top-level returned Collection are guaranteed to not be placed in the
+   * same InputSplit. This can be used to segregate your input files by the rack or host on which
+   * they are available, which, used in conjunction with {@link FileLocationResolver}, can improve
+   * the performance of your ExportSnapshot runs. To use this, pass the --custom-file-grouper
+   * argument with the fully qualified class name of an implementation of CustomFileGrouper that's
+   * on the classpath. If this argument is not used, no particular grouping logic will be applied.
+   */
+  @InterfaceAudience.Public
+  public interface CustomFileGrouper {
+    Collection<Collection<Pair<SnapshotFileInfo, Long>>>
+      getGroupedInputFiles(final Collection<Pair<SnapshotFileInfo, Long>> snapshotFiles);
+  }
+
+  private static class NoopCustomFileGrouper implements CustomFileGrouper {
+    @Override
+    public Collection<Collection<Pair<SnapshotFileInfo, Long>>>
+      getGroupedInputFiles(final Collection<Pair<SnapshotFileInfo, Long>> snapshotFiles) {
+      return ImmutableList.of(snapshotFiles);
+    }
+  }
+
+  /**
+   * If desired, you may implement a FileLocationResolver in order to influence the _location_
+   * metadata attached to each {@link InputSplit} that ExportSnapshot will submit to YARN. The
+   * method {@link #getLocationsForInputFiles(Collection)} method is called once for each InputSplit
+   * being constructed. Whatever is returned will ultimately be reported by that split's
+   * {@link InputSplit#getLocations()} method. This can be used to encourage YARN to schedule the
+   * ExportSnapshot's mappers on rack-local or host-local NodeManagers. To use this, pass the
+   * --file-location-resolver argument with the fully qualified class name of an implementation of
+   * FileLocationResolver that's on the classpath. If this argument is not used, no locations will
+   * be attached to the InputSplits.
+   */
+  @InterfaceAudience.Public
+  public interface FileLocationResolver {
+    Set<String> getLocationsForInputFiles(final Collection<Pair<SnapshotFileInfo, Long>> files);
+  }
+
+  static class NoopFileLocationResolver implements FileLocationResolver {
+    @Override
+    public Set<String> getLocationsForInputFiles(Collection<Pair<SnapshotFileInfo, Long>> files) {
+      return ImmutableSet.of();
+    }
+  }
+
   private static class ExportMapper
     extends Mapper<BytesWritable, NullWritable, NullWritable, NullWritable> {
     private static final Logger LOG = LoggerFactory.getLogger(ExportMapper.class);
@@ -710,8 +775,9 @@ private static Pair<SnapshotFileInfo, Long> getSnapshotFileAndSize(FileSystem fs
    * The algorithm used is pretty straightforward; the file list is sorted by size, and then each
    * group fetch the bigger file available, iterating through groups alternating the direction.
    */
-  static List<List<Pair<SnapshotFileInfo, Long>>>
-    getBalancedSplits(final List<Pair<SnapshotFileInfo, Long>> files, final int ngroups) {
+  static List<List<Pair<SnapshotFileInfo, Long>>> getBalancedSplits(
+    final Collection<Pair<SnapshotFileInfo, Long>> unsortedFiles, final int ngroups) {
+    List<Pair<SnapshotFileInfo, Long>> files = new ArrayList<>(unsortedFiles);
     // Sort files by size, from small to big
     Collections.sort(files, new Comparator<Pair<SnapshotFileInfo, Long>>() {
       public int compare(Pair<SnapshotFileInfo, Long> a, Pair<SnapshotFileInfo, Long> b) {
@@ -722,7 +788,6 @@ public int compare(Pair<SnapshotFileInfo, Long> a, Pair<SnapshotFileInfo, Long>
 
     // create balanced groups
     List<List<Pair<SnapshotFileInfo, Long>>> fileGroups = new LinkedList<>();
-    long[] sizeGroups = new long[ngroups];
     int hi = files.size() - 1;
     int lo = 0;
 
@@ -741,7 +806,6 @@ public int compare(Pair<SnapshotFileInfo, Long> a, Pair<SnapshotFileInfo, Long>
       Pair<SnapshotFileInfo, Long> fileInfo = files.get(hi--);
 
       // add the hi one
-      sizeGroups[g] += fileInfo.getSecond();
       group.add(fileInfo);
 
       // change direction when at the end or the beginning
@@ -755,16 +819,10 @@ public int compare(Pair<SnapshotFileInfo, Long> a, Pair<SnapshotFileInfo, Long>
       }
     }
 
-    if (LOG.isDebugEnabled()) {
-      for (int i = 0; i < sizeGroups.length; ++i) {
-        LOG.debug("export split=" + i + " size=" + StringUtils.humanReadableInt(sizeGroups[i]));
-      }
-    }
-
     return fileGroups;
   }
 
-  private static class ExportSnapshotInputFormat extends InputFormat<BytesWritable, NullWritable> {
+  static class ExportSnapshotInputFormat extends InputFormat<BytesWritable, NullWritable> {
     @Override
     public RecordReader<BytesWritable, NullWritable> createRecordReader(InputSplit split,
       TaskAttemptContext tac) throws IOException, InterruptedException {
@@ -778,37 +836,78 @@ public List<InputSplit> getSplits(JobContext context) throws IOException, Interr
       FileSystem fs = FileSystem.get(snapshotDir.toUri(), conf);
 
       List<Pair<SnapshotFileInfo, Long>> snapshotFiles = getSnapshotFiles(conf, fs, snapshotDir);
+
+      Collection<List<Pair<SnapshotFileInfo, Long>>> balancedGroups =
+        groupFilesForSplits(conf, snapshotFiles);
+
+      Class<? extends FileLocationResolver> fileLocationResolverClass =
+        conf.getClass(CONF_INPUT_FILE_LOCATION_RESOLVER_CLASS, NoopFileLocationResolver.class,
+          FileLocationResolver.class);
+      FileLocationResolver fileLocationResolver =
+        ReflectionUtils.newInstance(fileLocationResolverClass, conf);
+      LOG.info("FileLocationResolver {} will provide location metadata for each InputSplit",
+        fileLocationResolverClass);
+
+      List<InputSplit> splits = new ArrayList<>(balancedGroups.size());
+      for (Collection<Pair<SnapshotFileInfo, Long>> files : balancedGroups) {
+        splits.add(new ExportSnapshotInputSplit(files, fileLocationResolver));
+      }
+      return splits;
+    }
+
+    Collection<List<Pair<SnapshotFileInfo, Long>>> groupFilesForSplits(Configuration conf,
+      List<Pair<SnapshotFileInfo, Long>> snapshotFiles) {
       int mappers = conf.getInt(CONF_NUM_SPLITS, 0);
-      if (mappers == 0 && snapshotFiles.size() > 0) {
+      if (mappers == 0 && !snapshotFiles.isEmpty()) {
         mappers = 1 + (snapshotFiles.size() / conf.getInt(CONF_MAP_GROUP, 10));
         mappers = Math.min(mappers, snapshotFiles.size());
         conf.setInt(CONF_NUM_SPLITS, mappers);
         conf.setInt(MR_NUM_MAPS, mappers);
       }
 
-      List<List<Pair<SnapshotFileInfo, Long>>> groups = getBalancedSplits(snapshotFiles, mappers);
-      List<InputSplit> splits = new ArrayList(groups.size());
-      for (List<Pair<SnapshotFileInfo, Long>> files : groups) {
-        splits.add(new ExportSnapshotInputSplit(files));
-      }
-      return splits;
+      Class<? extends CustomFileGrouper> inputFileGrouperClass = conf.getClass(
+        CONF_INPUT_FILE_GROUPER_CLASS, NoopCustomFileGrouper.class, CustomFileGrouper.class);
+      CustomFileGrouper customFileGrouper =
+        ReflectionUtils.newInstance(inputFileGrouperClass, conf);
+      Collection<Collection<Pair<SnapshotFileInfo, Long>>> groups =
+        customFileGrouper.getGroupedInputFiles(snapshotFiles);
+
+      LOG.info("CustomFileGrouper {} split input files into {} groups", inputFileGrouperClass,
+        groups.size());
+      int mappersPerGroup = groups.isEmpty() ? 1 : Math.max(mappers / groups.size(), 1);
+      LOG.info(
+        "Splitting each group into {} InputSplits, "
+          + "to achieve closest possible amount of mappers to target of {}",
+        mappersPerGroup, mappers);
+
+      // Within each group, create splits of equal size. Groups are not mixed together.
+      return groups.stream().map(g -> getBalancedSplits(g, mappersPerGroup))
+        .flatMap(Collection::stream).collect(Collectors.toList());
     }
 
-    private static class ExportSnapshotInputSplit extends InputSplit implements Writable {
+    static class ExportSnapshotInputSplit extends InputSplit implements Writable {
+
       private List<Pair<BytesWritable, Long>> files;
+      private String[] locations;
       private long length;
 
       public ExportSnapshotInputSplit() {
         this.files = null;
+        this.locations = null;
       }
 
-      public ExportSnapshotInputSplit(final List<Pair<SnapshotFileInfo, Long>> snapshotFiles) {
-        this.files = new ArrayList(snapshotFiles.size());
+      public ExportSnapshotInputSplit(final Collection<Pair<SnapshotFileInfo, Long>> snapshotFiles,
+        FileLocationResolver fileLocationResolver) {
+        this.files = new ArrayList<>(snapshotFiles.size());
         for (Pair<SnapshotFileInfo, Long> fileInfo : snapshotFiles) {
           this.files.add(
             new Pair<>(new BytesWritable(fileInfo.getFirst().toByteArray()), fileInfo.getSecond()));
           this.length += fileInfo.getSecond();
         }
+        this.locations =
+          fileLocationResolver.getLocationsForInputFiles(snapshotFiles).toArray(new String[0]);
+        LOG.trace("This ExportSnapshotInputSplit has files {} of collective size {}, "
+          + "with location hints: {}", files, length, locations);
       }
 
       private List<Pair<BytesWritable, Long>> getSplitKeys() {
@@ -822,7 +921,7 @@ public long getLength() throws IOException, InterruptedException {
 
       @Override
       public String[] getLocations() throws IOException, InterruptedException {
-        return new String[] {};
+        return locations;
       }
 
       @Override
@@ -837,6 +936,12 @@ public void readFields(DataInput in) throws IOException {
           files.add(new Pair<>(fileInfo, size));
           length += size;
         }
+        int locationCount = in.readInt();
+        List<String> locations = new ArrayList<>(locationCount);
+        for (int i = 0; i < locationCount; ++i) {
+          locations.add(in.readUTF());
+        }
+        this.locations = locations.toArray(new String[0]);
       }
 
       @Override
@@ -846,6 +951,10 @@ public void write(DataOutput out) throws IOException {
           fileInfo.getFirst().write(out);
           out.writeLong(fileInfo.getSecond());
         }
+        out.writeInt(locations.length);
+        for (String location : locations) {
+          out.writeUTF(location);
+        }
       }
     }
 
@@ -905,7 +1014,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)
+    final String filesGroup, final int filesMode, final int mappers, final int bandwidthMB,
+    final String customFileGrouper, final String fileLocationResolver)
     throws IOException, InterruptedException, ClassNotFoundException {
     Configuration conf = getConf();
     if (filesGroup != null) conf.set(CONF_FILES_GROUP, filesGroup);
@@ -921,6 +1031,12 @@ 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 (customFileGrouper != null) {
+      conf.set(CONF_INPUT_FILE_GROUPER_CLASS, customFileGrouper);
+    }
+    if (fileLocationResolver != null) {
+      conf.set(CONF_INPUT_FILE_LOCATION_RESOLVER_CLASS, fileLocationResolver);
+    }
 
     String jobname = conf.get(CONF_MR_JOB_NAME, "ExportSnapshot-" + snapshotName);
     Job job = new Job(conf);
@@ -1021,6 +1137,8 @@ private void setPermissionParallel(final FileSystem outputFs, final short filesM
   private int filesMode = 0;
   private int mappers = 0;
   private boolean resetTtl = false;
+  private String customFileGrouper = null;
+  private String fileLocationResolver = null;
 
   @Override
   protected void processOptions(CommandLine cmd) {
@@ -1043,6 +1161,12 @@ 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.CUSTOM_FILE_GROUPER.getLongOpt())) {
+      customFileGrouper = cmd.getOptionValue(Options.CUSTOM_FILE_GROUPER.getLongOpt());
+    }
+    if (cmd.hasOption(Options.FILE_LOCATION_RESOLVER.getLongOpt())) {
+      fileLocationResolver = cmd.getOptionValue(Options.FILE_LOCATION_RESOLVER.getLongOpt());
+    }
   }
 
   /**
@@ -1211,7 +1335,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, customFileGrouper, fileLocationResolver);
 
       LOG.info("Finalize the Snapshot Export");
       if (!skipTmp) {
@@ -1269,6 +1393,8 @@ protected void addOptions() {
     addOption(Options.MAPPERS);
     addOption(Options.BANDWIDTH);
     addOption(Options.RESET_TTL);
+    addOption(Options.CUSTOM_FILE_GROUPER);
+    addOption(Options.FILE_LOCATION_RESOLVER);
   }
 
   public static void main(String[] args) {