Java Code Examples for org.apache.hadoop.io.DataOutputBuffer#reset()

private <E> E makeCopyForPassByValue(Serialization<E> serialization,
                                      E obj) throws IOException {
  Serializer<E> ser =
    serialization.getSerializer(GenericsUtil.getClass(obj));
  Deserializer<E> deser =
    serialization.getDeserializer(GenericsUtil.getClass(obj));

  DataOutputBuffer dof = threadLocalDataOutputBuffer.get();

  dof.reset();
  ser.open(dof);
  ser.serialize(obj);
  ser.close();
  obj = ReflectionUtils.newInstance(GenericsUtil.getClass(obj),
                                    getChainJobConf());
  ByteArrayInputStream bais =
    new ByteArrayInputStream(dof.getData(), 0, dof.getLength());
  deser.open(bais);
  deser.deserialize(obj);
  deser.close();
  return obj;
}
 

@Test
public void testReadWriteReplicaState() {
  try {
    DataOutputBuffer out = new DataOutputBuffer();
    DataInputBuffer in = new DataInputBuffer();
    for (HdfsServerConstants.ReplicaState repState : HdfsServerConstants.ReplicaState
        .values()) {
      repState.write(out);
      in.reset(out.getData(), out.getLength());
      HdfsServerConstants.ReplicaState result = HdfsServerConstants.ReplicaState
          .read(in);
      assertTrue("testReadWrite error !!!", repState == result);
      out.reset();
      in.reset();
    }
  } catch (Exception ex) {
    fail("testReadWrite ex error ReplicaState");
  }
}
 

private <E> E makeCopyForPassByValue(Serialization<E> serialization,
                                      E obj) throws IOException {
  Serializer<E> ser =
    serialization.getSerializer(GenericsUtil.getClass(obj));
  Deserializer<E> deser =
    serialization.getDeserializer(GenericsUtil.getClass(obj));

  DataOutputBuffer dof = threadLocalDataOutputBuffer.get();

  dof.reset();
  ser.open(dof);
  ser.serialize(obj);
  ser.close();
  obj = ReflectionUtils.newInstance(GenericsUtil.getClass(obj),
                                    getChainJobConf());
  ByteArrayInputStream bais =
    new ByteArrayInputStream(dof.getData(), 0, dof.getLength());
  deser.open(bais);
  deser.deserialize(obj);
  deser.close();
  return obj;
}
 

private static void run(Query query) throws IOException {

    DataOutputBuffer buffer = new DataOutputBuffer();
    DataInputBuffer in = new DataInputBuffer();
    QueryWritable outQw = new QueryWritable();
    QueryWritable inQw = new QueryWritable();

    long s = System.nanoTime();
    int count = 100000;
    for (int i = 0; i < count; i++) {
      outQw.setQuery(query);
      outQw.write(buffer);

      in.reset(buffer.getData(), 0, buffer.getLength());
      inQw.readFields(in);

      buffer.reset();
    }
    long e = System.nanoTime();
    System.out.println((e - s) / 1000000.0 / (double) count);
    // System.out.println((e - s) / (double) count);
  }
 

public void write(DataOutput out) throws IOException {
  if (firstKey == null) {
    Utils.writeVInt(out, 0);
    return;
  }

  DataOutputBuffer dob = new DataOutputBuffer();
  Utils.writeVInt(dob, firstKey.size());
  dob.write(firstKey.buffer());
  Utils.writeVInt(out, dob.size());
  out.write(dob.getData(), 0, dob.getLength());

  for (TFileIndexEntry entry : index) {
    dob.reset();
    entry.write(dob);
    Utils.writeVInt(out, dob.getLength());
    out.write(dob.getData(), 0, dob.getLength());
  }
}
 

@Test
public void testReadWriteReplicaState() {
  try {
    DataOutputBuffer out = new DataOutputBuffer();
    DataInputBuffer in = new DataInputBuffer();
    for (HdfsServerConstants.ReplicaState repState : HdfsServerConstants.ReplicaState
        .values()) {
      repState.write(out);
      in.reset(out.getData(), out.getLength());
      HdfsServerConstants.ReplicaState result = HdfsServerConstants.ReplicaState
          .read(in);
      assertTrue("testReadWrite error !!!", repState == result);
      out.reset();
      in.reset();
    }
  } catch (Exception ex) {
    fail("testReadWrite ex error ReplicaState");
  }
}
 

