Java Code Examples for org.apache.twill.zookeeper.ZKClientService#startAndWait()

/**
 * Invoked by jsvc to start the program.
 */
public void start() throws Exception {
  Injector injector = Guice.createInjector(
    new ConfigModule(conf),
    new ZKModule(),
    new DiscoveryModules().getDistributedModules(),
    new TransactionModules().getDistributedModules(),
    new TransactionClientModule()
  );

  ZKClientService zkClientService = injector.getInstance(ZKClientService.class);
  zkClientService.startAndWait();

  // start a tx server
  txService = injector.getInstance(TransactionService.class);
  try {
    LOG.info("Starting {}", getClass().getSimpleName());
    txService.startAndWait();
  } catch (Exception e) {
    System.err.println("Failed to start service: " + e.getMessage());
  }
}
 

@Provides
@Singleton
private DiscoveryService providesDiscoveryService(final ZKClientService zkClient,
                                                  final ZKDiscoveryService delegate) {
  return new DiscoveryService() {
    @Override
    public Cancellable register(Discoverable discoverable) {
      if (!zkClient.isRunning()) {
        zkClient.startAndWait();
      }
      return delegate.register(discoverable);
    }
  };
}
 

@Provides
@Singleton
private DiscoveryServiceClient providesDiscoveryServiceClient(final ZKClientService zkClient,
                                                              final ZKDiscoveryService delegate) {
  return new DiscoveryServiceClient() {
    @Override
    public ServiceDiscovered discover(String s) {
      if (!zkClient.isRunning()) {
        zkClient.startAndWait();
      }
      return delegate.discover(s);
    }
  };
}
 

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

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

@VisibleForTesting
public static void doMain(boolean verbose, Configuration conf) throws Exception {
  LOG.info("Starting tx server client test.");
  Injector injector = Guice.createInjector(
    new ConfigModule(conf),
    new ZKModule(),
    new DiscoveryModules().getDistributedModules(),
    new TransactionModules().getDistributedModules(),
    new TransactionClientModule()
  );

  ZKClientService zkClient = injector.getInstance(ZKClientService.class);
  zkClient.startAndWait();

  try {
    TransactionServiceClient client = injector.getInstance(TransactionServiceClient.class);
    LOG.info("Starting tx...");
    Transaction tx = client.startShort();
    if (verbose) {
      LOG.info("Started tx details: " + tx.toString());
    } else {
      LOG.info("Started tx: " + tx.getTransactionId() +
                 ", readPointer: " + tx.getReadPointer() +
                 ", invalids: " + tx.getInvalids().length +
                 ", inProgress: " + tx.getInProgress().length);
    }
    try {
      LOG.info("Checking if canCommit tx...");
      client.canCommitOrThrow(tx, Collections.<byte[]>emptyList());
      LOG.info("canCommit: success");
      LOG.info("Committing tx...");
      client.commitOrThrow(tx);
      LOG.info("Committed tx: success");
    } catch (TransactionConflictException e) {
      LOG.info("Aborting tx...");
      client.abort(tx);
      LOG.info("Aborted tx...");
    }
  } finally {
    zkClient.stopAndWait();
  }
}
 

