org.apache.twill.common.Threads Java Examples

@Ignore
@Test
public void testDistributedShell() throws InterruptedException {
  TwillRunner twillRunner = getTwillRunner();

  TwillController controller = twillRunner.prepare(new DistributedShell("pwd", "ls -al"))
                                          .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out)))
                                          .start();

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

  Assert.assertTrue(stopLatch.await(10, TimeUnit.SECONDS));
}
 

/**
 * Returns a {@link Runnable} that can be used as a {@link ListenableFuture} listener to trigger
 * further service action or completing the result future. Used by
 * {@link #doChain(boolean, com.google.common.util.concurrent.Service, com.google.common.util.concurrent.Service...)}
 */
private static Runnable createChainListener(final ListenableFuture<Service.State> future, final Service[] services,
                                            final AtomicInteger idx,
                                            final List<ListenableFuture<Service.State>> result,
                                            final SettableFuture<List<ListenableFuture<Service.State>>> resultFuture,
                                            final boolean doStart) {
  return new Runnable() {

    @Override
    public void run() {
      result.add(future);
      int nextIdx = idx.getAndIncrement();
      if (nextIdx == services.length) {
        resultFuture.set(result);
        return;
      }
      ListenableFuture<Service.State> actionFuture = doStart ? services[nextIdx].start() : services[nextIdx].stop();
      actionFuture.addListener(createChainListener(actionFuture, services, idx, result, resultFuture, doStart),
                               Threads.SAME_THREAD_EXECUTOR);
    }
  };
}
 

@Override
protected void doStart() {
  LOG.info("Start leader election on {}{} with guid {}", zkClient.getConnectString(), zkFolderPath, guid);

  executor = Executors.newSingleThreadExecutor(
    Threads.createDaemonThreadFactory("leader-election" + zkFolderPath.replace('/', '-')));

  executor.execute(new Runnable() {
    @Override
    public void run() {
      register();
      watcherCancellable = zkClient.addConnectionWatcher(wrapWatcher(new ConnectionWatcher()));
    }
  });
  notifyStarted();
}
 

@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
protected void doStart() throws Exception {
  LOG.info("Start application master with spec: {}",
    TwillRuntimeSpecificationAdapter.create().toJson(twillRuntimeSpec));

  // initialize the event handler, if it fails, it will fail the application.
  eventHandler.initialize(new BasicEventHandlerContext(twillRuntimeSpec));
  // call event handler started.
  eventHandler.started();

  instanceChangeExecutor = Executors.newSingleThreadExecutor(Threads.createDaemonThreadFactory("instanceChanger"));

  // Creates ZK path for runnable. It's ok if the path already exists.
  // That's for the case when the AM get killed and restarted
  ZKOperations.ignoreError(
    zkClient.create("/" + runId.getId() + "/runnables", null, CreateMode.PERSISTENT),
    KeeperException.NodeExistsException.class, null)
    .get();
  runningContainers.addWatcher(Constants.DISCOVERY_PATH_PREFIX);
  runnableContainerRequests = initContainerRequests();
}
 

@Override
public OperationFuture<String> create(final String path, @Nullable final byte[] data, final CreateMode createMode,
                                      final boolean createParent, final Iterable<ACL> acl) {
  // No retry for any SEQUENTIAL node, as some algorithms depends on only one sequential node being created.
  if (createMode == CreateMode.PERSISTENT_SEQUENTIAL || createMode == CreateMode.EPHEMERAL_SEQUENTIAL) {
    return super.create(path, data, createMode, createParent, acl);
  }

  final SettableOperationFuture<String> result = SettableOperationFuture.create(path, Threads.SAME_THREAD_EXECUTOR);
  Futures.addCallback(super.create(path, data, createMode, createParent, acl),
                      new OperationFutureCallback<String>(OperationType.CREATE, System.currentTimeMillis(),
                                                          path, result, new Supplier<OperationFuture<String>>() {
                        @Override
                        public OperationFuture<String> get() {
                          return FailureRetryZKClient.super.create(path, data, createMode, createParent, acl);
                        }
                      }));
  return result;
}
 

