Java Code Examples for org.apache.twill.common.Cancellable#cancel()

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

    Cancellable cancellable = register(discoveryService, "test_expires", "localhost", 54321);

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

    // Discover that registered host:port.
    Assert.assertTrue(waitTillExpected(1, discoverables));

    KillZKSession.kill(zkClient.getZooKeeperSupplier().get(), zkServer.getConnectionStr(), 10000);

    // Register one more endpoint to make sure state has been reflected after reconnection
    Cancellable cancellable2 = register(discoveryService, "test_expires", "localhost", 54322);

    // Reconnection would trigger re-registration.
    Assert.assertTrue(waitTillExpected(2, discoverables));

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

    // Verify that both are now gone.
    Assert.assertTrue(waitTillExpected(0, discoverables));
  } finally {
    closeServices(entry);
  }
}
 

@Test
public void simpleDiscoverable() throws Exception {
  final byte[] payload = "data".getBytes(StandardCharsets.UTF_8);
  Map.Entry<DiscoveryService, DiscoveryServiceClient> entry = create();
  try {
    DiscoveryService discoveryService = entry.getKey();
    DiscoveryServiceClient discoveryServiceClient = entry.getValue();

    // Register one service running on one host:port
    Cancellable cancellable = register(discoveryService, "foo", "localhost", 8090, payload);

    // Discover that registered host:port.
    ServiceDiscovered serviceDiscovered = discoveryServiceClient.discover("foo");
    Assert.assertTrue(waitTillExpected(1, serviceDiscovered));

    Discoverable discoverable = new Discoverable("foo", new InetSocketAddress("localhost", 8090), payload);

    // Check it exists.
    Assert.assertTrue(serviceDiscovered.contains(discoverable));

    // Remove the service
    cancellable.cancel();

    // There should be no service.
    Assert.assertTrue(waitTillExpected(0, serviceDiscovered));

    Assert.assertFalse(serviceDiscovered.contains(discoverable));
  } finally {
    closeServices(entry);
  }
}
 

@Override
protected <R> ProcessController<R> doLaunch(YarnLaunchContext launchContext) {
  Map<String, String> env = Maps.newHashMap(launchContext.getEnvironment());

  // Set extra environments
  env.put(EnvKeys.YARN_CONTAINER_ID, containerInfo.getId());
  env.put(EnvKeys.YARN_CONTAINER_HOST, containerInfo.getHost().getHostName());
  env.put(EnvKeys.YARN_CONTAINER_PORT, Integer.toString(containerInfo.getPort()));
  env.put(EnvKeys.YARN_CONTAINER_MEMORY_MB, Integer.toString(containerInfo.getMemoryMB()));
  env.put(EnvKeys.YARN_CONTAINER_VIRTUAL_CORES, Integer.toString(containerInfo.getVirtualCores()));

  launchContext.setEnvironment(env);

  LOG.info("Launching in container {} at {}:{}, {}",
           containerInfo.getId(), containerInfo.getHost().getHostName(),
           containerInfo.getPort(), launchContext.getCommands());
  final Cancellable cancellable = nmClient.start(containerInfo, launchContext);
  launched = true;

  return new ProcessController<R>() {
    @Override
    public void close() throws Exception {
      // no-op
    }

    @Override
    public R getReport() {
      // No reporting support for runnable launch yet.
      return null;

    }

    @Override
    public void cancel() {
      cancellable.cancel();
    }
  };
}
 

@Override
public void initialize(TwillContext context) {
  super.initialize(context);
  running = true;
  try {
    serverSocket = new ServerSocket(0);
    LOG.info("Server started: " + serverSocket.getLocalSocketAddress() +
             ", id: " + context.getInstanceId() +
             ", count: " + context.getInstanceCount());

    // Announce with service names as specified in app arguments and runnable arguments
    final List<Cancellable> cancellables = new ArrayList<>();
    for (String[] args : new String[][] {context.getApplicationArguments(), context.getArguments()}) {
      if (args.length > 0) {
        cancellables.add(context.announce(args[0], serverSocket.getLocalPort()));
      }
    }
    canceller = new Cancellable() {
      @Override
      public void cancel() {
        for (Cancellable c : cancellables) {
          c.cancel();
        }
      }
    };
  } catch (IOException e) {
    throw Throwables.propagate(e);
  }
}
 

@Override
protected void shutDown() throws Exception {
  LOG.info("Stopping KafkaClientService");
  scheduler.shutdownNow();
  for (Cancellable cancellable : publishers.values()) {
    cancellable.cancel();
  }
  consumer.stop();

  brokerService.stopAndWait();
  LOG.info("KafkaClientService stopped");
}
 

/**
 * Called to stop all consumers created. This method should only be
 * called by KafkaClientService who own this consumer.
 */
void stop() {
  LOG.info("Stopping Kafka consumer");
  List<Cancellable> cancels = Lists.newLinkedList();
  consumerCancels.drainTo(cancels);
  for (Cancellable cancel : cancels) {
    cancel.cancel();
  }
  consumers.invalidateAll();
  LOG.info("Kafka Consumer stopped");
}
 