public static void copy(DataInputBuffer src, DataOutputBuffer dst) throws IOException {
  byte[] b1 = src.getData();
  int s1 = src.getPosition();
  int l1 = src.getLength();
  dst.reset();
  dst.write(b1, s1, l1 - s1);
}
 

/** Fills up the rawKey object with the key returned by the Reader
 * @return true if there is a key returned; false, otherwise
 * @throws IOException
 */
public boolean nextRawKey() throws IOException {
  if (in == null) {
    int bufferSize = getBufferSize(conf); 
    Reader reader = new Reader(conf,
                               Reader.file(segmentPathName), 
                               Reader.bufferSize(bufferSize),
                               Reader.start(segmentOffset), 
                               Reader.length(segmentLength));
  
    //sometimes we ignore syncs especially for temp merge files
    if (ignoreSync) reader.ignoreSync();

    if (reader.getKeyClass() != keyClass)
      throw new IOException("wrong key class: " + reader.getKeyClass() +
                            " is not " + keyClass);
    if (reader.getValueClass() != valClass)
      throw new IOException("wrong value class: "+reader.getValueClass()+
                            " is not " + valClass);
    this.in = reader;
    rawKey = new DataOutputBuffer();
  }
  rawKey.reset();
  int keyLength = 
    in.nextRawKey(rawKey);
  return (keyLength >= 0);
}
 

void GzipConcatTest(Configuration conf,
    Class<? extends Decompressor> decomClass) throws IOException {
  Random r = new Random();
  long seed = r.nextLong();
  r.setSeed(seed);
  LOG.info(decomClass + " seed: " + seed);

  final int CONCAT = r.nextInt(4) + 3;
  final int BUFLEN = 128 * 1024;
  DataOutputBuffer dflbuf = new DataOutputBuffer();
  DataOutputBuffer chkbuf = new DataOutputBuffer();
  byte[] b = new byte[BUFLEN];
  for (int i = 0; i < CONCAT; ++i) {
    GZIPOutputStream gzout = new GZIPOutputStream(dflbuf);
    r.nextBytes(b);
    int len = r.nextInt(BUFLEN);
    int off = r.nextInt(BUFLEN - len);
    chkbuf.write(b, off, len);
    gzout.write(b, off, len);
    gzout.close();
  }
  final byte[] chk = Arrays.copyOf(chkbuf.getData(), chkbuf.getLength());

  CompressionCodec codec = ReflectionUtils.newInstance(GzipCodec.class, conf);
  Decompressor decom = codec.createDecompressor();
  assertNotNull(decom);
  assertEquals(decomClass, decom.getClass());
  DataInputBuffer gzbuf = new DataInputBuffer();
  gzbuf.reset(dflbuf.getData(), dflbuf.getLength());
  InputStream gzin = codec.createInputStream(gzbuf, decom);

  dflbuf.reset();
  IOUtils.copyBytes(gzin, dflbuf, 4096);
  final byte[] dflchk = Arrays.copyOf(dflbuf.getData(), dflbuf.getLength());
  assertArrayEquals(chk, dflchk);
}
 

@Test
public void testGzipCompatibility() throws IOException {
  Random r = new Random();
  long seed = r.nextLong();
  r.setSeed(seed);
  LOG.info("seed: " + seed);

  DataOutputBuffer dflbuf = new DataOutputBuffer();
  GZIPOutputStream gzout = new GZIPOutputStream(dflbuf);
  byte[] b = new byte[r.nextInt(128 * 1024 + 1)];
  r.nextBytes(b);
  gzout.write(b);
  gzout.close();

  DataInputBuffer gzbuf = new DataInputBuffer();
  gzbuf.reset(dflbuf.getData(), dflbuf.getLength());

  Configuration conf = new Configuration();
  conf.setBoolean(CommonConfigurationKeys.IO_NATIVE_LIB_AVAILABLE_KEY, false);
  CompressionCodec codec = ReflectionUtils.newInstance(GzipCodec.class, conf);
  Decompressor decom = codec.createDecompressor();
  assertNotNull(decom);
  assertEquals(BuiltInGzipDecompressor.class, decom.getClass());
  InputStream gzin = codec.createInputStream(gzbuf, decom);

  dflbuf.reset();
  IOUtils.copyBytes(gzin, dflbuf, 4096);
  final byte[] dflchk = Arrays.copyOf(dflbuf.getData(), dflbuf.getLength());
  assertArrayEquals(b, dflchk);
}
 

