org.apache.twill.zookeeper.ZKClient Java Examples

private Service createService(ZKClient zkClient, RunId runId) {
  return new AbstractTwillService(zkClient, runId) {

    private final CountDownLatch stopLatch = new CountDownLatch(1);

    @Override
    protected void doStart() throws Exception {
      LOG.info("Start");
    }

    @Override
    protected void doRun() throws Exception {
      stopLatch.await();
    }

    @Override
    protected void doStop() throws Exception {
      LOG.info("Stop");
    }

    @Override
    protected void triggerShutdown() {
      stopLatch.countDown();
    }
  };
}
 

/**
 * 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 BasicTwillContext(RunId runId, RunId appRunId, InetAddress host, String[] args, String[] appArgs,
                         TwillRunnableSpecification spec, int instanceId,
                         DiscoveryService discoveryService, DiscoveryServiceClient discoveryServiceClient,
                         ZKClient zkClient,
                         int instanceCount, int allowedMemoryMB, int virtualCores) {
  this.runId = runId;
  this.appRunId = appRunId;
  this.host = host;
  this.args = args;
  this.appArgs = appArgs;
  this.spec = spec;
  this.instanceId = instanceId;
  this.discoveryService = discoveryService;
  this.discoveryServiceClient = discoveryServiceClient;
  this.zkClient = zkClient;
  this.elections = new ElectionRegistry(zkClient);
  this.instanceCount = instanceCount;
  this.allowedMemoryMB = allowedMemoryMB;
  this.virtualCores = virtualCores;
}
 

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

RunningContainers(TwillRuntimeSpecification twillRuntimeSpec, String appId, TwillRunResources appMasterResources,
                  ZKClient zookeeperClient, Location applicationLocation,
                  Map<String, RuntimeSpecification> runnables,
                  EventHandler eventHandler) {
  containers = HashBasedTable.create();
  runnableInstances = Maps.newHashMap();
  completedContainerCount = Maps.newHashMap();
  startSequence = Lists.newLinkedList();
  containerLock = new ReentrantLock();
  containerChange = containerLock.newCondition();
  resourceReport = new DefaultResourceReport(appId, appMasterResources);
  zkClient = zookeeperClient;
  containerStats = HashMultimap.create();
  this.applicationLocation = applicationLocation;
  this.runnableNames = runnables.keySet();
  this.logLevels = new TreeMap<>();
  this.maxRetries = Maps.newHashMap(twillRuntimeSpec.getMaxRetries());
  this.numRetries = Maps.newHashMap();
  this.eventHandler = eventHandler;
}
 

public ApplicationMasterService(RunId runId, ZKClient zkClient,
                                TwillRuntimeSpecification twillRuntimeSpec, YarnAMClient amClient,
                                Configuration config, Location applicationLocation) throws Exception {
  super(zkClient, runId, config, applicationLocation);

  this.runId = runId;
  this.twillRuntimeSpec = twillRuntimeSpec;
  this.zkClient = zkClient;
  this.applicationLocation = applicationLocation;
  this.amClient = amClient;
  this.credentials = createCredentials();
  this.jvmOpts = loadJvmOptions();
  this.twillSpec = twillRuntimeSpec.getTwillSpecification();
  this.placementPolicyManager = new PlacementPolicyManager(twillSpec.getPlacementPolicies());
  this.environments = getEnvironments();

  this.amLiveNode = new ApplicationMasterLiveNodeData(Integer.parseInt(System.getenv(EnvKeys.YARN_APP_ID)),
                                                      Long.parseLong(System.getenv(EnvKeys.YARN_APP_ID_CLUSTER_TIME)),
                                                      amClient.getContainerId().toString(), getLocalizeFiles(),
                                                      twillRuntimeSpec.getKafkaZKConnect());

  this.expectedContainers = new ExpectedContainers(twillSpec);
  this.eventHandler = createEventHandler(twillSpec);
  this.runningContainers = createRunningContainers(amClient.getContainerId(), amClient.getHost());
}
 

@Inject
public TransactionService(Configuration conf,
                          ZKClient zkClient,
                          DiscoveryService discoveryService,
                          Provider<TransactionManager> txManagerProvider) {
  super(conf, discoveryService, txManagerProvider);
  this.conf = conf;
  this.zkClient = zkClient;
}
 

private void configureJob() throws Exception {
    job = Job.getInstance(getConf(),
            "UpdateStatistics-" + tableName + "-" + snapshotName);
    PhoenixMapReduceUtil.setInput(job, NullDBWritable.class,
            snapshotName, tableName, restoreDir);

    PhoenixConfigurationUtil.setMRJobType(job.getConfiguration(), MRJobType.UPDATE_STATS);

    // DO NOT allow mapper splits using statistics since it may result into many smaller chunks
    PhoenixConfigurationUtil.setSplitByStats(job.getConfiguration(), false);

    job.setJarByClass(UpdateStatisticsTool.class);
    job.setMapperClass(TableSnapshotMapper.class);
    job.setMapOutputKeyClass(NullWritable.class);
    job.setMapOutputValueClass(NullWritable.class);
    job.setOutputFormatClass(NullOutputFormat.class);
    job.setNumReduceTasks(0);
    job.setPriority(this.jobPriority);

    TableMapReduceUtil.addDependencyJars(job);
    TableMapReduceUtil.addDependencyJarsForClasses(job.getConfiguration(), PhoenixConnection.class, Chronology.class,
            CharStream.class, TransactionSystemClient.class, TransactionNotInProgressException.class,
            ZKClient.class, DiscoveryServiceClient.class, ZKDiscoveryService.class,
            Cancellable.class, TTransportException.class, SpanReceiver.class, TransactionProcessor.class, Gauge.class, MetricRegistriesImpl.class);
    LOGGER.info("UpdateStatisticsTool running for: " + tableName
            + " on snapshot: " + snapshotName + " with restore dir: " + restoreDir);
}
 

private TwillController getController(ZKClient zkClient, String appName, RunId runId) {
  AbstractTwillController controller = new AbstractTwillController(appName, runId,
                                                                   zkClient, false, ImmutableList.<LogHandler>of()) {

    @Override
    public void kill() {
      // No-op
    }

    @Override
    protected void instanceNodeUpdated(NodeData nodeData) {
      // No-op
    }

    @Override
    protected void instanceNodeFailed(Throwable cause) {
      // Shutdown if the instance node goes away
      if (cause instanceof KeeperException.NoNodeException) {
        forceShutDown();
      }
    }

    @Override
    public ResourceReport getResourceReport() {
      return null;
    }
  };
  controller.startAndWait();
  return controller;
}
 

protected TwillContainerControllerImpl(ZKClient zkClient, RunId runId, String runnable, int instanceId,
                                       ProcessController<Void> processController) {
  super(runId, zkClient);
  this.runnable = runnable;
  this.instanceId = instanceId;
  this.processController = processController;
  this.shutdownLatch = new CountDownLatch(1);
}
 

public TwillContainerLauncher(RuntimeSpecification runtimeSpec, ContainerInfo containerInfo,
                              ProcessLauncher.PrepareLaunchContext launchContext,
                              ZKClient zkClient, int instanceCount, JvmOptions jvmOpts,
                              int reservedMemory, double minHeapRatio,
                              Location secureStoreLocation) {
  this.runtimeSpec = runtimeSpec;
  this.containerInfo = containerInfo;
  this.launchContext = launchContext;
  this.zkClient = zkClient;
  this.instanceCount = instanceCount;
  this.jvmOpts = jvmOpts;
  this.reservedMemory = reservedMemory;
  this.minHeapRatio = minHeapRatio;
  this.secureStoreLocation = secureStoreLocation;
}
 

YarnTwillController(String appName, RunId runId, ZKClient zkClient, boolean logCollectionEnabled,
                    Iterable<LogHandler> logHandlers, Callable<ProcessController<YarnApplicationReport>> startUp,
                    long startTimeout, TimeUnit startTimeoutUnit) {
  super(appName, runId, zkClient, logCollectionEnabled, logHandlers);
  this.appName = appName;
  this.startUp = startUp;
  this.startTimeout = startTimeout;
  this.startTimeoutUnit = startTimeoutUnit;
}
 

/**
 * Creates an instance with an existing {@link ApplicationMasterLiveNodeData}.
 */