private void startClientAndTestPool(Configuration conf) throws Exception {
  Injector injector = Guice.createInjector(
    new ConfigModule(conf),
    new ZKModule(),
    new DiscoveryModules().getDistributedModules(),
    new TransactionModules().getDistributedModules(),
    new TransactionClientModule()
  );

  ZKClientService zkClient = injector.getInstance(ZKClientService.class);
  zkClient.startAndWait();

  final PooledClientProvider clientProvider = new PooledClientProvider(conf,
    injector.getInstance(DiscoveryServiceClient.class));

  // Test simple case of get + return. Note: this also initializes the provider's pool, which
  // takes about one second (discovery). Doing it before we test the threads makes it so that one
  // thread doesn't take exceptionally longer than the others.

  // Need to retry, since TransactionServiceMain#start returning doesn't indicate that the TransactionService
  // has registered itself for discovery yet
  Tests.waitFor("Failed to get client.", new Callable<Boolean>() {
    @SuppressWarnings({"unused", "EmptyTryBlock"})
    @Override
    public Boolean call() throws Exception {
      try (CloseableThriftClient closeableThriftClient = clientProvider.getCloseableClient()) {
        // do nothing with the client
      } catch (TException e) {
        // simply retry
        return false;
      }
      return true;
    }
  });

  //Now race to get MAX_CLIENT_COUNT+1 clients, exhausting the pool and requesting 1 more.
  List<Future<Integer>> clientIds = new ArrayList<>();
  // We want to ensure that all clients have been exhausted before releasing any.
  // Only once all the clients are fetched from the pool, will any be released. The last thread will reuse one of
  // these clients from the pool.
  CountDownLatch clientDoneLatch = new CountDownLatch(MAX_CLIENT_COUNT);
  ExecutorService executor = Executors.newFixedThreadPool(MAX_CLIENT_COUNT + 1);
  for (int i = 0; i < MAX_CLIENT_COUNT + 1; i++) {
    clientIds.add(executor.submit(new RetrieveClient(clientProvider, clientDoneLatch)));
  }

  Set<Integer> ids = new HashSet<>();
  for (Future<Integer> id : clientIds) {
    ids.add(id.get());
  }
  Assert.assertEquals(MAX_CLIENT_COUNT, ids.size());

  // Now, try it again with, with a countdown latch equal to the number of threads. All of them will only progress
  // past it, once they all acquire a client or time out while attempting to obtain one.
  // One of the threads should throw a TimeOutException, because the other threads don't release their clients
  // until then and the client thread pool isn't enough for the number of threads.
  clientDoneLatch = new CountDownLatch(MAX_CLIENT_COUNT + 1);
  for (int i = 0; i < MAX_CLIENT_COUNT + 1; i++) {
    clientIds.add(executor.submit(new RetrieveClient(clientProvider, clientDoneLatch)));
  }
  int numTimeoutExceptions = 0;
  for (Future<Integer> clientId : clientIds) {
    try {
      clientId.get();
    } catch (ExecutionException expected) {
      Assert.assertEquals(TimeoutException.class, expected.getCause().getClass());
      numTimeoutExceptions++;
    }
  }
  // expect that exactly one of the threads hit the TimeoutException
  Assert.assertEquals(String.format("Expected one thread to not obtain a client within %s milliseconds.",
                                    CLIENT_OBTAIN_TIMEOUT),
                      1, numTimeoutExceptions);

  executor.shutdownNow();
}
 

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

private ZKClientService createZKClient() {
  ZKClientService zkClient = ZKClientService.Builder.of(zkServer.getConnectionStr()).build();
  zkClient.startAndWait();

  return zkClient;
}
 

@Test(timeout = 150000)
public void testDisconnect() throws IOException, InterruptedException {
  File zkDataDir = tmpFolder.newFolder();
  InMemoryZKServer ownZKServer = InMemoryZKServer.builder().setDataDir(zkDataDir).build();
  ownZKServer.startAndWait();
  try {
    ZKClientService zkClient = ZKClientService.Builder.of(ownZKServer.getConnectionStr()).build();
    zkClient.startAndWait();

    try {
      final Semaphore leaderSem = new Semaphore(0);
      final Semaphore followerSem = new Semaphore(0);

      LeaderElection leaderElection = new LeaderElection(zkClient, "/testDisconnect", new ElectionHandler() {
        @Override
        public void leader() {
          leaderSem.release();
        }

        @Override
        public void follower() {
          followerSem.release();
        }
      });
      leaderElection.start();

      leaderSem.tryAcquire(20, TimeUnit.SECONDS);

      int zkPort = ownZKServer.getLocalAddress().getPort();

      // Disconnect by shutting the server and restart it on the same port
      ownZKServer.stopAndWait();

      // Right after disconnect, it should become follower
      followerSem.tryAcquire(20, TimeUnit.SECONDS);

      ownZKServer = InMemoryZKServer.builder().setDataDir(zkDataDir).setPort(zkPort).build();
      ownZKServer.startAndWait();

      // Right after reconnect, it should be leader again.
      leaderSem.tryAcquire(20, TimeUnit.SECONDS);

      // Now disconnect it again, but then cancel it before reconnect, it shouldn't become leader
      ownZKServer.stopAndWait();

      // Right after disconnect, it should become follower
      followerSem.tryAcquire(20, TimeUnit.SECONDS);

      ListenableFuture<?> cancelFuture = leaderElection.stop();

      ownZKServer = InMemoryZKServer.builder().setDataDir(zkDataDir).setPort(zkPort).build();
      ownZKServer.startAndWait();

      Futures.getUnchecked(cancelFuture);

      // After reconnect, it should not be leader
      Assert.assertFalse(leaderSem.tryAcquire(10, TimeUnit.SECONDS));
    } finally {
      zkClient.stopAndWait();
    }
  } finally {
    ownZKServer.stopAndWait();
  }
}
 

