org.apache.curator.framework.listen.Listenable Java Examples

/**
 * Constructor.
 *
 * @param genieHostInfo                    The genie local host information
 * @param serviceDiscovery                 The service discovery client
 * @param taskScheduler                    The task scheduler
 * @param listenableCuratorConnectionState The listenable curator client connection status
 * @param registry                         The metrics registry
 */
public AgentRoutingServiceCuratorDiscoveryImpl(
    final GenieHostInfo genieHostInfo,
    final ServiceDiscovery<Agent> serviceDiscovery,
    final TaskScheduler taskScheduler,
    final Listenable<ConnectionStateListener> listenableCuratorConnectionState,
    final MeterRegistry registry
) {
    this.localHostname = genieHostInfo.getHostname();
    this.serviceDiscovery = serviceDiscovery;
    this.taskScheduler = taskScheduler;
    this.registry = registry;

    // Schedule periodic reconciliation between in-memory connected set and Service Discovery state
    this.taskScheduler.schedule(this::reconcileRegistrationsTask, trigger);

    // Listen for Curator session state changes
    listenableCuratorConnectionState.addListener(this::handleConnectionStateChange);

    // Create gauge metric for agents connected and registered
    registry.gauge(CONNECTED_AGENTS_GAUGE_NAME, EMPTY_TAG_SET, this.connectedAgentsSet, Set::size);
    registry.gaugeMapSize(REGISTERED_AGENTS_GAUGE_NAME, EMPTY_TAG_SET, this.registeredAgentsMap);
}
 

/**
 * Get an implementation of {@link AgentRoutingService} if one hasn't already been defined.
 * This bean is created if Zookeeper is enabled, it uses Curator's {@link ServiceDiscovery}.
 *
 * @param genieHostInfo                    The local genie host information
 * @param serviceDiscovery                 The Zookeeper Curator service discovery
 * @param taskScheduler                    The task scheduler
 * @param listenableCuratorConnectionState the connection state listenable
 * @param registry                         The metrics registry
 * @return A {@link AgentRoutingServiceImpl} instance
 */
@Bean
@ConditionalOnBean(ServiceDiscovery.class)
@ConditionalOnMissingBean(AgentRoutingService.class)
public AgentRoutingService agentRoutingServiceCurator(
    final GenieHostInfo genieHostInfo,
    final ServiceDiscovery<AgentRoutingServiceCuratorDiscoveryImpl.Agent> serviceDiscovery,
    @Qualifier("genieTaskScheduler") final TaskScheduler taskScheduler,
    final Listenable<ConnectionStateListener> listenableCuratorConnectionState,
    final MeterRegistry registry
) {
    return new AgentRoutingServiceCuratorDiscoveryImpl(
        genieHostInfo,
        serviceDiscovery,
        taskScheduler,
        listenableCuratorConnectionState,
        registry
    );
}
 

/**
 * Test expected beans when Zookeeper is enabled.
 */
@Test
void expectedBeansWithZookeeperEnabled() {
    this.contextRunner
        .withUserConfiguration(ZookeeperMockConfig.class)
        .run(
            context -> {
                Assertions.assertThat(context).hasSingleBean(ZookeeperProperties.class);
                Assertions.assertThat(context).hasSingleBean(LeaderInitiatorFactoryBean.class);
                Assertions.assertThat(context).hasSingleBean(ServiceDiscovery.class);
                Assertions.assertThat(context).hasSingleBean(Listenable.class);
            }
        );
}
 

/**
 * Test expected beans when Zookeeper is not enabled.
 */
@Test
void expectedBeansWithZookeeperDisabled() {
    this.contextRunner.run(
        context -> {
            Assertions.assertThat(context).doesNotHaveBean(ZookeeperProperties.class);
            Assertions.assertThat(context).doesNotHaveBean(LeaderInitiatorFactoryBean.class);
            Assertions.assertThat(context).doesNotHaveBean(ServiceDiscovery.class);
            Assertions.assertThat(context).doesNotHaveBean(Listenable.class);
        }
    );
}
 