YarnTwillController(String appName, RunId runId, ZKClient zkClient,
                    final ApplicationMasterLiveNodeData amLiveNodeData, final YarnAppClient yarnAppClient) {
  super(appName, runId, zkClient, amLiveNodeData.getKafkaZKConnect() != null, Collections.<LogHandler>emptyList());
  this.appName = appName;
  this.amLiveNodeData = amLiveNodeData;
  this.startUp = () -> yarnAppClient.createProcessController(
    ApplicationId.newInstance(amLiveNodeData.getAppIdClusterTime(),
                              amLiveNodeData.getAppId()));
  this.startTimeout = Constants.APPLICATION_MAX_START_SECONDS;
  this.startTimeoutUnit = TimeUnit.SECONDS;
}
 

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

TwillContainerService(BasicTwillContext context, ContainerInfo containerInfo, ZKClient zkClient,
                      RunId runId, TwillRunnableSpecification specification, ClassLoader classLoader,
                      Configuration config, Location applicationLocation,
                      Map<String, String> defaultLogLevels, Map<String, String> logLevels) {
  super(zkClient, runId, config, applicationLocation);

  this.specification = specification;
  this.classLoader = classLoader;
  this.defaultLogLevels = ImmutableMap.copyOf(defaultLogLevels);
  this.oldLogLevels = new HashMap<>(defaultLogLevels);
  this.containerLiveNodeData = createLiveNodeData(
    containerInfo, isLoggerContext() ? logLevels : Collections.<String, String>emptyMap());
  this.context = context;
}
 