static void printDebug(GridmixRecord a, GridmixRecord b) throws IOException {
  DataOutputBuffer out = new DataOutputBuffer();
  a.write(out);
  System.out.println("A " +
      Arrays.toString(Arrays.copyOf(out.getData(), out.getLength())));
  out.reset();
  b.write(out);
  System.out.println("B " +
      Arrays.toString(Arrays.copyOf(out.getData(), out.getLength())));
}
 

static void setSerialize(GridmixRecord x, long seed, int size,
    DataOutputBuffer out) throws IOException {
  x.setSeed(seed);
  x.setSize(size);
  out.reset();
  x.write(out);
}
 

@Test
public void testGzipCompatibility() throws IOException {
  Random r = new Random();
  long seed = r.nextLong();
  r.setSeed(seed);
  LOG.info("seed: " + seed);

  DataOutputBuffer dflbuf = new DataOutputBuffer();
  GZIPOutputStream gzout = new GZIPOutputStream(dflbuf);
  byte[] b = new byte[r.nextInt(128 * 1024 + 1)];
  r.nextBytes(b);
  gzout.write(b);
  gzout.close();

  DataInputBuffer gzbuf = new DataInputBuffer();
  gzbuf.reset(dflbuf.getData(), dflbuf.getLength());

  Configuration conf = new Configuration();
  conf.setBoolean(CommonConfigurationKeys.IO_NATIVE_LIB_AVAILABLE_KEY, false);
  CompressionCodec codec = ReflectionUtils.newInstance(GzipCodec.class, conf);
  Decompressor decom = codec.createDecompressor();
  assertNotNull(decom);
  assertEquals(BuiltInGzipDecompressor.class, decom.getClass());
  InputStream gzin = codec.createInputStream(gzbuf, decom);

  dflbuf.reset();
  IOUtils.copyBytes(gzin, dflbuf, 4096);
  final byte[] dflchk = Arrays.copyOf(dflbuf.getData(), dflbuf.getLength());
  assertArrayEquals(b, dflchk);
}
 

/**
 * create inmemory segments
 *
 * @return
 * @throws IOException
 */
@SuppressWarnings("unchecked")
public List<TezMerger.Segment> createInMemStreams() throws IOException {
  int numberOfStreams = Math.max(2, rnd.nextInt(10));
  LOG.info("No of streams : " + numberOfStreams);

  SerializationFactory serializationFactory = new SerializationFactory(conf);
  Serializer keySerializer = serializationFactory.getSerializer(keyClass);
  Serializer valueSerializer = serializationFactory.getSerializer(valClass);

  LocalDirAllocator localDirAllocator =
      new LocalDirAllocator(TezRuntimeFrameworkConfigs.LOCAL_DIRS);
  InputContext context = createTezInputContext();
  MergeManager mergeManager = new MergeManager(conf, fs, localDirAllocator,
      context, null, null, null, null, null, 1024 * 1024 * 10, null, false, -1);

  DataOutputBuffer keyBuf = new DataOutputBuffer();
  DataOutputBuffer valBuf = new DataOutputBuffer();
  DataInputBuffer keyIn = new DataInputBuffer();
  DataInputBuffer valIn = new DataInputBuffer();
  keySerializer.open(keyBuf);
  valueSerializer.open(valBuf);

  List<TezMerger.Segment> segments = new LinkedList<TezMerger.Segment>();
  for (int i = 0; i < numberOfStreams; i++) {
    BoundedByteArrayOutputStream bout = new BoundedByteArrayOutputStream(1024 * 1024);
    InMemoryWriter writer =
        new InMemoryWriter(bout);
    Map<Writable, Writable> data = createData();
    //write data
    for (Map.Entry<Writable, Writable> entry : data.entrySet()) {
      keySerializer.serialize(entry.getKey());
      valueSerializer.serialize(entry.getValue());
      keyIn.reset(keyBuf.getData(), 0, keyBuf.getLength());
      valIn.reset(valBuf.getData(), 0, valBuf.getLength());
      writer.append(keyIn, valIn);
      originalData.put(entry.getKey(), entry.getValue());
      keyBuf.reset();
      valBuf.reset();
      keyIn.reset();
      valIn.reset();
    }
    IFile.Reader reader = new InMemoryReader(mergeManager, null, bout.getBuffer(), 0,
        bout.getBuffer().length);
    segments.add(new TezMerger.Segment(reader, null));

    data.clear();
    writer.close();
  }
  return segments;
}
 

