org.apache.twill.api.ResourceSpecification Java Examples

@Test
public void testFailureComplete() throws TimeoutException, ExecutionException, InterruptedException {
  TwillRunner twillRunner = getTwillRunner();

  // Start the app with an invalid ClassLoader. This will cause the AM fails to start.
  TwillController controller = twillRunner.prepare(new SleepTask(),
                                                   ResourceSpecification.Builder.with()
                                                     .setVirtualCores(1)
                                                     .setMemory(512, ResourceSpecification.SizeUnit.MEGA)
                                                     .setInstances(1).build())
    .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out, true)))
    .setClassLoader("InvalidClassLoader")
    .start();

  final CountDownLatch terminateLatch = new CountDownLatch(1);
  controller.onTerminated(new Runnable() {
    @Override
    public void run() {
      terminateLatch.countDown();
    }
  }, Threads.SAME_THREAD_EXECUTOR);

  Assert.assertTrue(terminateLatch.await(2, TimeUnit.MINUTES));
  Assert.assertEquals(ServiceController.TerminationStatus.FAILED, controller.getTerminationStatus());
}
 

@Override
public TwillSpecification configure() {
    return TwillSpecification.Builder.with()
            .setName(APPLICATION_NAME)
            .withRunnable()
                .add(PeriodicNotificationTwillRunnable.TWILL_RUNNABLE_NAME,
                        new PeriodicNotificationTwillRunnable(),
                        ResourceSpecification.Builder.with()
                            .setVirtualCores(2)
                            .setMemory(2, SizeUnit.GIGA)
                            .setInstances(1)
                            .build())
                .withLocalFiles()
                .add(PeriodicNotificationTwillRunnable.CONFIG_FILE_NAME, configFile)
                .apply()
            .anyOrder()
            .build();
}
 

@Test
public void testCodec() throws Exception {
  String expectedString =
          "{" +
                  "\"cores\":2," +
                  "\"memorySize\":1024," +
                  "\"instances\":2," +
                  "\"uplink\":100," +
                  "\"downlink\":100" +
          "}";
  final ResourceSpecification expected =
          new DefaultResourceSpecification(2, 1024, 2, 100, 100);
  final String actualString = gson.toJson(expected);
  Assert.assertEquals(expectedString, actualString);

  final JsonElement expectedJson = gson.toJsonTree(expected);
  final ResourceSpecification actual = gson.fromJson(expectedJson, DefaultResourceSpecification.class);
  final JsonElement actualJson = gson.toJsonTree(actual);

  Assert.assertEquals(expectedJson, actualJson);
  ReflectionAssert.assertLenientEquals(expected, actual);
}
 

private TwillRuntimeSpecificationAdapter() {
  gson = new GsonBuilder()
            .serializeNulls()
            .registerTypeAdapter(TwillRuntimeSpecification.class, new TwillRuntimeSpecificationCodec())
            .registerTypeAdapter(TwillSpecification.class, new TwillSpecificationCodec())
            .registerTypeAdapter(TwillSpecification.Order.class, new TwillSpecificationOrderCoder())
            .registerTypeAdapter(TwillSpecification.PlacementPolicy.class,
                                 new TwillSpecificationPlacementPolicyCoder())
            .registerTypeAdapter(EventHandlerSpecification.class, new EventHandlerSpecificationCoder())
            .registerTypeAdapter(RuntimeSpecification.class, new RuntimeSpecificationCodec())
            .registerTypeAdapter(TwillRunnableSpecification.class, new TwillRunnableSpecificationCodec())
            .registerTypeAdapter(ResourceSpecification.class, new ResourceSpecificationCodec())
            .registerTypeAdapter(LocalFile.class, new LocalFileCodec())
            .registerTypeAdapterFactory(new TwillSpecificationTypeAdapterFactory())
            .create();
}
 

private void maxRetriesRun(final int instances) throws TimeoutException, ExecutionException {
  TwillRunner runner = getTwillRunner();
  final int maxRetries = 3;
  final AtomicInteger retriesSeen = new AtomicInteger(0);

  ResourceSpecification resource = ResourceSpecification.Builder.with().setVirtualCores(1)
    .setMemory(512, ResourceSpecification.SizeUnit.MEGA).setInstances(instances).build();

  TwillController controller = runner.prepare(new FailingServer(), resource)
    .withMaxRetries(FailingServer.class.getSimpleName(), maxRetries)
    .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out, true)))
    .addLogHandler(new LogHandler() {
      @Override
      public void onLog(LogEntry logEntry) {
        if (logEntry.getMessage().contains("retries for instance")) {
          retriesSeen.incrementAndGet();
        }
      }
    })
    .start();

  controller.awaitTerminated(2, TimeUnit.MINUTES);
  Assert.assertEquals(maxRetries * instances, retriesSeen.get());
}
 