/**
 * Handles {@link SystemMessages#STOP_COMMAND} if the given message is a stop command. After this service is stopped,
 * the message node will be removed.
 *
 * @param message Message to process
 * @param messageRemover Runnable to remove the message node when this service is stopped
 * @return {@code true} if the given message is a stop command, {@code false} otherwise
 */
private boolean handleStopMessage(Message message, final Runnable messageRemover) {
  if (message.getType() != Message.Type.SYSTEM || !SystemMessages.STOP_COMMAND.equals(message.getCommand())) {
    return false;
  }

  // Stop this service.
  Futures.addCallback(stop(), new FutureCallback<State>() {
    @Override
    public void onSuccess(State result) {
      messageRemover.run();
    }

    @Override
    public void onFailure(Throwable t) {
      LOG.error("Stop service failed upon STOP command", t);
      messageRemover.run();
    }
  }, Threads.SAME_THREAD_EXECUTOR);
  return true;
}
 

@SuppressWarnings("unchecked")
@Override
protected void doStart() throws Exception {
  for (Map.Entry<String, String> entry : containerLiveNodeData.getLogLevels().entrySet()) {
    String loggerName = entry.getKey();
    String oldLogLevel = setLogLevel(loggerName, entry.getValue());
    if (!defaultLogLevels.containsKey(loggerName)) {
      oldLogLevels.put(loggerName, oldLogLevel);
    }
  }

  commandExecutor = Executors.newSingleThreadExecutor(
    Threads.createDaemonThreadFactory("runnable-command-executor"));

  Class<?> runnableClass = classLoader.loadClass(specification.getClassName());
  Preconditions.checkArgument(TwillRunnable.class.isAssignableFrom(runnableClass),
                              "Class %s is not instance of TwillRunnable.", specification.getClassName());

  runnable = Instances.newInstance((Class<TwillRunnable>) runnableClass);
  runnable.initialize(context);
}
 

@Override
protected void doStart() {
  // A single thread executor for all events
  ThreadPoolExecutor executor = new ThreadPoolExecutor(1, 1, 0L, TimeUnit.MILLISECONDS,
                                                       new LinkedBlockingQueue<Runnable>(),
                                                       Threads.createDaemonThreadFactory("zk-client-EventThread"));
  // Just discard the execution if the executor is closed
  executor.setRejectedExecutionHandler(new ThreadPoolExecutor.DiscardPolicy());
  eventExecutor = executor;

  try {
    zooKeeper.set(createZooKeeper());
  } catch (IOException e) {
    notifyFailed(e);
  }
}
 

@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());
}
 

@Test
public void testDebugPortAllRunnables() throws Exception {
  YarnTwillRunnerService runner = getTwillRunner();
  runner.start();

  TwillController controller = runner.prepare(new DummyApplication())
                                     .enableDebugging()
                                     .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out)))
                                     .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));
  Assert.assertTrue(waitForDebugPort(controller, "r1", 30));
  Assert.assertTrue(waitForDebugPort(controller, "r2", 30));
  controller.terminate().get(120, TimeUnit.SECONDS);
  // Sleep a bit before exiting.
  TimeUnit.SECONDS.sleep(2);
}
 

protected final void watchInstanceNode() {
  if (!shouldProcessZKEvent()) {
    return;
  }
  Futures.addCallback(zkClient.getData(getInstancePath(), new Watcher() {
    @Override
    public void process(WatchedEvent event) {
      if (!shouldProcessZKEvent()) {
        return;
      }
      switch (event.getType()) {
        case NodeDataChanged:
          watchInstanceNode();
          break;
        case NodeDeleted:
          instanceNodeFailed(KeeperException.create(KeeperException.Code.NONODE, getInstancePath()));
          break;
        default:
          LOG.info("Ignore ZK event for instance node: {}", event);
      }
    }
  }), instanceNodeDataCallback, Threads.SAME_THREAD_EXECUTOR);
}
 

