/**
  * Allocate a new builder. {@link ServiceProviderBuilder#providerStrategy} is set to {@link
  * RoundRobinStrategy}
  *
  * @return the builder
  */
 @Override
 public ServiceProviderBuilder<T> serviceProviderBuilder() {
   return new ServiceProviderBuilderImpl<T>(this)
       .providerStrategy(new RoundRobinStrategy<T>())
       .threadFactory(ThreadUtils.newThreadFactory("ServiceProvider"));
 }
 /**
  * Allocate a new service cache builder. The refresh padding is defaulted to 1 second.
  *
  * @return new cache builder
  */
 @Override
 public ServiceCacheBuilder<T> serviceCacheBuilder() {
   return new ServiceCacheBuilderImpl<T>(this)
       .threadFactory(ThreadUtils.newThreadFactory("ServiceCache"));
 }
public class TestLeaderSelectorWithExecutor extends BaseClassForTests {
  private static final ThreadFactory threadFactory = ThreadUtils.newThreadFactory("FeedGenerator");

  @Test
  public void test() throws Exception {
    Timing timing = new Timing();
    LeaderSelector leaderSelector = null;
    CuratorFramework client =
        CuratorFrameworkFactory.builder()
            .retryPolicy(new ExponentialBackoffRetry(100, 3))
            .connectString(server.getConnectString())
            .sessionTimeoutMs(timing.session())
            .connectionTimeoutMs(timing.connection())
            .build();
    try {
      client.start();

      MyLeaderSelectorListener listener = new MyLeaderSelectorListener();
      ExecutorService executorPool = Executors.newFixedThreadPool(20);
      leaderSelector = new LeaderSelector(client, "/test", threadFactory, executorPool, listener);

      leaderSelector.autoRequeue();
      leaderSelector.start();

      timing.sleepABit();

      Assert.assertEquals(listener.getLeaderCount(), 1);
    } finally {
      CloseableUtils.closeQuietly(leaderSelector);
      CloseableUtils.closeQuietly(client);
    }
  }

  private class MyLeaderSelectorListener implements LeaderSelectorListener {
    private volatile Thread ourThread;
    private final AtomicInteger leaderCount = new AtomicInteger(0);

    public int getLeaderCount() {
      return leaderCount.get();
    }

    @Override
    public void takeLeadership(CuratorFramework curatorFramework) throws Exception {
      ourThread = Thread.currentThread();
      try {
        leaderCount.incrementAndGet();
        while (!Thread.currentThread().isInterrupted()) {
          Thread.sleep(1000);
        }
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
      } finally {
        leaderCount.decrementAndGet();
      }
    }

    @Override
    public void stateChanged(CuratorFramework curatorFramework, ConnectionState newState) {
      if ((newState == ConnectionState.LOST) || (newState == ConnectionState.SUSPENDED)) {
        if (ourThread != null) {
          ourThread.interrupt();
        }
      }
    }
  }
}