@Test
public void testTaskCompleted() throws InterruptedException, TimeoutException, ExecutionException {
  TwillRunner twillRunner = getTwillRunner();
  TwillController controller = twillRunner.prepare(new SleepTask(),
                                              ResourceSpecification.Builder.with()
                                                .setVirtualCores(1)
                                                .setMemory(512, ResourceSpecification.SizeUnit.MEGA)
                                                .setInstances(3).build())
                                          .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out, true)))
                                          .start();

  final CountDownLatch runLatch = new CountDownLatch(1);
  controller.onRunning(new Runnable() {
    @Override
    public void run() {
      runLatch.countDown();
    }
  }, Threads.SAME_THREAD_EXECUTOR);

  Assert.assertTrue(runLatch.await(1, TimeUnit.MINUTES));
  controller.awaitTerminated(1, TimeUnit.MINUTES);
  Assert.assertEquals(ServiceController.TerminationStatus.SUCCEEDED, controller.getTerminationStatus());
}
 

@Override
public TwillSpecification configure() {
  return TwillSpecification.Builder.with()
    .setName("ResourceApplication")
    .withRunnable()
      .add("echo1", new EchoServer(), ResourceSpecification.Builder.with()
        .setVirtualCores(1)
        .setMemory(256, ResourceSpecification.SizeUnit.MEGA)
        .setInstances(2).build()).noLocalFiles()
      .add("echo2", new EchoServer(), ResourceSpecification.Builder.with()
        .setVirtualCores(2)
        .setMemory(512, ResourceSpecification.SizeUnit.MEGA)
        .setInstances(1).build()).noLocalFiles()
    .anyOrder()
    .build();
}
 

@Override
public TwillSpecification configure() {
  // Make the runnable request for container smaller than 128MB (the allocation minimum)
  ResourceSpecification res = ResourceSpecification.Builder.with()
    .setVirtualCores(1)
    .setMemory(16, ResourceSpecification.SizeUnit.MEGA)
    .build();

  return TwillSpecification.Builder.with()
    .setName("MaxHeapApp")
    .withRunnable()
    .add("sleep", new MaxHeapRunnable(12345), res).noLocalFiles()
    .add("sleep2", new MaxHeapRunnable(23456), res).noLocalFiles()
    .anyOrder()
    .build();
}
 

@Override
public TwillSpecification configure() {
  ResourceSpecification largeRes = ResourceSpecification.Builder.with()
    .setVirtualCores(1)
    .setMemory(1024, ResourceSpecification.SizeUnit.MEGA)
    .build();

  ResourceSpecification smallRes = ResourceSpecification.Builder.with()
    .setVirtualCores(1)
    .setMemory(512, ResourceSpecification.SizeUnit.MEGA)
    .build();

  return TwillSpecification.Builder.with()
    .setName("SleepApp")
    .withRunnable()
      .add("sleep1", new SleepRunnable(12345), largeRes).noLocalFiles()
      .add("sleep2", new SleepRunnable(12346), smallRes).noLocalFiles()
    .withOrder()
      .begin("sleep1")
      .nextWhenStarted("sleep2")
    .build();
}
 

public DefaultRuntimeSpecification(String name,
                                   TwillRunnableSpecification runnableSpec,
                                   ResourceSpecification resourceSpec,
                                   Collection<LocalFile> localFiles) {
  this.name = name;
  this.runnableSpec = runnableSpec;
  this.resourceSpec = resourceSpec;
  this.localFiles = Collections.unmodifiableList(new ArrayList<LocalFile>(localFiles));
}
 

@Test
public void testBuilderWithLists() throws Exception {
  final ResourceSpecification actual = ResourceSpecification.Builder.with()
          .setVirtualCores(5)
          .setMemory(4, ResourceSpecification.SizeUnit.GIGA)
          .setInstances(3)
          .setUplink(10, ResourceSpecification.SizeUnit.GIGA)
          .setDownlink(5, ResourceSpecification.SizeUnit.GIGA)
          .build();
  final DefaultResourceSpecification expected =
          new DefaultResourceSpecification(5, 4096, 3, 10240, 5120);
  ReflectionAssert.assertLenientEquals(expected, actual);
}
 