private YarnTwillController listenController(final YarnTwillController controller) {
  controller.onTerminated(new Runnable() {
    @Override
    public void run() {
      synchronized (YarnTwillRunnerService.this) {
        Iterables.removeIf(controllers.values(), new Predicate<YarnTwillController>() {
          @Override
          public boolean apply(YarnTwillController input) {
            return input == controller;
          }
        });
      }
    }
  }, Threads.SAME_THREAD_EXECUTOR);
  return controller;
}
 

/**
 * Start the publisher. This method must be called before other methods. This method is only to be called
 * by KafkaClientService who own this object.
 * @return A Cancellable for closing this publish.
 */
Cancellable start() {
  ExecutorService listenerExecutor
    = Executors.newSingleThreadExecutor(Threads.createDaemonThreadFactory("kafka-publisher"));

  // Listen to changes in broker list
  final BrokerListChangeListener listener = new BrokerListChangeListener(listenerCancelled, producer,
                                                                         ack, compression);
  Cancellable cancelChangeListener = brokerService.addChangeListener(listener, listenerExecutor);

  // Invoke the change listener at least once. Since every call to the listener is through the single thread
  // executor, there is no race and for sure the listener always see the latest change, either through this call
  // or from the BrokerService callback.
  Future<?> completion = listenerExecutor.submit(new Runnable() {
    @Override
    public void run() {
      listener.changed(brokerService);
    }
  });

  Futures.getUnchecked(completion);
  return new ProducerCancellable(listenerExecutor, listenerCancelled, cancelChangeListener, producer);
}
 

@Test
public void testDebugPortOneRunnable() throws Exception {
  YarnTwillRunnerService runner = getTwillRunner();
  runner.start();

  TwillController controller = runner.prepare(new DummyApplication())
                                     .enableDebugging("r1")
                                     .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out)))
                                     .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));
  Assert.assertTrue(waitForDebugPort(controller, "r1", 30));
  controller.terminate().get(120, TimeUnit.SECONDS);
  // Sleep a bit before exiting.
  TimeUnit.SECONDS.sleep(2);
}
 

@Override
protected void startUp() {
  scheduler = Executors.newSingleThreadScheduledExecutor(Threads.createDaemonThreadFactory("location-cache-cleanup"));
  scheduler.execute(new Runnable() {
    @Override
    public void run() {
      long currentTime = System.currentTimeMillis();
      cleanup(currentTime);

      // By default, run the cleanup at half of the expiry
      long scheduleDelay = expiry / 2;
      for (PendingCleanup pendingCleanup : pendingCleanups) {
        // If there is any pending cleanup that needs to be cleanup early, schedule the run earlier.
        if (pendingCleanup.getExpireTime() - currentTime < scheduleDelay) {
          scheduleDelay = pendingCleanup.getExpireTime() - currentTime;
        }
      }
      scheduler.schedule(this, scheduleDelay, TimeUnit.MILLISECONDS);
    }
  });
}
 

/**
 * Constructs ZKDiscoveryService using the provided zookeeper client for storing service registry under namespace.
 * @param zkClient of zookeeper quorum
 * @param namespace under which the service registered would be stored in zookeeper.
 *                  If namespace is {@code null}, no namespace will be used.
 */
public ZKDiscoveryService(ZKClient zkClient, String namespace) {
  this.closed = new AtomicBoolean();
  this.discoverables = HashMultimap.create();
  this.lock = new ReentrantLock();
  this.retryExecutor = Executors.newSingleThreadScheduledExecutor(
    Threads.createDaemonThreadFactory("zk-discovery-retry"));
  this.zkClient = namespace == null ? zkClient : ZKClients.namespace(zkClient, namespace);
  this.services = CacheBuilder.newBuilder()
    .removalListener(new RemovalListener<String, ServiceDiscoveredCacheEntry>() {
      @Override
      public void onRemoval(RemovalNotification<String, ServiceDiscoveredCacheEntry> notification) {
        ServiceDiscoveredCacheEntry entry = notification.getValue();
        if (entry != null) {
          entry.cancel();
        }
      }
    })
    .build(createServiceLoader());
  this.watcherCancellable = this.zkClient.addConnectionWatcher(createConnectionWatcher());
}
 