@Test
public void testRace() throws InterruptedException {
  ExecutorService executor = Executors.newFixedThreadPool(2);
  final AtomicInteger leaderCount = new AtomicInteger(0);
  final CountDownLatch completeLatch = new CountDownLatch(2);

  // Starts two threads and try to compete for leader and immediate drop leadership.
  // This is to test the case when a follower tries to watch for leader node, but the leader is already gone
  for (int i = 0; i < 2; i++) {
    final ZKClientService zkClient = ZKClientService.Builder.of(zkServer.getConnectionStr()).build();
    zkClient.startAndWait();
    executor.execute(new Runnable() {
      @Override
      public void run() {
        try {
          for (int i = 0; i < 1000; i++) {
            final CountDownLatch leaderLatch = new CountDownLatch(1);
            LeaderElection election = new LeaderElection(zkClient, "/testRace", new ElectionHandler() {
              @Override
              public void leader() {
                leaderCount.incrementAndGet();
                leaderLatch.countDown();
              }

              @Override
              public void follower() {
                // no-op
              }
            });
            election.startAndWait();
            Uninterruptibles.awaitUninterruptibly(leaderLatch);
            election.stopAndWait();
          }
          completeLatch.countDown();
        } finally {
          zkClient.stopAndWait();
        }
      }
    });
  }

  try {
    Assert.assertTrue(completeLatch.await(2, TimeUnit.MINUTES));
  } finally {
    executor.shutdownNow();
  }
}
 

@Test
public void testZKCleanup() throws Exception {
  final ZKClientService zkClient = ZKClientService.Builder.of(getZKConnectionString() + "/twill").build();
  zkClient.startAndWait();

  try {
    TwillRunner runner = getTwillRunner();

    // Start an application and stop it.
    TwillController controller = runner.prepare(new EchoServer())
      .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out, true)))
      .withApplicationArguments("echo")
      .withArguments("EchoServer", "echo2")
      .start();

    Iterable<Discoverable> echoServices = controller.discoverService("echo");
    Assert.assertTrue(waitForSize(echoServices, 1, 120));

    controller.terminate().get();

    // Verify the ZK node gets cleanup
    waitFor(null, new Callable<Stat>() {
      @Override
      public Stat call() throws Exception {
        return zkClient.exists("/EchoServer").get();
      }
    }, 10000, 100, TimeUnit.MILLISECONDS);

    // Start two instances of the application and stop one of it
    List<TwillController> controllers = new ArrayList<>();
    for (int i = 0; i < 2; i++) {
      controllers.add(runner.prepare(new EchoServer())
                        .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out, true)))
                        .withApplicationArguments("echo")
                        .withArguments("EchoServer", "echo2")
                        .start());
    }

    // There should be two instances up and running.
    echoServices = controllers.get(1).discoverService("echo");
    Assert.assertTrue(waitForSize(echoServices, 2, 120));

    // Stop one instance of the app
    controllers.get(0).terminate().get();

    // Verify the ZK node should still be there
    Assert.assertNotNull(zkClient.exists("/EchoServer").get());

    // We should still be able to do discovery, which depends on the ZK node.
    echoServices = controllers.get(1).discoverService("echo");
    Assert.assertTrue(waitForSize(echoServices, 1, 120));

    // Stop second instance of the app
    controllers.get(1).terminate().get();

    // Verify the ZK node gets cleanup
    waitFor(null, new Callable<Stat>() {
      @Override
      public Stat call() throws Exception {
        return zkClient.exists("/EchoServer").get();
      }
    }, 10000, 100, TimeUnit.MILLISECONDS);

  } finally {
    zkClient.stopAndWait();
  }
}
 