@Test
public void testRecoveryFromOtherVersions() throws IOException {
  final String user = "someuser";
  final ApplicationId appId = ApplicationId.newInstance(12345, 1);
  final File tmpDir = new File(System.getProperty("test.build.data",
      System.getProperty("java.io.tmpdir")),
      TestShuffleHandler.class.getName());
  Configuration conf = new Configuration();
  conf.setInt(ShuffleHandler.SHUFFLE_PORT_CONFIG_KEY, 0);
  conf.setInt(ShuffleHandler.MAX_SHUFFLE_CONNECTIONS, 3);
  ShuffleHandler shuffle = new ShuffleHandler();
  // emulate aux services startup with recovery enabled
  shuffle.setRecoveryPath(new Path(tmpDir.toString()));
  tmpDir.mkdirs();
  try {
    shuffle.init(conf);
    shuffle.start();

    // setup a shuffle token for an application
    DataOutputBuffer outputBuffer = new DataOutputBuffer();
    outputBuffer.reset();
    Token<JobTokenIdentifier> jt = new Token<JobTokenIdentifier>(
        "identifier".getBytes(), "password".getBytes(), new Text(user),
        new Text("shuffleService"));
    jt.write(outputBuffer);
    shuffle.initializeApplication(new ApplicationInitializationContext(user,
        appId, ByteBuffer.wrap(outputBuffer.getData(), 0,
            outputBuffer.getLength())));

    // verify we are authorized to shuffle
    int rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_OK, rc);

    // emulate shuffle handler restart
    shuffle.close();
    shuffle = new ShuffleHandler();
    shuffle.setRecoveryPath(new Path(tmpDir.toString()));
    shuffle.init(conf);
    shuffle.start();

    // verify we are still authorized to shuffle to the old application
    rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_OK, rc);
    Version version = Version.newInstance(1, 0);
    Assert.assertEquals(version, shuffle.getCurrentVersion());
  
    // emulate shuffle handler restart with compatible version
    Version version11 = Version.newInstance(1, 1);
    // update version info before close shuffle
    shuffle.storeVersion(version11);
    Assert.assertEquals(version11, shuffle.loadVersion());
    shuffle.close();
    shuffle = new ShuffleHandler();
    shuffle.setRecoveryPath(new Path(tmpDir.toString()));
    shuffle.init(conf);
    shuffle.start();
    // shuffle version will be override by CURRENT_VERSION_INFO after restart
    // successfully.
    Assert.assertEquals(version, shuffle.loadVersion());
    // verify we are still authorized to shuffle to the old application
    rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_OK, rc);
  
    // emulate shuffle handler restart with incompatible version
    Version version21 = Version.newInstance(2, 1);
    shuffle.storeVersion(version21);
    Assert.assertEquals(version21, shuffle.loadVersion());
    shuffle.close();
    shuffle = new ShuffleHandler();
    shuffle.setRecoveryPath(new Path(tmpDir.toString()));
    shuffle.init(conf);
  
    try {
      shuffle.start();
      Assert.fail("Incompatible version, should expect fail here.");
    } catch (ServiceStateException e) {
      Assert.assertTrue("Exception message mismatch", 
      e.getMessage().contains("Incompatible version for state DB schema:"));
    } 
  
  } finally {
    if (shuffle != null) {
      shuffle.close();
    }
    FileUtil.fullyDelete(tmpDir);
  }
}
 

