001/*
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *     http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package org.apache.hadoop.hbase.mapreduce;
019
020import java.io.ByteArrayOutputStream;
021import java.io.DataInput;
022import java.io.DataOutput;
023import java.io.IOException;
024import java.lang.reflect.InvocationTargetException;
025import java.util.ArrayList;
026import java.util.List;
027import java.util.UUID;
028import org.apache.hadoop.conf.Configuration;
029import org.apache.hadoop.fs.FileSystem;
030import org.apache.hadoop.fs.Path;
031import org.apache.hadoop.hbase.HDFSBlocksDistribution;
032import org.apache.hadoop.hbase.HDFSBlocksDistribution.HostAndWeight;
033import org.apache.hadoop.hbase.HRegionLocation;
034import org.apache.hadoop.hbase.PrivateCellUtil;
035import org.apache.hadoop.hbase.client.ClientSideRegionScanner;
036import org.apache.hadoop.hbase.client.Connection;
037import org.apache.hadoop.hbase.client.ConnectionFactory;
038import org.apache.hadoop.hbase.client.IsolationLevel;
039import org.apache.hadoop.hbase.client.RegionInfo;
040import org.apache.hadoop.hbase.client.RegionLocator;
041import org.apache.hadoop.hbase.client.Result;
042import org.apache.hadoop.hbase.client.Scan;
043import org.apache.hadoop.hbase.client.Scan.ReadType;
044import org.apache.hadoop.hbase.client.TableDescriptor;
045import org.apache.hadoop.hbase.io.ImmutableBytesWritable;
046import org.apache.hadoop.hbase.mob.MobUtils;
047import org.apache.hadoop.hbase.regionserver.HRegion;
048import org.apache.hadoop.hbase.snapshot.RestoreSnapshotHelper;
049import org.apache.hadoop.hbase.snapshot.SnapshotDescriptionUtils;
050import org.apache.hadoop.hbase.snapshot.SnapshotManifest;
051import org.apache.hadoop.hbase.util.Bytes;
052import org.apache.hadoop.hbase.util.CommonFSUtils;
053import org.apache.hadoop.hbase.util.RegionSplitter;
054import org.apache.hadoop.io.Writable;
055import org.apache.hadoop.mapreduce.Job;
056import org.apache.yetus.audience.InterfaceAudience;
057import org.slf4j.Logger;
058import org.slf4j.LoggerFactory;
059
060import org.apache.hbase.thirdparty.com.google.common.collect.Lists;
061
062import org.apache.hadoop.hbase.shaded.protobuf.ProtobufUtil;
063import org.apache.hadoop.hbase.shaded.protobuf.generated.MapReduceProtos.TableSnapshotRegionSplit;
064import org.apache.hadoop.hbase.shaded.protobuf.generated.SnapshotProtos.SnapshotDescription;
065import org.apache.hadoop.hbase.shaded.protobuf.generated.SnapshotProtos.SnapshotRegionManifest;
066
067/**
068 * Hadoop MR API-agnostic implementation for mapreduce over table snapshots.
069 */
070@InterfaceAudience.Private
071public class TableSnapshotInputFormatImpl {
072  // TODO: Snapshots files are owned in fs by the hbase user. There is no
073  // easy way to delegate access.
074
075  public static final Logger LOG = LoggerFactory.getLogger(TableSnapshotInputFormatImpl.class);
076
077  private static final String SNAPSHOT_NAME_KEY = "hbase.TableSnapshotInputFormat.snapshot.name";
078  // key for specifying the root dir of the restored snapshot
079  protected static final String RESTORE_DIR_KEY = "hbase.TableSnapshotInputFormat.restore.dir";
080
081  /** See {@link #getBestLocations(Configuration, HDFSBlocksDistribution, int)} */
082  private static final String LOCALITY_CUTOFF_MULTIPLIER =
083    "hbase.tablesnapshotinputformat.locality.cutoff.multiplier";
084  private static final float DEFAULT_LOCALITY_CUTOFF_MULTIPLIER = 0.8f;
085
086  /**
087   * For MapReduce jobs running multiple mappers per region, determines what split algorithm we
088   * should be using to find split points for scanners.
089   */
090  public static final String SPLIT_ALGO = "hbase.mapreduce.split.algorithm";
091  /**
092   * For MapReduce jobs running multiple mappers per region, determines number of splits to generate
093   * per region.
094   */
095  public static final String NUM_SPLITS_PER_REGION = "hbase.mapreduce.splits.per.region";
096
097  /**
098   * Whether to calculate the block location for splits. Default to true. If the computing layer
099   * runs outside of HBase cluster, the block locality does not master. Setting this value to false
100   * could skip the calculation and save some time. Set access modifier to "public" so that these
101   * could be accessed by test classes of both org.apache.hadoop.hbase.mapred and
102   * org.apache.hadoop.hbase.mapreduce.
103   */
104  public static final String SNAPSHOT_INPUTFORMAT_LOCALITY_ENABLED_KEY =
105    "hbase.TableSnapshotInputFormat.locality.enabled";
106  public static final boolean SNAPSHOT_INPUTFORMAT_LOCALITY_ENABLED_DEFAULT = true;
107
108  /**
109   * Whether to calculate the Snapshot region location by region location from meta. It is much
110   * faster than computing block locations for splits.
111   */
112  public static final String SNAPSHOT_INPUTFORMAT_LOCALITY_BY_REGION_LOCATION =
113    "hbase.TableSnapshotInputFormat.locality.by.region.location";
114
115  public static final boolean SNAPSHOT_INPUTFORMAT_LOCALITY_BY_REGION_LOCATION_DEFAULT = false;
116
117  /**
118   * In some scenario, scan limited rows on each InputSplit for sampling data extraction
119   */
120  public static final String SNAPSHOT_INPUTFORMAT_ROW_LIMIT_PER_INPUTSPLIT =
121    "hbase.TableSnapshotInputFormat.row.limit.per.inputsplit";
122
123  /**
124   * Whether to enable scan metrics on Scan, default to true
125   */
126  public static final String SNAPSHOT_INPUTFORMAT_SCAN_METRICS_ENABLED =
127    "hbase.TableSnapshotInputFormat.scan_metrics.enabled";
128
129  public static final boolean SNAPSHOT_INPUTFORMAT_SCAN_METRICS_ENABLED_DEFAULT = true;
130
131  /**
132   * The {@link ReadType} which should be set on the {@link Scan} to read the HBase Snapshot,
133   * default STREAM.
134   */
135  public static final String SNAPSHOT_INPUTFORMAT_SCANNER_READTYPE =
136    "hbase.TableSnapshotInputFormat.scanner.readtype";
137  public static final ReadType SNAPSHOT_INPUTFORMAT_SCANNER_READTYPE_DEFAULT = ReadType.STREAM;
138
139  /**
140   * Implementation class for InputSplit logic common between mapred and mapreduce.
141   */
142  public static class InputSplit implements Writable {
143
144    private TableDescriptor htd;
145    private RegionInfo regionInfo;
146    private String[] locations;
147    private String scan;
148    private String restoreDir;
149
150    // constructor for mapreduce framework / Writable
151    public InputSplit() {
152    }
153
154    public InputSplit(TableDescriptor htd, RegionInfo regionInfo, List<String> locations, Scan scan,
155      Path restoreDir) {
156      this.htd = htd;
157      this.regionInfo = regionInfo;
158      if (locations == null || locations.isEmpty()) {
159        this.locations = new String[0];
160      } else {
161        this.locations = locations.toArray(new String[locations.size()]);
162      }
163      try {
164        this.scan = scan != null ? TableMapReduceUtil.convertScanToString(scan) : "";
165      } catch (IOException e) {
166        LOG.warn("Failed to convert Scan to String", e);
167      }
168
169      this.restoreDir = restoreDir.toString();
170    }
171
172    public TableDescriptor getHtd() {
173      return htd;
174    }
175
176    public String getScan() {
177      return scan;
178    }
179
180    public String getRestoreDir() {
181      return restoreDir;
182    }
183
184    public long getLength() {
185      // TODO: We can obtain the file sizes of the snapshot here.
186      return 0;
187    }
188
189    public String[] getLocations() {
190      return locations;
191    }
192
193    public TableDescriptor getTableDescriptor() {
194      return htd;
195    }
196
197    public RegionInfo getRegionInfo() {
198      return regionInfo;
199    }
200
201    // TODO: We should have ProtobufSerialization in Hadoop, and directly use PB objects instead of
202    // doing this wrapping with Writables.
203    @Override
204    public void write(DataOutput out) throws IOException {
205      TableSnapshotRegionSplit.Builder builder = TableSnapshotRegionSplit.newBuilder()
206        .setTable(ProtobufUtil.toTableSchema(htd)).setRegion(ProtobufUtil.toRegionInfo(regionInfo));
207
208      for (String location : locations) {
209        builder.addLocations(location);
210      }
211
212      TableSnapshotRegionSplit split = builder.build();
213
214      ByteArrayOutputStream baos = new ByteArrayOutputStream();
215      split.writeTo(baos);
216      baos.close();
217      byte[] buf = baos.toByteArray();
218      out.writeInt(buf.length);
219      out.write(buf);
220
221      Bytes.writeByteArray(out, Bytes.toBytes(scan));
222      Bytes.writeByteArray(out, Bytes.toBytes(restoreDir));
223
224    }
225
226    @Override
227    public void readFields(DataInput in) throws IOException {
228      int len = in.readInt();
229      byte[] buf = new byte[len];
230      in.readFully(buf);
231      TableSnapshotRegionSplit split = TableSnapshotRegionSplit.parser().parseFrom(buf);
232      this.htd = ProtobufUtil.toTableDescriptor(split.getTable());
233      this.regionInfo = ProtobufUtil.toRegionInfo(split.getRegion());
234      List<String> locationsList = split.getLocationsList();
235      this.locations = locationsList.toArray(new String[locationsList.size()]);
236
237      this.scan = Bytes.toString(Bytes.readByteArray(in));
238      this.restoreDir = Bytes.toString(Bytes.readByteArray(in));
239    }
240  }
241
242  /**
243   * Implementation class for RecordReader logic common between mapred and mapreduce.
244   */
245  public static class RecordReader {
246    private InputSplit split;
247    private Scan scan;
248    private Result result = null;
249    private ImmutableBytesWritable row = null;
250    private ClientSideRegionScanner scanner;
251    private int numOfCompleteRows = 0;
252    private int rowLimitPerSplit;
253
254    public ClientSideRegionScanner getScanner() {
255      return scanner;
256    }
257
258    public void initialize(InputSplit split, Configuration conf) throws IOException {
259      this.scan = TableMapReduceUtil.convertStringToScan(split.getScan());
260      this.split = split;
261      this.rowLimitPerSplit = conf.getInt(SNAPSHOT_INPUTFORMAT_ROW_LIMIT_PER_INPUTSPLIT, 0);
262      TableDescriptor htd = split.htd;
263      RegionInfo hri = this.split.getRegionInfo();
264      FileSystem fs = CommonFSUtils.getCurrentFileSystem(conf);
265
266      // region is immutable, this should be fine,
267      // otherwise we have to set the thread read point
268      scan.setIsolationLevel(IsolationLevel.READ_UNCOMMITTED);
269      // disable caching of data blocks
270      scan.setCacheBlocks(false);
271
272      scanner =
273        new ClientSideRegionScanner(conf, fs, new Path(split.restoreDir), htd, hri, scan, null);
274    }
275
276    public boolean nextKeyValue() throws IOException {
277      result = scanner.next();
278      if (result == null) {
279        // we are done
280        return false;
281      }
282
283      if (rowLimitPerSplit > 0 && ++this.numOfCompleteRows > rowLimitPerSplit) {
284        return false;
285      }
286      if (this.row == null) {
287        this.row = new ImmutableBytesWritable();
288      }
289      this.row.set(result.getRow());
290      return true;
291    }
292
293    public ImmutableBytesWritable getCurrentKey() {
294      return row;
295    }
296
297    public Result getCurrentValue() {
298      return result;
299    }
300
301    public long getPos() {
302      return 0;
303    }
304
305    public float getProgress() {
306      return 0; // TODO: use total bytes to estimate
307    }
308
309    public void close() {
310      if (this.scanner != null) {
311        this.scanner.close();
312      }
313    }
314  }
315
316  public static List<InputSplit> getSplits(Configuration conf) throws IOException {
317    String snapshotName = getSnapshotName(conf);
318
319    Path rootDir = CommonFSUtils.getRootDir(conf);
320    FileSystem fs = rootDir.getFileSystem(conf);
321
322    SnapshotManifest manifest = getSnapshotManifest(conf, snapshotName, rootDir, fs);
323
324    List<RegionInfo> regionInfos = getRegionInfosFromManifest(manifest);
325
326    // TODO: mapred does not support scan as input API. Work around for now.
327    Scan scan = extractScanFromConf(conf);
328    // the temp dir where the snapshot is restored
329    Path restoreDir = new Path(conf.get(RESTORE_DIR_KEY));
330
331    RegionSplitter.SplitAlgorithm splitAlgo = getSplitAlgo(conf);
332
333    int numSplits = conf.getInt(NUM_SPLITS_PER_REGION, 1);
334
335    return getSplits(scan, manifest, regionInfos, restoreDir, conf, splitAlgo, numSplits);
336  }
337
338  public static RegionSplitter.SplitAlgorithm getSplitAlgo(Configuration conf) throws IOException {
339    String splitAlgoClassName = conf.get(SPLIT_ALGO);
340    if (splitAlgoClassName == null) {
341      return null;
342    }
343    try {
344      return Class.forName(splitAlgoClassName).asSubclass(RegionSplitter.SplitAlgorithm.class)
345        .getDeclaredConstructor().newInstance();
346    } catch (ClassNotFoundException | InstantiationException | IllegalAccessException
347      | NoSuchMethodException | InvocationTargetException e) {
348      throw new IOException("SplitAlgo class " + splitAlgoClassName + " is not found", e);
349    }
350  }
351
352  public static List<RegionInfo> getRegionInfosFromManifest(SnapshotManifest manifest) {
353    List<SnapshotRegionManifest> regionManifests = manifest.getRegionManifests();
354    if (regionManifests == null) {
355      throw new IllegalArgumentException("Snapshot seems empty");
356    }
357
358    List<RegionInfo> regionInfos = Lists.newArrayListWithCapacity(regionManifests.size());
359
360    for (SnapshotRegionManifest regionManifest : regionManifests) {
361      RegionInfo hri = ProtobufUtil.toRegionInfo(regionManifest.getRegionInfo());
362      if (hri.isOffline() && (hri.isSplit() || hri.isSplitParent())) {
363        continue;
364      }
365      // The mob region is a dummy region used only to organise mob files under mobdir. It has no
366      // region directory under the table dir to open. See HBASE-30365.
367      if (MobUtils.isMobRegionInfo(hri)) {
368        continue;
369      }
370      regionInfos.add(hri);
371    }
372    return regionInfos;
373  }
374
375  public static SnapshotManifest getSnapshotManifest(Configuration conf, String snapshotName,
376    Path rootDir, FileSystem fs) throws IOException {
377    Path snapshotDir = SnapshotDescriptionUtils.getCompletedSnapshotDir(snapshotName, rootDir);
378    SnapshotDescription snapshotDesc = SnapshotDescriptionUtils.readSnapshotInfo(fs, snapshotDir);
379    return SnapshotManifest.open(conf, fs, snapshotDir, snapshotDesc);
380  }
381
382  public static Scan extractScanFromConf(Configuration conf) throws IOException {
383    Scan scan = null;
384    if (conf.get(TableInputFormat.SCAN) != null) {
385      scan = TableMapReduceUtil.convertStringToScan(conf.get(TableInputFormat.SCAN));
386    } else if (conf.get(org.apache.hadoop.hbase.mapred.TableInputFormat.COLUMN_LIST) != null) {
387      String[] columns =
388        conf.get(org.apache.hadoop.hbase.mapred.TableInputFormat.COLUMN_LIST).split(" ");
389      scan = new Scan();
390      for (String col : columns) {
391        scan.addFamily(Bytes.toBytes(col));
392      }
393    } else {
394      throw new IllegalArgumentException("Unable to create scan");
395    }
396
397    if (scan.getReadType() == ReadType.DEFAULT) {
398      LOG.info(
399        "Provided Scan has DEFAULT ReadType," + " updating STREAM for Snapshot-based InputFormat");
400      // Update the "DEFAULT" ReadType to be "STREAM" to try to improve the default case.
401      scan.setReadType(conf.getEnum(SNAPSHOT_INPUTFORMAT_SCANNER_READTYPE,
402        SNAPSHOT_INPUTFORMAT_SCANNER_READTYPE_DEFAULT));
403    }
404
405    return scan;
406  }
407
408  public static List<InputSplit> getSplits(Scan scan, SnapshotManifest manifest,
409    List<RegionInfo> regionManifests, Path restoreDir, Configuration conf) throws IOException {
410    return getSplits(scan, manifest, regionManifests, restoreDir, conf, null, 1);
411  }
412
413  public static List<InputSplit> getSplits(Scan scan, SnapshotManifest manifest,
414    List<RegionInfo> regionManifests, Path restoreDir, Configuration conf,
415    RegionSplitter.SplitAlgorithm sa, int numSplits) throws IOException {
416    // load table descriptor
417    TableDescriptor htd = manifest.getTableDescriptor();
418
419    Path tableDir = CommonFSUtils.getTableDir(restoreDir, htd.getTableName());
420
421    boolean localityEnabled = conf.getBoolean(SNAPSHOT_INPUTFORMAT_LOCALITY_ENABLED_KEY,
422      SNAPSHOT_INPUTFORMAT_LOCALITY_ENABLED_DEFAULT);
423
424    boolean scanMetricsEnabled = conf.getBoolean(SNAPSHOT_INPUTFORMAT_SCAN_METRICS_ENABLED,
425      SNAPSHOT_INPUTFORMAT_SCAN_METRICS_ENABLED_DEFAULT);
426    scan.setScanMetricsEnabled(scanMetricsEnabled);
427
428    boolean useRegionLoc = conf.getBoolean(SNAPSHOT_INPUTFORMAT_LOCALITY_BY_REGION_LOCATION,
429      SNAPSHOT_INPUTFORMAT_LOCALITY_BY_REGION_LOCATION_DEFAULT);
430
431    Connection connection = null;
432    RegionLocator regionLocator = null;
433    if (localityEnabled && useRegionLoc) {
434      Configuration newConf = new Configuration(conf);
435      newConf.setInt("hbase.hconnection.threads.max", 1);
436      try {
437        connection = ConnectionFactory.createConnection(newConf);
438        regionLocator = connection.getRegionLocator(htd.getTableName());
439
440        /* Get all locations for the table and cache it */
441        regionLocator.getAllRegionLocations();
442      } finally {
443        if (connection != null) {
444          connection.close();
445        }
446      }
447    }
448
449    List<InputSplit> splits = new ArrayList<>();
450    for (RegionInfo hri : regionManifests) {
451      // load region descriptor
452      List<String> hosts = null;
453      if (localityEnabled) {
454        if (regionLocator != null) {
455          /* Get Location from the local cache */
456          HRegionLocation location = regionLocator.getRegionLocation(hri.getStartKey(), false);
457
458          hosts = new ArrayList<>(1);
459          hosts.add(location.getHostname());
460        } else {
461          hosts = calculateLocationsForInputSplit(conf, htd, hri, tableDir);
462        }
463      }
464
465      if (numSplits > 1) {
466        byte[][] sp = sa.split(hri.getStartKey(), hri.getEndKey(), numSplits, true);
467        for (int i = 0; i < sp.length - 1; i++) {
468          if (
469            PrivateCellUtil.overlappingKeys(scan.getStartRow(), scan.getStopRow(), sp[i], sp[i + 1])
470          ) {
471
472            Scan boundedScan = new Scan(scan);
473            if (scan.getStartRow().length == 0) {
474              boundedScan.withStartRow(sp[i]);
475            } else {
476              boundedScan.withStartRow(
477                Bytes.compareTo(scan.getStartRow(), sp[i]) > 0 ? scan.getStartRow() : sp[i]);
478            }
479
480            if (scan.getStopRow().length == 0) {
481              boundedScan.withStopRow(sp[i + 1]);
482            } else {
483              boundedScan.withStopRow(
484                Bytes.compareTo(scan.getStopRow(), sp[i + 1]) < 0 ? scan.getStopRow() : sp[i + 1]);
485            }
486
487            splits.add(new InputSplit(htd, hri, hosts, boundedScan, restoreDir));
488          }
489        }
490      } else {
491        if (
492          PrivateCellUtil.overlappingKeys(scan.getStartRow(), scan.getStopRow(), hri.getStartKey(),
493            hri.getEndKey())
494        ) {
495
496          splits.add(new InputSplit(htd, hri, hosts, scan, restoreDir));
497        }
498      }
499    }
500
501    return splits;
502  }
503
504  /**
505   * Compute block locations for snapshot files (which will get the locations for referred hfiles)
506   * only when localityEnabled is true.
507   */
508  private static List<String> calculateLocationsForInputSplit(Configuration conf,
509    TableDescriptor htd, RegionInfo hri, Path tableDir) throws IOException {
510    return getBestLocations(conf, HRegion.computeHDFSBlocksDistribution(conf, htd, hri, tableDir));
511  }
512
513  /**
514   * This computes the locations to be passed from the InputSplit. MR/Yarn schedulers does not take
515   * weights into account, thus will treat every location passed from the input split as equal. We
516   * do not want to blindly pass all the locations, since we are creating one split per region, and
517   * the region's blocks are all distributed throughout the cluster unless favorite node assignment
518   * is used. On the expected stable case, only one location will contain most of the blocks as
519   * local. On the other hand, in favored node assignment, 3 nodes will contain highly local blocks.
520   * Here we are doing a simple heuristic, where we will pass all hosts which have at least 80%
521   * (hbase.tablesnapshotinputformat.locality.cutoff.multiplier) as much block locality as the top
522   * host with the best locality. Return at most numTopsAtMost locations if there are more than
523   * that.
524   */
525  private static List<String> getBestLocations(Configuration conf,
526    HDFSBlocksDistribution blockDistribution, int numTopsAtMost) {
527    HostAndWeight[] hostAndWeights = blockDistribution.getTopHostsWithWeights();
528
529    if (hostAndWeights.length == 0) { // no matter what numTopsAtMost is
530      return null;
531    }
532
533    if (numTopsAtMost < 1) { // invalid if numTopsAtMost < 1, correct it to be 1
534      numTopsAtMost = 1;
535    }
536    int top = Math.min(numTopsAtMost, hostAndWeights.length);
537    List<String> locations = new ArrayList<>(top);
538    HostAndWeight topHost = hostAndWeights[0];
539    locations.add(topHost.getHost());
540
541    if (top == 1) { // only care about the top host
542      return locations;
543    }
544
545    // When top >= 2,
546    // do the heuristic: filter all hosts which have at least cutoffMultiplier % of block locality
547    double cutoffMultiplier =
548      conf.getFloat(LOCALITY_CUTOFF_MULTIPLIER, DEFAULT_LOCALITY_CUTOFF_MULTIPLIER);
549
550    double filterWeight = topHost.getWeight() * cutoffMultiplier;
551
552    for (int i = 1; i <= top - 1; i++) {
553      if (hostAndWeights[i].getWeight() >= filterWeight) {
554        locations.add(hostAndWeights[i].getHost());
555      } else {
556        // As hostAndWeights is in descending order,
557        // we could break the loop as long as we meet a weight which is less than filterWeight.
558        break;
559      }
560    }
561
562    return locations;
563  }
564
565  public static List<String> getBestLocations(Configuration conf,
566    HDFSBlocksDistribution blockDistribution) {
567    // 3 nodes will contain highly local blocks. So default to 3.
568    return getBestLocations(conf, blockDistribution, 3);
569  }
570
571  public static String getSnapshotName(Configuration conf) {
572    String snapshotName = conf.get(SNAPSHOT_NAME_KEY);
573    if (snapshotName == null) {
574      throw new IllegalArgumentException("Snapshot name must be provided");
575    }
576    return snapshotName;
577  }
578
579  /**
580   * Configures the job to use TableSnapshotInputFormat to read from a snapshot.
581   * @param conf         the job to configuration
582   * @param snapshotName the name of the snapshot to read from
583   * @param restoreDir   a temporary directory to restore the snapshot into. Current user should
584   *                     have write permissions to this directory, and this should not be a
585   *                     subdirectory of rootdir. After the job is finished, restoreDir can be
586   *                     deleted.
587   * @throws IOException if an error occurs
588   */
589  public static void setInput(Configuration conf, String snapshotName, Path restoreDir)
590    throws IOException {
591    setInput(conf, snapshotName, restoreDir, null, 1);
592  }
593
594  /**
595   * Configures the job to use TableSnapshotInputFormat to read from a snapshot.
596   * @param conf               the job to configure
597   * @param snapshotName       the name of the snapshot to read from
598   * @param restoreDir         a temporary directory to restore the snapshot into. Current user
599   *                           should have write permissions to this directory, and this should not
600   *                           be a subdirectory of rootdir. After the job is finished, restoreDir
601   *                           can be deleted.
602   * @param numSplitsPerRegion how many input splits to generate per one region
603   * @param splitAlgo          SplitAlgorithm to be used when generating InputSplits
604   * @throws IOException if an error occurs
605   */
606  public static void setInput(Configuration conf, String snapshotName, Path restoreDir,
607    RegionSplitter.SplitAlgorithm splitAlgo, int numSplitsPerRegion) throws IOException {
608    conf.set(SNAPSHOT_NAME_KEY, snapshotName);
609    if (numSplitsPerRegion < 1) {
610      throw new IllegalArgumentException(
611        "numSplits must be >= 1, " + "illegal numSplits : " + numSplitsPerRegion);
612    }
613    if (splitAlgo == null && numSplitsPerRegion > 1) {
614      throw new IllegalArgumentException("Split algo can't be null when numSplits > 1");
615    }
616    if (splitAlgo != null) {
617      conf.set(SPLIT_ALGO, splitAlgo.getClass().getName());
618    }
619    conf.setInt(NUM_SPLITS_PER_REGION, numSplitsPerRegion);
620    Path rootDir = CommonFSUtils.getRootDir(conf);
621    FileSystem fs = rootDir.getFileSystem(conf);
622
623    restoreDir = new Path(restoreDir, UUID.randomUUID().toString());
624
625    RestoreSnapshotHelper.copySnapshotForScanner(conf, fs, rootDir, restoreDir, snapshotName);
626    conf.set(RESTORE_DIR_KEY, restoreDir.toString());
627  }
628
629  /**
630   * clean restore directory after snapshot scan job
631   * @param job          the snapshot scan job
632   * @param snapshotName the name of the snapshot to read from
633   * @throws IOException if an error occurs
634   */
635  public static void cleanRestoreDir(Job job, String snapshotName) throws IOException {
636    Configuration conf = job.getConfiguration();
637    Path restoreDir = new Path(conf.get(RESTORE_DIR_KEY));
638    FileSystem fs = restoreDir.getFileSystem(conf);
639    if (!fs.exists(restoreDir)) {
640      LOG.warn("{} doesn't exist on file system, maybe it's already been cleaned", restoreDir);
641      return;
642    }
643    if (!fs.delete(restoreDir, true)) {
644      LOG.warn("Failed clean restore dir {} for snapshot {}", restoreDir, snapshotName);
645    }
646    LOG.debug("Clean restore directory {} for {}", restoreDir, snapshotName);
647  }
648}