@Test
public void testBuilder() throws Exception {
  final ResourceSpecification actual = ResourceSpecification.Builder.with()
          .setVirtualCores(5)
          .setMemory(4, ResourceSpecification.SizeUnit.GIGA)
          .setInstances(3)
          .setUplink(10, ResourceSpecification.SizeUnit.GIGA)
          .setDownlink(5, ResourceSpecification.SizeUnit.GIGA)
          .build();
  final DefaultResourceSpecification expected =
          new DefaultResourceSpecification(5, 4096, 3, 10240, 5120);
  ReflectionAssert.assertLenientEquals(expected, actual);
}
 

private Resource createCapability(ResourceSpecification resourceSpec) {
  Resource capability = Records.newRecord(Resource.class);

  if (!YarnUtils.setVirtualCores(capability, resourceSpec.getVirtualCores())) {
    LOG.debug("Virtual cores limit not supported.");
  }

  capability.setMemory(resourceSpec.getMemorySize());
  return capability;
}
 

@Override
public ResourceSpecification deserialize(JsonElement json, Type typeOfT,
                                         JsonDeserializationContext context) throws JsonParseException {
  JsonObject jsonObj = json.getAsJsonObject();
  final List<String> hosts = context.deserialize(jsonObj.getAsJsonArray("hosts"), List.class);
  final List<String> racks = context.deserialize(jsonObj.getAsJsonArray("racks"), List.class);
  return new DefaultResourceSpecification(jsonObj.get("cores").getAsInt(),
                                          jsonObj.get("memorySize").getAsInt(),
                                          jsonObj.get("instances").getAsInt(),
                                          jsonObj.get("uplink").getAsInt(),
                                          jsonObj.get("downlink").getAsInt());
}
 

@Test
public void testFailureRestart() throws Exception {
  TwillRunner runner = getTwillRunner();

  ResourceSpecification resource = ResourceSpecification.Builder.with()
    .setVirtualCores(1)
    .setMemory(512, ResourceSpecification.SizeUnit.MEGA)
    .setInstances(2)
    .build();
  TwillController controller = runner.prepare(new FailureRunnable(), resource)
    .withApplicationArguments("failure")
    .withArguments(FailureRunnable.class.getSimpleName(), "failure2")
    .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out, true)))
    .start();

  Iterable<Discoverable> discoverables = controller.discoverService("failure");
  Assert.assertTrue(waitForSize(discoverables, 2, 120));

  // Make sure we see the right instance IDs
  Assert.assertEquals(Sets.newHashSet(0, 1), getInstances(discoverables));

  // Kill server with instanceId = 0
  controller.sendCommand(FailureRunnable.class.getSimpleName(), Command.Builder.of("kill0").build());

  // Make sure the runnable is killed.
  Assert.assertTrue(waitForSize(discoverables, 1, 120));

  // Wait for the restart
  Assert.assertTrue(waitForSize(discoverables, 2, 120));

  // Make sure we see the right instance IDs
  Assert.assertEquals(Sets.newHashSet(0, 1), getInstances(discoverables));

  controller.terminate().get(120, TimeUnit.SECONDS);
}
 

@Override
public TwillSpecification configure() {
  return TwillSpecification.Builder.with()
    .setName("TimeoutApplication")
    .withRunnable()
    .add(new TimeoutRunnable(),
         ResourceSpecification.Builder.with()
           .setVirtualCores(1)
           .setMemory(8, ResourceSpecification.SizeUnit.GIGA).build())
    .noLocalFiles()
    .anyOrder()
    .withEventHandler(new Handler(parentFolderPath))
    .build();
}
 

@Override
public JsonElement serialize(ResourceSpecification src, Type typeOfSrc, JsonSerializationContext context) {
  JsonObject json = new JsonObject();
  json.addProperty("cores", src.getVirtualCores());
  json.addProperty("memorySize", src.getMemorySize());
  json.addProperty("instances", src.getInstances());
  json.addProperty("uplink", src.getUplink());
  json.addProperty("downlink", src.getDownlink());
  return json;
}
 