/**
 * The Curator-client connection state listenable.
 *
 * @param client              The curator framework client to use
 * @return {@link ServiceDiscovery} bean for instances of type {@link AgentRoutingServiceCuratorDiscoveryImpl.Agent}
 */
@Bean
@ConditionalOnMissingBean(Listenable.class)
Listenable<ConnectionStateListener> listenableCuratorConnectionState(
    final CuratorFramework client
) {
    return client.getConnectionStateListenable();
}
 

/**
 * Return the cache listenable
 *
 * @return listenable
 */
public Listenable<NodeCacheListener> getListenable()
{
    Preconditions.checkState(state.get() != State.CLOSED, "Closed");

    return listeners;
}
 

@Test
public void assertAddDataListener() {
    TreeCache treeCache = mock(TreeCache.class);
    @SuppressWarnings("unchecked")
    Listenable<TreeCacheListener> listeners = mock(Listenable.class);
    TreeCacheListener listener = mock(TreeCacheListener.class);
    when(treeCache.getListenable()).thenReturn(listeners);
    when(regCenter.getRawCache("/test_job")).thenReturn(treeCache);
    jobNodeStorage.addDataListener(listener);
    verify(listeners).addListener(listener);
}
 

@Test
public void assertAddConnectionStateListener() {
    CuratorFramework client = mock(CuratorFramework.class);
    @SuppressWarnings("unchecked")
    Listenable<ConnectionStateListener> listeners = mock(Listenable.class);
    ConnectionStateListener listener = mock(ConnectionStateListener.class);
    when(client.getConnectionStateListenable()).thenReturn(listeners);
    when(regCenter.getRawClient()).thenReturn(client);
    jobNodeStorage.addConnectionStateListener(listener);
    verify(listeners).addListener(listener);
}
 

@Override
public CuratorFramework newClient(final String connectString, final int sessionTimeoutMs,
                                  final int connectionTimeoutMs, final RetryPolicy retryPolicy,
                                  final ACLProvider aclProvider,
                                  final List<AuthInfo> authorization) {
  final CuratorFramework curator = mock(CuratorFramework.class);

  final RetryLoop retryLoop = mock(RetryLoop.class);
  when(retryLoop.shouldContinue()).thenReturn(false);

  final CuratorZookeeperClient czkClient = mock(CuratorZookeeperClient.class);
  when(czkClient.newRetryLoop()).thenReturn(retryLoop);

  when(curator.getZookeeperClient()).thenReturn(czkClient);

  @SuppressWarnings("unchecked") final Listenable<ConnectionStateListener> mockListener =
      (Listenable<ConnectionStateListener>) mock(Listenable.class);

  when(curator.getConnectionStateListenable()).thenReturn(mockListener);

  final GetChildrenBuilder builder = mock(GetChildrenBuilder.class);
  when(curator.getChildren()).thenReturn(builder);

  try {
    when(builder.forPath(anyString())).thenThrow(
        new KeeperException.ConnectionLossException());
  } catch (Exception ignored) {
    // never throws
  }
  when(curator.newNamespaceAwareEnsurePath(anyString())).thenReturn(mock(EnsurePath.class));

  return curator;
}
 

@Override
public Listenable<UnhandledErrorListener> getUnhandledErrorListenable() {
    throw new UnsupportedOperationException("Not implemented in MockCurator");
}
 

@Override
public Listenable<CuratorCacheListener> listenable()
{
    return listenerManager;
}
 

@Override
public Listenable<CuratorCacheListener> listenable()
{
    return listenerManager;
}
 

@Override
public Listenable<ModeledCacheListener<T>> listenable()
{
    return cache.listenable();
}
 

@Override
public Listenable<UnhandledErrorListener> getUnhandledErrorListenable()
{
    return client.getUnhandledErrorListenable();
}
 

@Override
public Listenable<CuratorListener> getCuratorListenable()
{
    return client.getCuratorListenable();
}
 

@Override
public Listenable<UnhandledErrorListener> getUnhandledErrorListenable()
{
    return unhandledErrorListeners;
}
 