@Override
protected void configure() {
  /**
   * ZKClientService is provided by the provider method
   * {@link #provideZKClientService(org.apache.hadoop.conf.Configuration)}.
   */
  bind(ZKClient.class).to(ZKClientService.class);
}
 

/**
 * Creates a distributed lock instance.
 *
 * @param zkClient the {@link ZKClient} to interact with the ZooKeeper used for the lock coordination
 * @param path the path in ZooKeeper where the lock coordination happens
 */
public ReentrantDistributedLock(ZKClient zkClient, String path) {
  this.zkClient = zkClient;
  this.path = path.startsWith("/") ? path : "/" + path;
  this.localLockNode = new ThreadLocal<String>();
  this.lock = new ReentrantLock();
}
 

RewatchOnExpireWatcher(ZKClient client, ActionType actionType, String path, Watcher delegate) {
  this.client = client;
  this.actionType = actionType;
  this.path = path;
  this.delegate = delegate;
  this.lastResult = new AtomicMarkableReference<Object>(null, false);
}
 

public ElectionRegistry(ZKClient zkClient) {
  this.zkClient = zkClient;
  Multimap<String, LeaderElection> multimap = HashMultimap.create();
  this.registry = Multimaps.synchronizedMultimap(multimap);
}
 

@Inject
public TestTransactionService(Configuration conf, ZKClient zkClient,
                              DiscoveryService discoveryService,
                              Provider<TransactionManager> txManagerProvider) {
  super(conf, zkClient, discoveryService, txManagerProvider);
}
 

public ZKKafkaClientService(ZKClient zkClient) {
  this.brokerService = new ZKBrokerService(zkClient);
  this.publishers = Collections.synchronizedMap(new IdentityHashMap<WeakReference<KafkaPublisher>, Cancellable>());
  this.referenceQueue = new ReferenceQueue<KafkaPublisher>();
  this.consumer = new SimpleKafkaConsumer(brokerService);
}
 

public RewatchOnExpireZKClient(ZKClient delegate) {
  super(delegate);
}
 

public NamespaceZKClient(ZKClient delegate, String namespace) {
  super(delegate);
  this.namespace = namespace;
  this.delegate = delegate;
  this.connectString = delegate.getConnectString() + namespace;
}
 

public FailureRetryZKClient(ZKClient delegate, RetryStrategy retryStrategy) {
  super(delegate);
  this.retryStrategy = retryStrategy;
}
 

protected AbstractZKServiceController(RunId runId, ZKClient zkClient) {
  super(runId);
  this.zkClient = zkClient;
  this.instanceNodeDataCallback = new InstanceNodeDataCallback();
  this.messageFutures = Lists.newLinkedList();
}
 

protected AbstractTwillService(final ZKClient zkClient, RunId runId) {
  this.zkClient = zkClient;
  this.runId = runId;
}
 

AppMasterTwillZKPathService(ZKClient zkClient, RunId runId) {
  super(zkClient, runId);
  this.zkClient = zkClient;
}
 

public LeaderElection(ZKClient zkClient, String prefix, ElectionHandler handler) {
  this.guid = UUID.randomUUID().toString();
  this.zkClient = zkClient;
  this.zkFolderPath = prefix.startsWith("/") ? prefix : "/" + prefix;
  this.handler = handler;
}
 

public TwillZKPathService(ZKClient zkClient, RunId runId) {
  this.zkClient = zkClient;
  this.path = "/" + runId.getId();
}
 

protected AbstractYarnTwillService(ZKClient zkClient, RunId runId,
                                   Configuration config, Location applicationLocation) {
  super(zkClient, runId);
  this.config = config;
  this.applicationLocation = applicationLocation;
}