@Override
public RuntimeSpecification deserialize(JsonElement json, Type typeOfT,
                                        JsonDeserializationContext context) throws JsonParseException {
  JsonObject jsonObj = json.getAsJsonObject();

  String name = jsonObj.get("name").getAsString();
  TwillRunnableSpecification runnable = context.deserialize(jsonObj.get("runnable"),
                                                             TwillRunnableSpecification.class);
  ResourceSpecification resources = context.deserialize(jsonObj.get("resources"),
                                                        ResourceSpecification.class);
  Collection<LocalFile> files = context.deserialize(jsonObj.get("files"),
                                                    new TypeToken<Collection<LocalFile>>() { }.getType());

  return new DefaultRuntimeSpecification(name, runnable, resources, files);
}
 

@Override
public JsonElement serialize(RuntimeSpecification src, Type typeOfSrc, JsonSerializationContext context) {
  JsonObject json = new JsonObject();
  json.addProperty("name", src.getName());
  json.add("runnable", context.serialize(src.getRunnableSpecification(), TwillRunnableSpecification.class));
  json.add("resources", context.serialize(src.getResourceSpecification(), ResourceSpecification.class));
  json.add("files", context.serialize(src.getLocalFiles(), new TypeToken<Collection<LocalFile>>() { }.getType()));

  return json;
}
 

@Test
public void maxRetriesWithIncreasedInstances() throws InterruptedException, ExecutionException {
  TwillRunner runner = getTwillRunner();
  final int maxRetries = 3;
  final AtomicInteger retriesSeen = new AtomicInteger(0);
  final CountDownLatch retriesExhausted = new CountDownLatch(1);
  final CountDownLatch allRunning = new CountDownLatch(1);

  // start with 2 instances
  ResourceSpecification resource = ResourceSpecification.Builder.with().setVirtualCores(1)
    .setMemory(512, ResourceSpecification.SizeUnit.MEGA).setInstances(2).build();

  TwillController controller = runner.prepare(new FailingInstanceServer(), resource)
    .withMaxRetries(FailingInstanceServer.class.getSimpleName(), maxRetries)
    .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out, true))).addLogHandler(new LogHandler() {
      @Override
      public void onLog(LogEntry logEntry) {
        if (logEntry.getMessage().contains(" retries for instance")) {
          retriesSeen.incrementAndGet();
        }
        if (logEntry.getMessage().contains("Retries exhausted")) {
          retriesExhausted.countDown();
        }
        if (logEntry.getMessage().contains("fully provisioned with 2 instances")) {
          allRunning.countDown();
        }
      }
    }).start();

  try {
    // wait for initial instances to have started
    allRunning.await();

    /*
     * now increase the number of instances. these should fail since there instance ids are > 1. afterwards, the
     * number of retries should be 3 since only this one instance failed.
     */
    controller.changeInstances(FailingInstanceServer.class.getSimpleName(), 3);
    retriesExhausted.await();
    Assert.assertEquals(3, retriesSeen.get());

  } finally {
    controller.terminate().get();
  }
}
 

@Override
public ResourceSpecification getResourceSpecification() {
  return resourceSpec;
}
 

public SingleRunnableApplication(TwillRunnable runnable, ResourceSpecification resourceSpec) {
  this.runnable = runnable;
  this.resourceSpec = resourceSpec;
}
 

@Override
public TwillPreparer prepare(TwillRunnable runnable) {
  return prepare(runnable, ResourceSpecification.BASIC);
}
 