@Test
public void testRecovery() throws IOException {
  final String user = "someuser";
  final ApplicationId appId = ApplicationId.newInstance(12345, 1);
  final JobID jobId = JobID.downgrade(TypeConverter.fromYarn(appId));
  final File tmpDir = new File(System.getProperty("test.build.data",
      System.getProperty("java.io.tmpdir")),
      TestShuffleHandler.class.getName());
  Configuration conf = new Configuration();
  conf.setInt(ShuffleHandler.SHUFFLE_PORT_CONFIG_KEY, 0);
  conf.setInt(ShuffleHandler.MAX_SHUFFLE_CONNECTIONS, 3);
  ShuffleHandler shuffle = new ShuffleHandler();
  // emulate aux services startup with recovery enabled
  shuffle.setRecoveryPath(new Path(tmpDir.toString()));
  tmpDir.mkdirs();
  try {
    shuffle.init(conf);
    shuffle.start();

    // setup a shuffle token for an application
    DataOutputBuffer outputBuffer = new DataOutputBuffer();
    outputBuffer.reset();
    Token<JobTokenIdentifier> jt = new Token<JobTokenIdentifier>(
        "identifier".getBytes(), "password".getBytes(), new Text(user),
        new Text("shuffleService"));
    jt.write(outputBuffer);
    shuffle.initializeApplication(new ApplicationInitializationContext(user,
        appId, ByteBuffer.wrap(outputBuffer.getData(), 0,
          outputBuffer.getLength())));

    // verify we are authorized to shuffle
    int rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_OK, rc);

    // emulate shuffle handler restart
    shuffle.close();
    shuffle = new ShuffleHandler();
    shuffle.setRecoveryPath(new Path(tmpDir.toString()));
    shuffle.init(conf);
    shuffle.start();

    // verify we are still authorized to shuffle to the old application
    rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_OK, rc);

    // shutdown app and verify access is lost
    shuffle.stopApplication(new ApplicationTerminationContext(appId));
    rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_UNAUTHORIZED, rc);

    // emulate shuffle handler restart
    shuffle.close();
    shuffle = new ShuffleHandler();
    shuffle.setRecoveryPath(new Path(tmpDir.toString()));
    shuffle.init(conf);
    shuffle.start();

    // verify we still don't have access
    rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_UNAUTHORIZED, rc);
  } finally {
    if (shuffle != null) {
      shuffle.close();
    }
    FileUtil.fullyDelete(tmpDir);
  }
}
 

@SuppressWarnings("unchecked")
private <T extends org.apache.hadoop.mapreduce.InputSplit> 
int writeNewSplits(JobContext job, Path submitSplitFile
                   ) throws IOException, InterruptedException, 
                            ClassNotFoundException {
  JobConf conf = job.getJobConf();
  org.apache.hadoop.mapreduce.InputFormat<?,?> input =
    ReflectionUtils.newInstance(job.getInputFormatClass(), job.getJobConf());
  
  List<org.apache.hadoop.mapreduce.InputSplit> splits = input.getSplits(job);
  T[] array = (T[])
    splits.toArray(new org.apache.hadoop.mapreduce.InputSplit[splits.size()]);

  // sort the splits into order based on size, so that the biggest
  // go first
  Arrays.sort(array, new NewSplitComparator());
  DataOutputStream out = writeSplitsFileHeader(conf, submitSplitFile, 
                                               array.length);
  try {
    if (array.length != 0) {
      DataOutputBuffer buffer = new DataOutputBuffer();
      RawSplit rawSplit = new RawSplit();
      SerializationFactory factory = new SerializationFactory(conf);
      Serializer<T> serializer = 
        factory.getSerializer((Class<T>) array[0].getClass());
      serializer.open(buffer);
      for(T split: array) {
        rawSplit.setClassName(split.getClass().getName());
        buffer.reset();
        serializer.serialize(split);
        rawSplit.setDataLength(split.getLength());
        rawSplit.setBytes(buffer.getData(), 0, buffer.getLength());
        rawSplit.setLocations(split.getLocations());
        rawSplit.write(out);
      }
      serializer.close();
    }
  } finally {
    out.close();
  }
  return array.length;
}
 