@Override
protected void doStart() {
  // A single thread executor for all events
  ThreadPoolExecutor executor = new ThreadPoolExecutor(1, 1, 0L, TimeUnit.MILLISECONDS,
                                                       new LinkedBlockingQueue<Runnable>(),
                                                       Threads.createDaemonThreadFactory("zk-client-EventThread"));
  // Just discard the execution if the executor is closed
  executor.setRejectedExecutionHandler(new ThreadPoolExecutor.DiscardPolicy());
  eventExecutor = executor;

  try {
    zooKeeper.set(createZooKeeper());
  } catch (IOException e) {
    notifyFailed(e);
  }
}
 

/**
 * Sends a {@link Message} to the remote service. Returns a future that will be completed when the message
 * has been processed.
 * @param message The message to send.
 * @param result Object to set into the future when message is being processed.
 * @param <V> Type of the result.
 * @return A {@link ListenableFuture} that will be completed when the message has been processed.
 */
protected final synchronized <V> ListenableFuture<V> sendMessage(Message message, V result) {
  if (!isRunning()) {
    return Futures.immediateFailedFuture(new IllegalStateException("Cannot send message to non-running application"));
  }
  final ListenableFuture<V> messageFuture = ZKMessages.sendMessage(zkClient, getMessagePrefix(), message, result);
  messageFutures.add(messageFuture);
  messageFuture.addListener(new Runnable() {
    @Override
    public void run() {
      // If the completion is triggered when stopping, do nothing.
      if (state() == State.STOPPING) {
        return;
      }
      synchronized (AbstractZKServiceController.this) {
        messageFutures.remove(messageFuture);
      }
    }
  }, Threads.SAME_THREAD_EXECUTOR);

  return messageFuture;
}
 

/**
 * Performs the actual logic of chain Service start/stop.
 */
private static ListenableFuture<List<ListenableFuture<Service.State>>> doChain(boolean doStart,
                                                                               Service firstService,
                                                                               Service...moreServices) {
  SettableFuture<List<ListenableFuture<Service.State>>> resultFuture = SettableFuture.create();
  List<ListenableFuture<Service.State>> result = Lists.newArrayListWithCapacity(moreServices.length + 1);

  ListenableFuture<Service.State> future = doStart ? firstService.start() : firstService.stop();
  future.addListener(createChainListener(future, moreServices, new AtomicInteger(0), result, resultFuture, doStart),
                     Threads.SAME_THREAD_EXECUTOR);
  return resultFuture;
}
 

@Test
public void testControllerListener() throws InterruptedException {
  InMemoryZKServer zkServer = InMemoryZKServer.builder().build();
  zkServer.startAndWait();

  LOG.info("ZKServer: " + zkServer.getConnectionStr());
  try {
    RunId runId = RunIds.generate();
    ZKClientService zkClientService = ZKClientService.Builder.of(zkServer.getConnectionStr()).build();
    zkClientService.startAndWait();

    Service service = createService(zkClientService, runId);
    service.startAndWait();

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

    Assert.assertTrue(runLatch.await(2, TimeUnit.SECONDS));

    service.stopAndWait();

    zkClientService.stopAndWait();
  } finally {
    zkServer.stopAndWait();
  }
}
 

private void actOnExists(final String path, final Runnable action) {
  // Watch for node existence.
  final AtomicBoolean nodeExists = new AtomicBoolean(false);
  Futures.addCallback(zkClient.exists(path, new Watcher() {
    @Override
    public void process(WatchedEvent event) {
      if (!shouldProcessZKEvent()) {
        return;
      }
      // When node is created, call the action.
      // Other event type would be handled by the action.
      if (event.getType() == Event.EventType.NodeCreated && nodeExists.compareAndSet(false, true)) {
        action.run();
      }
    }
  }), new FutureCallback<Stat>() {
    @Override
    public void onSuccess(Stat result) {
      if (result != null && nodeExists.compareAndSet(false, true)) {
        action.run();
      }
    }

    @Override
    public void onFailure(Throwable t) {
      LOG.error("Failed in exists call to {}. Shutting down service.", path, t);
      forceShutDown();
    }
  }, Threads.SAME_THREAD_EXECUTOR);
}
 