@Override
public TwillSpecification configure() {
  final String yarnDeploymentPolicyStr = yarnConfig.get(DEPLOYMENT_POLICY, DISTRIBUTED.name());

  TwillSpecification.PlacementPolicy.Type yarnDeploymentPolicy = DISTRIBUTED;
  try {
    yarnDeploymentPolicy =
      TwillSpecification.PlacementPolicy.Type.valueOf(yarnDeploymentPolicyStr.toUpperCase());
  } catch(IllegalArgumentException e) {
    logger.error("Invalid Deployment Policy is provided: {}, reverting to {}", yarnDeploymentPolicyStr, DISTRIBUTED);
  }


  URI kerberosKeytabURI = null;
  if (!Strings.isNullOrEmpty(keytabFileLocation)) {
    kerberosKeytabURI = new File(keytabFileLocation).toURI();
  }

  String cpu = yarnConfig.get(YARN_CPU);
  String memoryOnHeap = yarnConfig.get(YARN_MEMORY_ON_HEAP);
  String memoryOffHeap = yarnConfig.get(YARN_MEMORY_OFF_HEAP);
  String containerCount = yarnConfig.get(YARN_CONTAINER_COUNT);
  if (cpu == null) {
    throw new IllegalArgumentException("No cpu was specified in yarn config");
  }
  if (memoryOnHeap == null) {
    throw new IllegalArgumentException("No memory was specified in yarn config");
  }
  if (memoryOffHeap == null) {
    throw new IllegalArgumentException("No memory was specified in yarn config");
  }
  if (containerCount == null) {
    throw new IllegalArgumentException("No container count was specified in yarn config");
  }
  int cpuInt = Integer.valueOf(cpu);
  int memoryOnHeapInt = Integer.valueOf(memoryOnHeap);
  int memoryOffHeapInt = Integer.valueOf(memoryOffHeap);
  int containerCountInt = Integer.valueOf(containerCount);

  TwillSpecification.Builder.MoreFile files = TwillSpecification.Builder.with()
      .setName(yarnConfig.get(YARN_APP_NAME, YARN_APPLICATION_NAME_DEFAULT))
      .withRunnable()
      .add(YARN_RUNNABLE_NAME, new AppBundleRunnable(),
          ResourceSpecification.Builder.with()
              .setVirtualCores(cpuInt)
              .setMemory(memoryOnHeapInt+memoryOffHeapInt, ResourceSpecification.SizeUnit.MEGA)
              .setInstances(containerCountInt).build())
      .withLocalFiles()
      .add(YARN_BUNDLED_JAR_NAME, yarnBundledJarPath.toUri(), false);


  // Adding Dremio jars as resources
  if (kerberosKeytabURI != null) {
    files = files.add(KEYTAB_FILE_NAME, kerberosKeytabURI, false);
  }

  for (File classpathJarName : classpathJarNames) {
    files = files.add(classpathJarName.getName(), classpathJarName, false);
  }

  return files.apply().withPlacementPolicy()
      .add(yarnDeploymentPolicy, YARN_RUNNABLE_NAME)
      .anyOrder().build();
}
 

@Override
public TwillPreparer prepare(TwillRunnable runnable, ResourceSpecification resourceSpecification) {
  return prepare(new SingleRunnableApplication(runnable, resourceSpecification));
}
 

@Test
public void testResourceReportWithFailingContainers() throws InterruptedException, IOException,
  TimeoutException, ExecutionException {
  TwillRunner runner = getTwillRunner();

  ResourceSpecification resourceSpec = ResourceSpecification.Builder.with()
    .setVirtualCores(1)
    .setMemory(256, ResourceSpecification.SizeUnit.MEGA)
    .setInstances(2)
    .build();
  TwillController controller = runner.prepare(new BuggyServer(), resourceSpec)
    .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out, true)))
    .withApplicationArguments("echo")
    .withArguments("BuggyServer", "echo2")
    .start();

  final CountDownLatch running = new CountDownLatch(1);
  controller.onRunning(new Runnable() {
    @Override
    public void run() {
      running.countDown();
    }
  }, Threads.SAME_THREAD_EXECUTOR);

  Assert.assertTrue(running.await(120, TimeUnit.SECONDS));

  Iterable<Discoverable> echoServices = controller.discoverService("echo");
  Assert.assertTrue(waitForSize(echoServices, 2, 120));
  // check that we have 2 runnables.
  ResourceReport report = getResourceReport(controller, 10000);
  Assert.assertEquals(2, report.getRunnableResources("BuggyServer").size());

  // cause a divide by 0 in one server
  Discoverable discoverable = echoServices.iterator().next();
  try (
    Socket socket = new Socket(discoverable.getSocketAddress().getAddress(),
                               discoverable.getSocketAddress().getPort())
  ) {
    PrintWriter writer = new PrintWriter(new OutputStreamWriter(socket.getOutputStream(), Charsets.UTF_8), true);
    writer.println("0");
  }

  // takes some time for app master to find out the container completed...
  int count = 0;
  while (count < 100) {
    report = getResourceReport(controller, 10000);
    // check that we have 1 runnable, not 2.
    if (report.getRunnableResources("BuggyServer").size() == 1) {
      break;
    }
    LOG.info("Wait for BuggyServer to have 1 instance left. Trial {}.", count);
    count++;
    TimeUnit.SECONDS.sleep(1);
  }
  Assert.assertTrue("Still has 2 contains running after 100 seconds", count < 100);

  controller.terminate().get(100, TimeUnit.SECONDS);
  // Sleep a bit before exiting.
  TimeUnit.SECONDS.sleep(2);
}
 