@Test
public void testRecovery() throws IOException {
  final String user = "someuser";
  final ApplicationId appId = ApplicationId.newInstance(12345, 1);
  final JobID jobId = JobID.downgrade(TypeConverter.fromYarn(appId));
  final File tmpDir = new File(System.getProperty("test.build.data",
      System.getProperty("java.io.tmpdir")),
      TestShuffleHandler.class.getName());
  Configuration conf = new Configuration();
  conf.set(HADOOP_TMP_DIR, TEST_DIR.getAbsolutePath());
  conf.setInt(ShuffleHandler.SHUFFLE_PORT_CONFIG_KEY, 0);
  conf.setInt(ShuffleHandler.MAX_SHUFFLE_CONNECTIONS, 3);
  ShuffleHandler shuffle = new ShuffleHandler();
  // emulate aux services startup with recovery enabled
  shuffle.setRecoveryPath(new Path(tmpDir.toString()));
  tmpDir.mkdirs();
  try {
    shuffle.init(conf);
    shuffle.start();

    // setup a shuffle token for an application
    DataOutputBuffer outputBuffer = new DataOutputBuffer();
    outputBuffer.reset();
    Token<JobTokenIdentifier> jt = new Token<JobTokenIdentifier>(
        "identifier".getBytes(), "password".getBytes(), new Text(user),
        new Text("shuffleService"));
    jt.write(outputBuffer);
    shuffle.initializeApplication(new ApplicationInitializationContext(user,
        appId, ByteBuffer.wrap(outputBuffer.getData(), 0,
          outputBuffer.getLength())));

    // verify we are authorized to shuffle
    int rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_OK, rc);

    // emulate shuffle handler restart
    shuffle.close();
    shuffle = new ShuffleHandler();
    shuffle.setRecoveryPath(new Path(tmpDir.toString()));
    shuffle.init(conf);
    shuffle.start();

    // verify we are still authorized to shuffle to the old application
    rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_OK, rc);

    // shutdown app and verify access is lost
    shuffle.stopApplication(new ApplicationTerminationContext(appId));
    rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_UNAUTHORIZED, rc);

    // emulate shuffle handler restart
    shuffle.close();
    shuffle = new ShuffleHandler();
    shuffle.setRecoveryPath(new Path(tmpDir.toString()));
    shuffle.init(conf);
    shuffle.start();

    // verify we still don't have access
    rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_UNAUTHORIZED, rc);
  } finally {
    if (shuffle != null) {
      shuffle.close();
    }
    FileUtil.fullyDelete(tmpDir);
  }
}
 

@Test
public void testRecoveryFromOtherVersions() throws IOException {
  final String user = "someuser";
  final ApplicationId appId = ApplicationId.newInstance(12345, 1);
  final File tmpDir = new File(System.getProperty("test.build.data",
      System.getProperty("java.io.tmpdir")),
      TestShuffleHandler.class.getName());
  Configuration conf = new Configuration();
  conf.set(HADOOP_TMP_DIR, TEST_DIR.getAbsolutePath());
  conf.setInt(ShuffleHandler.SHUFFLE_PORT_CONFIG_KEY, 0);
  conf.setInt(ShuffleHandler.MAX_SHUFFLE_CONNECTIONS, 3);
  ShuffleHandler shuffle = new ShuffleHandler();
  // emulate aux services startup with recovery enabled
  shuffle.setRecoveryPath(new Path(tmpDir.toString()));
  tmpDir.mkdirs();
  try {
    shuffle.init(conf);
    shuffle.start();

    // setup a shuffle token for an application
    DataOutputBuffer outputBuffer = new DataOutputBuffer();
    outputBuffer.reset();
    Token<JobTokenIdentifier> jt = new Token<JobTokenIdentifier>(
        "identifier".getBytes(), "password".getBytes(), new Text(user),
        new Text("shuffleService"));
    jt.write(outputBuffer);
    shuffle.initializeApplication(new ApplicationInitializationContext(user,
        appId, ByteBuffer.wrap(outputBuffer.getData(), 0,
            outputBuffer.getLength())));

    // verify we are authorized to shuffle
    int rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_OK, rc);

    // emulate shuffle handler restart
    shuffle.close();
    shuffle = new ShuffleHandler();
    shuffle.setRecoveryPath(new Path(tmpDir.toString()));
    shuffle.init(conf);
    shuffle.start();

    // verify we are still authorized to shuffle to the old application
    rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_OK, rc);
    Version version = Version.newInstance(1, 0);
    Assert.assertEquals(version, shuffle.getCurrentVersion());

    // emulate shuffle handler restart with compatible version
    Version version11 = Version.newInstance(1, 1);
    // update version info before close shuffle
    shuffle.storeVersion(version11);
    Assert.assertEquals(version11, shuffle.loadVersion());
    shuffle.close();
    shuffle = new ShuffleHandler();
    shuffle.setRecoveryPath(new Path(tmpDir.toString()));
    shuffle.init(conf);
    shuffle.start();
    // shuffle version will be override by CURRENT_VERSION_INFO after restart
    // successfully.
    Assert.assertEquals(version, shuffle.loadVersion());
    // verify we are still authorized to shuffle to the old application
    rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_OK, rc);

    // emulate shuffle handler restart with incompatible version
    Version version21 = Version.newInstance(2, 1);
    shuffle.storeVersion(version21);
    Assert.assertEquals(version21, shuffle.loadVersion());
    shuffle.close();
    shuffle = new ShuffleHandler();
    shuffle.setRecoveryPath(new Path(tmpDir.toString()));
    shuffle.init(conf);
  
    try {
      shuffle.start();
      Assert.fail("Incompatible version, should expect fail here.");
    } catch (ServiceStateException e) {
      Assert.assertTrue("Exception message mismatch",
      e.getMessage().contains("Incompatible version for state DB schema:"));
    }

  } finally {
    if (shuffle != null) {
      shuffle.close();
    }
    FileUtil.fullyDelete(tmpDir);
  }
}
 