@Override
protected void startUp() throws Exception {
  scheduler = Executors.newSingleThreadScheduledExecutor(Threads.createDaemonThreadFactory("kafka-client-cleanup"));
  scheduler.scheduleAtFixedRate(this, PUBLISHER_CLEANUP_SECONDS, PUBLISHER_CLEANUP_SECONDS, TimeUnit.SECONDS);

  // Start broker service to get auto-updated brokers information.
  brokerService.startAndWait();
}
 

@Test
public void testController() throws ExecutionException, InterruptedException, TimeoutException {
  InMemoryZKServer zkServer = InMemoryZKServer.builder().build();
  zkServer.startAndWait();

  LOG.info("ZKServer: " + zkServer.getConnectionStr());

  try {
    RunId runId = RunIds.generate();
    ZKClientService zkClientService = ZKClientService.Builder.of(zkServer.getConnectionStr()).build();
    zkClientService.startAndWait();

    Service service = createService(zkClientService, runId);
    service.startAndWait();

    TwillController controller = getController(zkClientService, "testController", runId);
    controller.sendCommand(Command.Builder.of("test").build()).get(2, TimeUnit.SECONDS);
    controller.terminate().get(2, TimeUnit.SECONDS);

    final CountDownLatch terminateLatch = new CountDownLatch(1);
    service.addListener(new ServiceListenerAdapter() {
      @Override
      public void terminated(Service.State from) {
        terminateLatch.countDown();
      }
    }, Threads.SAME_THREAD_EXECUTOR);

    Assert.assertTrue(service.state() == Service.State.TERMINATED || terminateLatch.await(2, TimeUnit.SECONDS));

    zkClientService.stopAndWait();

  } finally {
    zkServer.stopAndWait();
  }
}
 

@Test
public void testControllerBefore() throws InterruptedException, ExecutionException, TimeoutException {
  InMemoryZKServer zkServer = InMemoryZKServer.builder().build();
  zkServer.startAndWait();

  LOG.info("ZKServer: " + zkServer.getConnectionStr());
  try {
    RunId runId = RunIds.generate();
    ZKClientService zkClientService = ZKClientService.Builder.of(zkServer.getConnectionStr()).build();
    zkClientService.startAndWait();

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

    Service service = createService(zkClientService, runId);
    service.start();

    Assert.assertTrue(runLatch.await(2, TimeUnit.SECONDS));

    try {
      controller.awaitTerminated(2, TimeUnit.SECONDS);
      Assert.fail("Service should not be terminated");
    } catch (TimeoutException e) {
      // Expected
    }

    service.stop();
    controller.awaitTerminated(120, TimeUnit.SECONDS);

  } finally {
    zkServer.stopAndWait();
  }
}
 

private ClusterEnriched startClusterAsync(Cluster cluster) throws YarnProvisioningHandlingException {
  YarnConfiguration yarnConfiguration = new YarnConfiguration();
  updateYarnConfiguration(cluster, yarnConfiguration);

  List<Property> props = cluster.getClusterConfig().getSubPropertyList();
  // only to show those props on UI/API
  defaultsConfigurator.getDistroTypeDefaultsConfigurator(
    cluster.getClusterConfig().getDistroType(), cluster.getClusterConfig().getIsSecure()).mergeProperties(props);
  List<Property> cleansedProperties = new ArrayList<>();
  for (Property prop : props) {
    if (!EXCLUDED.contains(prop.getKey())) {
      cleansedProperties.add(prop);
    }
  }

  // async call - unfortunately I can not add event handlers before start of TwillController
  // which means we can miss failures
  TwillController twillController = yarnController.startCluster(yarnConfiguration, cleansedProperties);

  String runId = twillController.getRunId().getId();
  RunId dRunId = new RunId(runId);
  cluster.setState(ClusterState.STARTING);
  cluster.setRunId(dRunId);

  OnRunningRunnable onRunning = new OnRunningRunnable(cluster);

  twillController.onRunning(onRunning, Threads.SAME_THREAD_EXECUTOR);

  initOnTerminatingThread(cluster, twillController);

  return getClusterInfo(cluster);
}
 