@Test
public void testRunnablesGetAllowedResourcesInEnv() throws InterruptedException, IOException,
  TimeoutException, ExecutionException {
  TwillRunner runner = getTwillRunner();

  ResourceSpecification resourceSpec = ResourceSpecification.Builder.with()
    .setVirtualCores(1)
    .setMemory(2048, ResourceSpecification.SizeUnit.MEGA)
    .setInstances(1)
    .build();
  TwillController controller = runner.prepare(new EnvironmentEchoServer(), resourceSpec)
    .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out, true)))
    .withApplicationArguments("envecho")
    .withArguments("EnvironmentEchoServer", "echo2")
    .start();

  final CountDownLatch running = new CountDownLatch(1);
  controller.onRunning(new Runnable() {
    @Override
    public void run() {
      running.countDown();
    }
  }, Threads.SAME_THREAD_EXECUTOR);

  Assert.assertTrue(running.await(120, TimeUnit.SECONDS));

  Iterable<Discoverable> envEchoServices = controller.discoverService("envecho");
  Assert.assertTrue(waitForSize(envEchoServices, 1, 120));

  // TODO: check virtual cores once yarn adds the ability
  Map<String, String> expectedValues = Maps.newHashMap();
  expectedValues.put(EnvKeys.YARN_CONTAINER_MEMORY_MB, "2048");
  expectedValues.put(EnvKeys.TWILL_INSTANCE_COUNT, "1");

  // check environment of the runnable.
  Discoverable discoverable = envEchoServices.iterator().next();
  for (Map.Entry<String, String> expected : expectedValues.entrySet()) {
    try (
      Socket socket = new Socket(discoverable.getSocketAddress().getHostName(),
                                 discoverable.getSocketAddress().getPort())
    ) {
      PrintWriter writer = new PrintWriter(new OutputStreamWriter(socket.getOutputStream(), Charsets.UTF_8), true);
      LineReader reader = new LineReader(new InputStreamReader(socket.getInputStream(), Charsets.UTF_8));
      writer.println(expected.getKey());
      Assert.assertEquals(expected.getValue(), reader.readLine());
    }
  }

  controller.terminate().get(120, TimeUnit.SECONDS);
  // Sleep a bit before exiting.
  TimeUnit.SECONDS.sleep(2);
}
 

/**
 * Verify the cluster configuration (number and capability of node managers) required for the tests.
 */
@BeforeClass
public static void verifyClusterCapability() throws InterruptedException {
  // Ignore verifications if it is running against older Hadoop versions which does not support blacklists.
  Assume.assumeTrue(YarnUtils.getHadoopVersion().equals(YarnUtils.HadoopVersions.HADOOP_22));

  // All runnables in this test class use same resource specification for the sake of convenience.
  resource = ResourceSpecification.Builder.with()
    .setVirtualCores(RUNNABLE_CORES)
    .setMemory(RUNNABLE_MEMORY, ResourceSpecification.SizeUnit.MEGA)
    .build();
  twoInstancesResource = ResourceSpecification.Builder.with()
    .setVirtualCores(RUNNABLE_CORES)
    .setMemory(RUNNABLE_MEMORY, ResourceSpecification.SizeUnit.MEGA)
    .setInstances(2)
    .build();

  // The tests need exactly three NodeManagers in the cluster.
  int trials = 0;
  while (trials++ < 20) {
    try {
      nodeReports = TWILL_TESTER.getNodeReports();
      if (nodeReports != null && nodeReports.size() == 3) {
        break;
      }
    } catch (Exception e) {
      LOG.error("Failed to get node reports", e);
    }
    LOG.warn("NodeManagers != 3. {}", nodeReports);
    TimeUnit.SECONDS.sleep(1);
  }

  // All NodeManagers should have enough capacity available to accommodate at least two runnables.
  for (NodeReport nodeReport : nodeReports) {
    Resource capability = nodeReport.getCapability();
    Resource used = nodeReport.getUsed();
    Assert.assertNotNull(capability);
    if (used != null) {
      Assert.assertTrue(2 * resource.getMemorySize() < capability.getMemory() - used.getMemory());
    } else {
      Assert.assertTrue(2 * resource.getMemorySize() < capability.getMemory());
    }
  }
}
 