@Test
public void testRecoveryFromOtherVersions() throws IOException {
  final String user = "someuser";
  final ApplicationId appId = ApplicationId.newInstance(12345, 1);
  final File tmpDir = new File(System.getProperty("test.build.data",
      System.getProperty("java.io.tmpdir")),
      TestShuffleHandler.class.getName());
  Configuration conf = new Configuration();
  conf.setInt(ShuffleHandler.SHUFFLE_PORT_CONFIG_KEY, 0);
  conf.setInt(ShuffleHandler.MAX_SHUFFLE_CONNECTIONS, 3);
  ShuffleHandler shuffle = new ShuffleHandler();
  // emulate aux services startup with recovery enabled
  shuffle.setRecoveryPath(new Path(tmpDir.toString()));
  tmpDir.mkdirs();
  try {
    shuffle.init(conf);
    shuffle.start();

    // setup a shuffle token for an application
    DataOutputBuffer outputBuffer = new DataOutputBuffer();
    outputBuffer.reset();
    Token<JobTokenIdentifier> jt = new Token<JobTokenIdentifier>(
        "identifier".getBytes(), "password".getBytes(), new Text(user),
        new Text("shuffleService"));
    jt.write(outputBuffer);
    shuffle.initializeApplication(new ApplicationInitializationContext(user,
        appId, ByteBuffer.wrap(outputBuffer.getData(), 0,
            outputBuffer.getLength())));

    // verify we are authorized to shuffle
    int rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_OK, rc);

    // emulate shuffle handler restart
    shuffle.close();
    shuffle = new ShuffleHandler();
    shuffle.setRecoveryPath(new Path(tmpDir.toString()));
    shuffle.init(conf);
    shuffle.start();

    // verify we are still authorized to shuffle to the old application
    rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_OK, rc);
    Version version = Version.newInstance(1, 0);
    Assert.assertEquals(version, shuffle.getCurrentVersion());
  
    // emulate shuffle handler restart with compatible version
    Version version11 = Version.newInstance(1, 1);
    // update version info before close shuffle
    shuffle.storeVersion(version11);
    Assert.assertEquals(version11, shuffle.loadVersion());
    shuffle.close();
    shuffle = new ShuffleHandler();
    shuffle.setRecoveryPath(new Path(tmpDir.toString()));
    shuffle.init(conf);
    shuffle.start();
    // shuffle version will be override by CURRENT_VERSION_INFO after restart
    // successfully.
    Assert.assertEquals(version, shuffle.loadVersion());
    // verify we are still authorized to shuffle to the old application
    rc = getShuffleResponseCode(shuffle, jt);
    Assert.assertEquals(HttpURLConnection.HTTP_OK, rc);
  
    // emulate shuffle handler restart with incompatible version
    Version version21 = Version.newInstance(2, 1);
    shuffle.storeVersion(version21);
    Assert.assertEquals(version21, shuffle.loadVersion());
    shuffle.close();
    shuffle = new ShuffleHandler();
    shuffle.setRecoveryPath(new Path(tmpDir.toString()));
    shuffle.init(conf);
  
    try {
      shuffle.start();
      Assert.fail("Incompatible version, should expect fail here.");
    } catch (ServiceStateException e) {
      Assert.assertTrue("Exception message mismatch", 
      e.getMessage().contains("Incompatible version for state DB schema:"));
    } 
  
  } finally {
    if (shuffle != null) {
      shuffle.close();
    }
    FileUtil.fullyDelete(tmpDir);
  }
}