@Test
public void testKafkaClient() throws Exception {
  String topic = "testClient";

  Thread t1 = createPublishThread(kafkaClient, topic, Compression.GZIP, "GZIP Testing message", 10);
  Thread t2 = createPublishThread(kafkaClient, topic, Compression.NONE, "Testing message", 10);

  t1.start();
  t2.start();

  Thread t3 = createPublishThread(kafkaClient, topic, Compression.SNAPPY, "Snappy Testing message", 10);
  t2.join();
  t3.start();

  final CountDownLatch latch = new CountDownLatch(30);
  final CountDownLatch stopLatch = new CountDownLatch(1);
  Cancellable cancel = kafkaClient.getConsumer().prepare().add(topic, 0, 0).consume(new KafkaConsumer
    .MessageCallback() {
    @Override
    public long onReceived(Iterator<FetchedMessage> messages) {
      long nextOffset = -1;
      while (messages.hasNext()) {
        FetchedMessage message = messages.next();
        nextOffset = message.getNextOffset();
        LOG.info(Charsets.UTF_8.decode(message.getPayload()).toString());
        latch.countDown();
      }
      return nextOffset;
    }

    @Override
    public void finished() {
      stopLatch.countDown();
    }
  });

  Assert.assertTrue(latch.await(5, TimeUnit.SECONDS));
  cancel.cancel();
  Assert.assertTrue(stopLatch.await(1, TimeUnit.SECONDS));
}
 

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

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

    List<Cancellable> cancellables = Lists.newArrayList();

    cancellables.add(register(discoveryService, "service1", "localhost", 1));
    cancellables.add(register(discoveryService, "service1", "localhost", 2));
    cancellables.add(register(discoveryService, "service1", "localhost", 3));
    cancellables.add(register(discoveryService, "service1", "localhost", 4));
    cancellables.add(register(discoveryService, "service1", "localhost", 5));

    cancellables.add(register(discoveryService, "service2", "localhost", 1));
    cancellables.add(register(discoveryService, "service2", "localhost", 2));
    cancellables.add(register(discoveryService, "service2", "localhost", 3));

    cancellables.add(register(discoveryService, "service3", "localhost", 1));
    cancellables.add(register(discoveryService, "service3", "localhost", 2));

    ServiceDiscovered serviceDiscovered = discoveryServiceClient.discover("service1");
    Assert.assertTrue(waitTillExpected(5, serviceDiscovered));

    serviceDiscovered = discoveryServiceClient.discover("service2");
    Assert.assertTrue(waitTillExpected(3, serviceDiscovered));

    serviceDiscovered = discoveryServiceClient.discover("service3");
    Assert.assertTrue(waitTillExpected(2, serviceDiscovered));

    cancellables.add(register(discoveryService, "service3", "localhost", 3));
    Assert.assertTrue(waitTillExpected(3, serviceDiscovered)); // Shows live iterator.

    for (Cancellable cancellable : cancellables) {
      cancellable.cancel();
    }

    Assert.assertTrue(waitTillExpected(0, discoveryServiceClient.discover("service1")));
    Assert.assertTrue(waitTillExpected(0, discoveryServiceClient.discover("service2")));
    Assert.assertTrue(waitTillExpected(0, discoveryServiceClient.discover("service3")));
  } finally {
    closeServices(entry);
  }
}
 

@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 testKafkaClientSkipNext() throws Exception {
  String topic = "testClientSkipNext";
  // Publish 30 messages with indecies the same as offsets within the range 0 - 29
  Thread t1 = createPublishThread(kafkaClient, topic, Compression.GZIP, "GZIP Testing message", 10);
  t1.start();
  t1.join();
  Thread t2 = createPublishThread(kafkaClient, topic, Compression.NONE, "Testing message", 10, 10);
  t2.start();
  t2.join();
  Thread t3 = createPublishThread(kafkaClient, topic, Compression.SNAPPY, "Snappy Testing message", 10, 20);
  t3.start();
  t3.join();

  final CountDownLatch stopLatch = new CountDownLatch(1);
  final BlockingQueue<Long> offsetQueue = 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;
      if (messages.hasNext()) {
        FetchedMessage message = messages.next();
        nextOffset = message.getNextOffset() + 1;
        offsetQueue.offer(message.getOffset());
        LOG.info(Charsets.UTF_8.decode(message.getPayload()).toString());
        return nextOffset;
      }
      return nextOffset;
    }

    @Override
    public void finished() {
      stopLatch.countDown();
    }
  });
  // 15 messages should be in the queue since onReceived returns `message.getNextOffset() + 1` as next offset to read
  for (long i = 0; i < 30; i += 2) {
    Assert.assertEquals(i, (long) offsetQueue.poll(60, TimeUnit.SECONDS));
  }
  Assert.assertNull(offsetQueue.poll(2, TimeUnit.SECONDS));
  cancel.cancel();
  Assert.assertTrue(stopLatch.await(1, TimeUnit.SECONDS));
}