@Test
public void testEchoServer() throws Exception {
  TwillRunner runner = getTwillRunner();

  TwillController controller = runner.prepare(new EchoServer(),
                                              ResourceSpecification.Builder.with()
                                                       .setVirtualCores(1)
                                                       .setMemory(1, ResourceSpecification.SizeUnit.GIGA)
                                                       .setInstances(2)
                                                       .build())
                                      .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out, true)))
                                      .withApplicationArguments("echo")
                                      .withArguments("EchoServer", "echo2")
                                      .start();

  final CountDownLatch running = new CountDownLatch(1);
  controller.onRunning(new Runnable() {
    @Override
    public void run() {
      running.countDown();
    }
  }, Threads.SAME_THREAD_EXECUTOR);

  Assert.assertTrue(running.await(120, TimeUnit.SECONDS));

  Iterable<Discoverable> echoServices = controller.discoverService("echo");
  Assert.assertTrue(waitForSize(echoServices, 2, 120));

  for (Discoverable discoverable : echoServices) {
    String msg = "Hello: " + discoverable.getSocketAddress();

    try (
      Socket socket = new Socket(discoverable.getSocketAddress().getAddress(),
                                 discoverable.getSocketAddress().getPort())
    ) {
      PrintWriter writer = new PrintWriter(new OutputStreamWriter(socket.getOutputStream(), Charsets.UTF_8), true);
      LineReader reader = new LineReader(new InputStreamReader(socket.getInputStream(), Charsets.UTF_8));

      writer.println(msg);
      Assert.assertEquals(msg, reader.readLine());
    }
  }

  // Increase number of instances
  controller.changeInstances("EchoServer", 3).get(60, TimeUnit.SECONDS);
  Assert.assertTrue(waitForSize(echoServices, 3, 120));

  echoServices = controller.discoverService("echo2");

  // Decrease number of instances
  controller.changeInstances("EchoServer", 1).get(60, TimeUnit.SECONDS);
  Assert.assertTrue(waitForSize(echoServices, 1, 120));

  // Increase number of instances again
  controller.changeInstances("EchoServer", 2).get(60, TimeUnit.SECONDS);
  Assert.assertTrue(waitForSize(echoServices, 2, 120));

  // Test restart on instances for runnable
  Map<Integer, String> instanceIdToContainerId = Maps.newHashMap();
  ResourceReport report = waitForAfterRestartResourceReport(controller, "EchoServer", 15L,
                                                            TimeUnit.MINUTES, 2, null);
  Assert.assertTrue(report != null);
  Collection<TwillRunResources> runResources = report.getRunnableResources("EchoServer");
  for (TwillRunResources twillRunResources : runResources) {
    instanceIdToContainerId.put(twillRunResources.getInstanceId(), twillRunResources.getContainerId());
  }

  controller.restartAllInstances("EchoServer").get(60, TimeUnit.SECONDS);
  Assert.assertTrue(waitForSize(echoServices, 2, 120));

  report = waitForAfterRestartResourceReport(controller, "EchoServer", 15L, TimeUnit.MINUTES, 2,
                                             instanceIdToContainerId);
  Assert.assertTrue(report != null);

  // Make sure still only one app is running
  Iterable<TwillRunner.LiveInfo> apps = runner.lookupLive();
  Assert.assertTrue(waitForSize(apps, 1, 120));

  // Creates a new runner service to check it can regain control over running app.
  TwillRunnerService runnerService = TWILL_TESTER.createTwillRunnerService();
  runnerService.start();

  try {
    Iterable <TwillController> controllers = runnerService.lookup("EchoServer");
    Assert.assertTrue(waitForSize(controllers, 1, 120));

    for (TwillController c : controllers) {
      LOG.info("Stopping application: " + c.getRunId());
      c.terminate().get(30, TimeUnit.SECONDS);
    }

    Assert.assertTrue(waitForSize(apps, 0, 120));
  } finally {
    runnerService.stop();
  }

  // Sleep a bit before exiting.
  TimeUnit.SECONDS.sleep(2);
}