@Test
public void testKafkaClientReconnect() throws Exception {
  String topic = "backoff";
  Properties kafkaServerConfig = generateKafkaConfig(zkServer.getConnectionStr() + "/backoff");
  EmbeddedKafkaServer server = new EmbeddedKafkaServer(kafkaServerConfig);

  ZKClientService zkClient = ZKClientService.Builder.of(zkServer.getConnectionStr() + "/backoff").build();
  zkClient.startAndWait();
  try {
    zkClient.create("/", null, CreateMode.PERSISTENT).get();

    ZKKafkaClientService kafkaClient = new ZKKafkaClientService(zkClient);
    kafkaClient.startAndWait();

    try {
      server.startAndWait();
      try {
        // Publish a messages
        createPublishThread(kafkaClient, topic, Compression.NONE, "First message", 1).start();

        // Create a consumer
        final BlockingQueue<String> queue = new LinkedBlockingQueue<>();
        Cancellable cancel = kafkaClient.getConsumer().prepare().add(topic, 0, 0)
          .consume(new KafkaConsumer.MessageCallback() {
            @Override
            public long onReceived(Iterator<FetchedMessage> messages) {
              long nextOffset = -1L;
              while (messages.hasNext()) {
                FetchedMessage message = messages.next();
                nextOffset = message.getNextOffset();
                queue.offer(Charsets.UTF_8.decode(message.getPayload()).toString());
              }
              return nextOffset;
            }

            @Override
            public void finished() {
            }
          });

        // Wait for the first message
        Assert.assertEquals("0 First message", queue.poll(60, TimeUnit.SECONDS));

        // Shutdown the server
        server.stopAndWait();

        // Start the server again.
        // Needs to create a new instance with the same config since guava service cannot be restarted
        server = new EmbeddedKafkaServer(kafkaServerConfig);
        server.startAndWait();

        // Wait a little while to make sure changes is reflected in broker service
        TimeUnit.SECONDS.sleep(3);

        // Publish another message
        createPublishThread(kafkaClient, topic, Compression.NONE, "Second message", 1).start();

        // Should be able to get the second message
        Assert.assertEquals("0 Second message", queue.poll(60, TimeUnit.SECONDS));

        cancel.cancel();
      } finally {
        kafkaClient.stopAndWait();
      }
    } finally {
      server.stopAndWait();
    }
  } finally {
    zkClient.stopAndWait();
  }
}
 

@Test
public void testBrokerChange() throws Exception {
  // Create a new namespace in ZK for Kafka server for this test case
  String connectionStr = zkServer.getConnectionStr() + "/broker_change";
  ZKClientService zkClient = ZKClientService.Builder.of(connectionStr).build();
  zkClient.startAndWait();
  zkClient.create("/", null, CreateMode.PERSISTENT).get();

  // Start a new kafka server
  File logDir = TMP_FOLDER.newFolder();
  EmbeddedKafkaServer server = new EmbeddedKafkaServer(generateKafkaConfig(connectionStr, logDir));
  server.startAndWait();

  // Start a Kafka client
  KafkaClientService kafkaClient = new ZKKafkaClientService(zkClient);
  kafkaClient.startAndWait();

  // Attach a consumer
  final BlockingQueue<String> consumedMessages = Queues.newLinkedBlockingQueue();
  kafkaClient.getConsumer()
    .prepare().addFromBeginning("test", 0).consume(new KafkaConsumer.MessageCallback() {
    @Override
    public long onReceived(Iterator<FetchedMessage> messages) {
      long nextOffset = -1L;
      while (messages.hasNext()) {
        FetchedMessage message = messages.next();
        nextOffset = message.getNextOffset();
        consumedMessages.add(Charsets.UTF_8.decode(message.getPayload()).toString());
      }
      return nextOffset;
    }

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

  // Get a publisher and publish a message
  KafkaPublisher publisher = kafkaClient.getPublisher(KafkaPublisher.Ack.FIRE_AND_FORGET, Compression.NONE);
  publisher.prepare("test").add(Charsets.UTF_8.encode("Message 0"), 0).send().get();

  // Should receive one message
  Assert.assertEquals("Message 0", consumedMessages.poll(5, TimeUnit.SECONDS));

  // Now shutdown and restart the server on different port
  server.stopAndWait();
  server = new EmbeddedKafkaServer(generateKafkaConfig(connectionStr, logDir));
  server.startAndWait();

  // Wait a little while to make sure changes is reflected in broker service
  TimeUnit.SECONDS.sleep(3);

  // Now publish again with the same publisher. It should succeed and the consumer should receive the message.
  publisher.prepare("test").add(Charsets.UTF_8.encode("Message 1"), 0).send().get();
  Assert.assertEquals("Message 1", consumedMessages.poll(5, TimeUnit.SECONDS));

  kafkaClient.stopAndWait();
  zkClient.stopAndWait();
  server.stopAndWait();
}