/**
 * Re-attach the {@link TreeCacheListener TreeCacheListeners}.
 */
private void reattachTreeCacheListeners() {

	TreeCache treeCache = zookeeperServiceWatch.getCache();

	Listenable<TreeCacheListener> listenable = treeCache.getListenable();

	/**
	 * All registered TreeCacheListeners except {@link ZookeeperServiceWatch}.
	 * Usually, "otherListeners" will be empty because Spring Cloud Zookeeper only
	 * adds "zookeeperServiceWatch" bean as {@link TreeCacheListener}.
	 */
	List<TreeCacheListener> otherListeners = new LinkedList<>();

	if (listenable instanceof ListenerContainer) {
		ListenerContainer<TreeCacheListener> listenerContainer = (ListenerContainer) listenable;
		listenerContainer.forEach(listener -> {
			/**
			 * Even though "listener" is an instance of
			 * {@link ZookeeperServiceWatch}, "zookeeperServiceWatch" bean that
			 * was enhanced by AOP is different from the former, thus it's
			 * required to exclude "listener".
			 */
			if (!(listener instanceof ZookeeperServiceWatch)) {
				otherListeners.add(listener);
			}
			return null;
		});

		// remove all TreeCacheListeners temporarily
		listenerContainer.clear();
		// re-store zookeeperServiceWatch, and make sure it's first one
		// now "beforeChildEvent" is available for Spring AOP
		listenerContainer.addListener(zookeeperServiceWatch);
		// re-store others
		otherListeners.forEach(listenerContainer::addListener);
	}
	else {
		if (logger.isWarnEnabled()) {
			logger.warn(
					"Tell me which version Curator framework current application used? I will do better :D");
		}
	}
}
 

@Override
public Listenable<ConnectionStateListener> getConnectionStateListenable() {
    return new MockListenable<>();
}
 

@Override
public Listenable<CuratorListener> getCuratorListenable() {
    throw new UnsupportedOperationException("Not implemented in MockCurator");
}
 

public Listenable<ModeledCacheListener<T>> listenable()
{
    return listenerContainer;
}
 

/**
 * Allows catching unhandled errors in asynchronous operations.
 *
 * TODO: consider making public.
 */
@VisibleForTesting
public Listenable<UnhandledErrorListener> getUnhandledErrorListenable()
{
    return errorListeners;
}
 

@Override
public Listenable<QueuePutListener<T>> getPutListenerContainer()
{
    return queue.getPutListenerContainer();
}
 

/**
 * Return the manager for put listeners
 *
 * @return put listener container
 */
@Override
public Listenable<QueuePutListener<T>> getPutListenerContainer()
{
    return queue.getPutListenerContainer();
}
 

/**
 * Return the manager for put listeners
 *
 * @return put listener container
 */
@Override
public Listenable<QueuePutListener<T>> getPutListenerContainer()
{
    return putListenerContainer;
}
 

/**
 * Return the manager for put listeners
 *
 * @return put listener container
 */
@Override
public Listenable<QueuePutListener<T>> getPutListenerContainer()
{
    return queue.getPutListenerContainer();
}
 

@Override
public Listenable<ConnectionStateListener> getConnectionStateListenable() {
  return client.getConnectionStateListenable();
}
 

@Override
public Listenable<ConnectionStateListener> getConnectionStateListenable() {
  return client.getConnectionStateListenable();
}
 

@Override
public Listenable<WorkflowListener> getListenable()
{
    return listenerContainer;
}
 