/**
 * Logs if the given future failed.
 */
private <T> void logIfFailed(ListenableFuture<T> future) {
  Futures.addCallback(future, new FutureCallback<T>() {
    @Override
    public void onSuccess(T result) {
      // All-good
    }

    @Override
    public void onFailure(Throwable t) {
      LOG.error("Operation failed for service {} with runId {}", getServiceName(), runId, t);
    }
  }, Threads.SAME_THREAD_EXECUTOR);
}
 

@Test
public void testAppSessionExpire() throws InterruptedException, ExecutionException, TimeoutException {
  TwillRunner runner = getTwillRunner();
  TwillController controller = runner.prepare(new SleepRunnable(600))
                                     .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);

  // Wait for application running
  Assert.assertTrue(runLatch.await(60, TimeUnit.SECONDS));

  // Find the app master ZK session and expire it two times, 10 seconds apart.
  for (int i = 0; i < 2; i++) {
    Assert.assertTrue(expireAppMasterZKSession(controller, 10, TimeUnit.SECONDS));

    try {
      controller.awaitTerminated(10, TimeUnit.SECONDS);
      Assert.fail("Unexpected application termination.");
    } catch (TimeoutException e) {
      // OK, expected.
    }
  }

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

@Test
public void testDisableLogCollection() throws Exception {
  final AtomicBoolean logReceived = new AtomicBoolean();

  // Start the LogRunnable by turning off log collection
  TwillRunner runner = getTwillRunner();
  TwillController controller = runner.prepare(new LogRunnable())
    .withConfiguration(Collections.singletonMap(Configs.Keys.LOG_COLLECTION_ENABLED, "false"))
    .addLogHandler(new LogHandler() {
      @Override
      public void onLog(LogEntry logEntry) {
        logReceived.set(true);
      }
    })
    .start();

  // Make sure the runnable gets executed
  try {
    final CountDownLatch latch = new CountDownLatch(1);
    controller.discoverService("log").watchChanges(new ServiceDiscovered.ChangeListener() {
      @Override
      public void onChange(ServiceDiscovered serviceDiscovered) {
        if (Iterables.size(serviceDiscovered) == 1) {
          latch.countDown();
        }
      }
    }, Threads.SAME_THREAD_EXECUTOR);
    Assert.assertTrue(latch.await(120, TimeUnit.SECONDS));
  } finally {
    controller.terminate().get(120, TimeUnit.SECONDS);
  }

  // Should receive no log
  Assert.assertFalse("Not expecting logs collected", logReceived.get());
}
 

@Override
protected final void startUp() throws Exception {
  // Single thread executor that will discard task silently if it is already terminated, which only
  // happens when this service is shutting down.
  messageCallbackExecutor = new ThreadPoolExecutor(1, 1,
                                                   0L, TimeUnit.MILLISECONDS,
                                                   new LinkedBlockingQueue<Runnable>(),
                                                   Threads.createDaemonThreadFactory("message-callback"),
                                                   new ThreadPoolExecutor.DiscardPolicy());

  // Watch for session expiration, recreate the live node if reconnected after expiration.
  watcherCancellable = zkClient.addConnectionWatcher(new Watcher() {
    private boolean expired = false;

    @Override
    public void process(WatchedEvent event) {
      if (event.getState() == Event.KeeperState.Expired) {
        LOG.warn("ZK Session expired for service {} with runId {}.", getServiceName(), runId.getId());
        expired = true;
      } else if (event.getState() == Event.KeeperState.SyncConnected && expired) {
        LOG.info("Reconnected after expiration for service {} with runId {}", getServiceName(), runId.getId());
        expired = false;
        logIfFailed(createLiveNode());
      }
    }
  });

  // Create the live node, if succeeded, start the service, otherwise fail out.
  createLiveNode().get();

  // Create node for messaging
  ZKOperations.ignoreError(zkClient.create(getZKPath("messages"), null, CreateMode.PERSISTENT),
                           KeeperException.NodeExistsException.class, null).get();

  doStart();

  // Starts watching for messages
  watchMessages();
}