org.apache.twill.zookeeper.ZKClients Java Examples

@Provides
@Singleton
private ZKClientService provideZKClientService(Configuration conf) {
  String zkStr = conf.get(TxConstants.Service.CFG_DATA_TX_ZOOKEEPER_QUORUM);
  if (zkStr == null) {
    // Default to HBase one.
    zkStr = conf.get(TxConstants.HBase.ZOOKEEPER_QUORUM);
  }

  int timeOut = conf.getInt(TxConstants.HBase.ZK_SESSION_TIMEOUT, TxConstants.HBase.DEFAULT_ZK_SESSION_TIMEOUT);
  ZKClientService zkClientService = new TephraZKClientService(zkStr, timeOut, null,
                                                              ArrayListMultimap.<String, byte[]>create());
  return ZKClientServices.delegate(
    ZKClients.reWatchOnExpire(
      ZKClients.retryOnFailure(zkClientService, RetryStrategies.exponentialDelay(500, 2000, 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());
}
 

public AbstractTwillController(String appName, RunId runId, ZKClient zkClient, boolean logCollectionEnabled,
                               Iterable<LogHandler> logHandlers) {
  super(runId, zkClient);
  this.appName = appName;
  this.runId = runId;
  this.logHandlers = new ConcurrentLinkedQueue<>();

  // When addressing TWILL-147, need to check if the given ZKClient is
  // actually used by the Kafka used for log collection
  if (logCollectionEnabled) {
    this.kafkaClient = new ZKKafkaClientService(ZKClients.namespace(zkClient, "/" + runId.getId() + "/kafka"));
    Iterables.addAll(this.logHandlers, logHandlers);
  } else {
    this.kafkaClient = null;
    if (!Iterables.isEmpty(logHandlers)) {
      LOG.warn("Log collection is disabled for application {} with runId {}. " +
                 "Adding log handler won't get any logs.", appName, runId);
    }
  }
}
 

@Override
public PhoenixTransactionService getTransactionService(Configuration config, ConnectionInfo connInfo, int port) {
    config.setInt(TxConstants.Service.CFG_DATA_TX_BIND_PORT, port);
    int retryTimeOut = config.getInt(TxConstants.Service.CFG_DATA_TX_CLIENT_DISCOVERY_TIMEOUT_SEC, 
            TxConstants.Service.DEFAULT_DATA_TX_CLIENT_DISCOVERY_TIMEOUT_SEC);
    ZKClientService zkClient = ZKClientServices.delegate(
      ZKClients.reWatchOnExpire(
        ZKClients.retryOnFailure(
          ZKClientService.Builder.of(connInfo.getZookeeperConnectionString())
            .setSessionTimeout(config.getInt(HConstants.ZK_SESSION_TIMEOUT,
                    HConstants.DEFAULT_ZK_SESSION_TIMEOUT))
            .build(),
          RetryStrategies.exponentialDelay(500, retryTimeOut, TimeUnit.MILLISECONDS)
        )
      )
    );

    DiscoveryService discovery = new ZKDiscoveryService(zkClient);
    TransactionManager txManager = new TransactionManager(config, new HDFSTransactionStateStorage(config, 
            new SnapshotCodecProvider(config), new TxMetricsCollector()), new TxMetricsCollector());
    TransactionService txService = new TransactionService(config, zkClient, discovery, Providers.of(txManager));
    TephraTransactionService service = new TephraTransactionService(zkClient, txService);
    service.start();
    return service;
}
 

@BeforeClass
public static void beforeClass() {
  zkServer = InMemoryZKServer.builder().setTickTime(100000).build();
  zkServer.startAndWait();

  zkClient = ZKClientServices.delegate(
    ZKClients.retryOnFailure(
      ZKClients.reWatchOnExpire(
        ZKClientService.Builder.of(zkServer.getConnectionStr()).build()),
      RetryStrategies.fixDelay(1, TimeUnit.SECONDS)));
  zkClient.startAndWait();
}
 

/**
 * Creates a {@link ZKClientService}.
 */
protected final ZKClientService createZKClient() {
  TwillRuntimeSpecification twillRuntimeSpec = getTwillRuntimeSpecification();

  return ZKClientServices.delegate(
    ZKClients.namespace(
      ZKClients.reWatchOnExpire(
        ZKClients.retryOnFailure(
          ZKClientService.Builder.of(twillRuntimeSpec.getZkConnectStr()).build(),
          RetryStrategies.fixDelay(1, TimeUnit.SECONDS)
        )
      ), "/" + twillRuntimeSpec.getTwillAppName()
    )
  );
}
 

private ZKClientService getZKClientService(String zkConnect) {
  return ZKClientServices.delegate(
    ZKClients.reWatchOnExpire(
      ZKClients.retryOnFailure(ZKClientService.Builder.of(zkConnect)
                                 .setSessionTimeout(ZK_TIMEOUT)
                                 .build(), RetryStrategies.exponentialDelay(100, 2000, TimeUnit.MILLISECONDS))));
}
 

private void updateController(final String appName, final RunId runId, final AtomicBoolean cancelled) {
  String instancePath = String.format("/%s/instances/%s", appName, runId.getId());

  // Fetch the content node.
  Futures.addCallback(zkClientService.getData(instancePath), new FutureCallback<NodeData>() {
    @Override
    public void onSuccess(NodeData result) {
      if (cancelled.get()) {
        return;
      }

      ApplicationMasterLiveNodeData amLiveNodeData = ApplicationMasterLiveNodeDecoder.decode(result);
      if (amLiveNodeData == null) {
        return;
      }

      synchronized (YarnTwillRunnerService.this) {
        if (!controllers.contains(appName, runId)) {
          ZKClient zkClient = ZKClients.namespace(zkClientService, "/" + appName);
          YarnAppClient yarnAppClient = new VersionDetectYarnAppClientFactory().create(new Configuration(yarnConfig));

          YarnTwillController controller = listenController(
            new YarnTwillController(appName, runId, zkClient, amLiveNodeData, yarnAppClient));
          controllers.put(appName, runId, controller);
          controller.start();
        }
      }
    }

    @Override
    public void onFailure(Throwable t) {
      LOG.warn("Failed in fetching application instance node.", t);
    }
  }, Threads.SAME_THREAD_EXECUTOR);
}
 

@Override
public void start() {
  Preconditions.checkNotNull(zkConnectStr);

  eventConverter = new LogEventConverter(hostname, runnableName);
  scheduler = Executors.newSingleThreadScheduledExecutor(Threads.createDaemonThreadFactory(PUBLISH_THREAD_NAME));

  zkClientService = ZKClientServices.delegate(
    ZKClients.reWatchOnExpire(
      ZKClients.retryOnFailure(ZKClientService.Builder.of(zkConnectStr).build(),
                               RetryStrategies.fixDelay(1, TimeUnit.SECONDS))));

  kafkaClient = new ZKKafkaClientService(zkClientService);
  Futures.addCallback(Services.chainStart(zkClientService, kafkaClient),
                      new FutureCallback<List<ListenableFuture<Service.State>>>() {
    @Override
    public void onSuccess(List<ListenableFuture<Service.State>> result) {
      for (ListenableFuture<Service.State> future : result) {
        Preconditions.checkState(Futures.getUnchecked(future) == Service.State.RUNNING,
                                 "Service is not running.");
      }
      addInfo("Kafka client started: " + zkConnectStr);
      scheduler.scheduleWithFixedDelay(flushTask, 0, flushPeriod, TimeUnit.MILLISECONDS);
    }

    @Override
    public void onFailure(Throwable t) {
      // Fail to talk to kafka. Other than logging, what can be done?
      addError("Failed to start kafka appender.", t);
    }
  }, Threads.SAME_THREAD_EXECUTOR);

  super.start();
}
 

@Test (timeout = 30000)
public void testDoubleRegister() throws Exception {
  Map.Entry<DiscoveryService, DiscoveryServiceClient> entry = create();
  try {
    DiscoveryService discoveryService = entry.getKey();
    DiscoveryServiceClient discoveryServiceClient = entry.getValue();

    // Register on the same host port, it shouldn't fail.
    Cancellable cancellable = register(discoveryService, "test_double_reg", "localhost", 54321);
    Cancellable cancellable2 = register(discoveryService, "test_double_reg", "localhost", 54321);

    ServiceDiscovered discoverables = discoveryServiceClient.discover("test_double_reg");

    Assert.assertTrue(waitTillExpected(1, discoverables));

    cancellable.cancel();
    cancellable2.cancel();

    // Register again with two different clients, but killing session of the first one.
    final ZKClientService zkClient2 = ZKClientServices.delegate(
      ZKClients.retryOnFailure(
        ZKClients.reWatchOnExpire(
          ZKClientService.Builder.of(zkServer.getConnectionStr()).build()),
        RetryStrategies.fixDelay(1, TimeUnit.SECONDS)));
    zkClient2.startAndWait();

    try (ZKDiscoveryService discoveryService2 = new ZKDiscoveryService(zkClient2)) {
      cancellable2 = register(discoveryService2, "test_multi_client", "localhost", 54321);

      // Schedule a thread to shutdown zkClient2.
      new Thread() {
        @Override
        public void run() {
          try {
            TimeUnit.SECONDS.sleep(2);
            zkClient2.stopAndWait();
          } catch (InterruptedException e) {
            LOG.error(e.getMessage(), e);
          }
        }
      }.start();

      // This call would block until zkClient2 is shutdown.
      cancellable = register(discoveryService, "test_multi_client", "localhost", 54321);
      cancellable.cancel();
    } finally {
      zkClient2.stopAndWait();
    }
  } finally {
    closeServices(entry);
  }
}
 

/**
 * Launches runnables in the provisioned containers.
 */
private void launchRunnable(List<? extends ProcessLauncher<YarnContainerInfo>> launchers,
                            Queue<ProvisionRequest> provisioning) {
  for (ProcessLauncher<YarnContainerInfo> processLauncher : launchers) {
    LOG.info("Container allocated: {}", processLauncher.getContainerInfo().<Object>getContainer());
    ProvisionRequest provisionRequest = provisioning.peek();
    if (provisionRequest == null) {
      continue;
    }

    String runnableName = provisionRequest.getRuntimeSpec().getName();
    LOG.info("Starting runnable {} in {}", runnableName, processLauncher.getContainerInfo().getContainer());

    int containerCount = expectedContainers.getExpected(runnableName);

    // Setup container environment variables
    Map<String, String> env = new LinkedHashMap<>();
    if (environments.containsKey(runnableName)) {
      env.putAll(environments.get(runnableName));
    }

    ProcessLauncher.PrepareLaunchContext launchContext = processLauncher.prepareLaunch(env,
                                                                                       amLiveNode.getLocalFiles(),
                                                                                       credentials);
    TwillContainerLauncher launcher = new TwillContainerLauncher(
      twillSpec.getRunnables().get(runnableName), processLauncher.getContainerInfo(), launchContext,
      ZKClients.namespace(zkClient, getZKNamespace(runnableName)),
      containerCount, jvmOpts, twillRuntimeSpec.getReservedMemory(runnableName),
      twillRuntimeSpec.getMinHeapRatio(runnableName), getSecureStoreLocation());

    runningContainers.start(runnableName, processLauncher.getContainerInfo(), launcher);

    // Need to call complete to workaround bug in YARN AMRMClient
    if (provisionRequest.containerAcquired()) {
      amClient.completeContainerRequest(provisionRequest.getRequestId());
    }

    /*
     * The provisionRequest will either contain a single container (ALLOCATE_ONE_INSTANCE_AT_A_TIME), or all the
     * containers to satisfy the expectedContainers count. In the later case, the provision request is complete once
     * all the containers have run at which point we poll() to remove the provisioning request.
     */
    if (expectedContainers.getExpected(runnableName) == runningContainers.count(runnableName) ||
      provisioning.peek().getType().equals(AllocationSpecification.Type.ALLOCATE_ONE_INSTANCE_AT_A_TIME)) {
      provisioning.poll();
    }
    if (expectedContainers.getExpected(runnableName) == runningContainers.count(runnableName)) {
      LOG.info("Runnable {} fully provisioned with {} instances.", runnableName, containerCount);
    }
  }
}
 

/**
 * Returns a {@link ZKClient} that namespaced under the given run id.
 */
private static ZKClient getAppRunZKClient(ZKClient zkClient, RunId appRunId) {
  return ZKClients.namespace(zkClient, String.format("/%s", appRunId));
}
 

private static ZKClient getContainerZKClient(ZKClient zkClient, RunId appRunId, String runnableName) {
  return ZKClients.namespace(zkClient, String.format("/%s/runnables/%s", appRunId, runnableName));
}