@Provides
@Singleton
CuratorFramework provideCuratorFramework(
    ShutdownRegistry shutdownRegistry,
    @ServiceDiscoveryBindings.ZooKeeper Iterable<InetSocketAddress> zooKeeperCluster,
    ACLProvider aclProvider,
    StatsProvider statsProvider) {

  String connectString =
      FluentIterable.from(zooKeeperCluster)
          .transform(InetSocketAddressHelper::toString)
          .join(Joiner.on(','));

  if (zooKeeperConfig.getChrootPath().isPresent()) {
    connectString = connectString + zooKeeperConfig.getChrootPath().get();
  }

  // export current connection state
  for (ConnectionState connectionState : ConnectionState.values()) {
    statsProvider.makeGauge(
        zkConnectionGaugeName(connectionState),
        new Supplier<Integer>() {
          @Override
          public Integer get() {
            return connectionState.equals(currentState) ? 1 : 0;
          }
        }
    );
  }

  // connection state counter
  AtomicLong zkConnectionConnectedCounter =
      statsProvider.makeCounter(zkConnectionStateCounterName(ConnectionState.CONNECTED));
  AtomicLong zkConnectionReadonlyCounter =
      statsProvider.makeCounter(zkConnectionStateCounterName(ConnectionState.READ_ONLY));
  AtomicLong zkConnectionSuspendedCounter =
      statsProvider.makeCounter(zkConnectionStateCounterName(ConnectionState.SUSPENDED));
  AtomicLong zkConnectionReconnectedCounter =
      statsProvider.makeCounter(zkConnectionStateCounterName(ConnectionState.RECONNECTED));
  AtomicLong zkConnectionLostCounter =
      statsProvider.makeCounter(zkConnectionStateCounterName(ConnectionState.LOST));

  // This emulates the default BackoffHelper configuration used by the legacy commons/zookeeper
  // stack. BackoffHelper is unbounded, this dies after around 5 minutes using the 10 retries.
  // NB: BoundedExponentialBackoffRetry caps max retries at 29 if you send it a larger value.
  RetryPolicy retryPolicy =
      new BoundedExponentialBackoffRetry(
          Amount.of(1, Time.SECONDS).as(Time.MILLISECONDS),
          Amount.of(1, Time.MINUTES).as(Time.MILLISECONDS),
          10);

  CuratorFrameworkFactory.Builder builder = CuratorFrameworkFactory.builder()
      .dontUseContainerParents() // Container nodes are only available in ZK 3.5+.
      .connectString(connectString)
      .canBeReadOnly(false) // We must be able to write to perform leader election.
      .sessionTimeoutMs(zooKeeperConfig.getSessionTimeout().as(Time.MILLISECONDS))
      .connectionTimeoutMs(zooKeeperConfig.getConnectionTimeout().as(Time.MILLISECONDS))
      .retryPolicy(retryPolicy)
      .aclProvider(aclProvider);

  if (zooKeeperConfig.getCredentials().isPresent()) {
    Credentials credentials = zooKeeperConfig.getCredentials().get();
    builder.authorization(credentials.scheme(), credentials.authToken());
  }

  CuratorFramework curatorFramework = builder.build();
  Listenable<ConnectionStateListener> connectionStateListener = curatorFramework
      .getConnectionStateListenable();
  connectionStateListener.addListener((CuratorFramework client, ConnectionState newState) -> {
    currentState = newState;
    switch (newState) {
      case CONNECTED:
        zkConnectionConnectedCounter.getAndIncrement();
        break;
      case READ_ONLY:
        zkConnectionReadonlyCounter.getAndIncrement();
        break;
      case SUSPENDED:
        zkConnectionSuspendedCounter.getAndIncrement();
        break;
      case RECONNECTED:
        zkConnectionReconnectedCounter.getAndIncrement();
        break;
      case LOST:
        zkConnectionLostCounter.getAndIncrement();
        break;
      default:
        currentState = null;
        break;
    }
  });

  // TODO(John Sirois): It would be nice to use a Service to control the lifecycle here, but other
  // services (org.apache.aurora.scheduler.http.JettyServerModule.RedirectMonitor) rely on this
  // service being started 1st which is not deterministic as things stand.  Find a way to leverage
  // the Service system for services with Service dependencies.
  curatorFramework.start();
  shutdownRegistry.addAction(curatorFramework::close);

  return curatorFramework;
}
 

@Bean
@SuppressWarnings("unchecked")
Listenable<ConnectionStateListener> listenableCuratorConnectionState() {
    return Mockito.mock(Listenable.class);
}