Examples with ZKClientService org.apache.twill.zookeeper.ZKClientService used on opensource projects

Example 6 with ZKClientService

use of org.apache.twill.zookeeper.ZKClientService in project cdap by caskdata.

the class ResourceBalancerServiceTest method testServiceStopFailure.

@Test
public void testServiceStopFailure() throws Exception {
    ZKClientService zkClient = ZKClientService.Builder.of(zkServer.getConnectionStr()).build();
    zkClient.startAndWait();
    try (ZKDiscoveryService discoveryService = new ZKDiscoveryService(zkClient)) {
        // Test the failure on stop case
        final TestBalancerService stopFailureService = new TestBalancerService("test", 4, zkClient, discoveryService, discoveryService, false, true);
        stopFailureService.startAndWait();
        // Should get four partitions
        Tasks.waitFor(ImmutableSet.of(0, 1, 2, 3), new Callable<Set<Integer>>() {

            @Override
            public Set<Integer> call() throws Exception {
                return stopFailureService.getPartitions();
            }
        }, 10, TimeUnit.SECONDS, 100, TimeUnit.MILLISECONDS);
        // Register a new discoverable, this should trigger a partition change in the resource balancer service
        Cancellable cancellable = discoveryService.register(new Discoverable("test", new InetSocketAddress(InetAddress.getLoopbackAddress(), 1234)));
        try {
            // When there is exception thrown by the underlying service during partition change,
            // the resource balancer service should fail.
            Tasks.waitFor(Service.State.FAILED, new Callable<Service.State>() {

                @Override
                public Service.State call() throws Exception {
                    return stopFailureService.state();
                }
            }, 10, TimeUnit.SECONDS, 100, TimeUnit.MILLISECONDS);
        } finally {
            cancellable.cancel();
        }
    } finally {
        zkClient.stopAndWait();
    }
}

Example 7 with ZKClientService

use of org.apache.twill.zookeeper.ZKClientService in project cdap by caskdata.

the class ResourceBalancerServiceTest method testResourceBalancerService.

@Test
public void testResourceBalancerService() throws Exception {
    // Simple test for resource balancer does react to discovery changes correct
    // More detailed tests are in ResourceCoordinatorTest, which the ResourceBalancerService depends on
    ZKClientService zkClient = ZKClientService.Builder.of(zkServer.getConnectionStr()).build();
    zkClient.startAndWait();
    try (ZKDiscoveryService discoveryService = new ZKDiscoveryService(zkClient)) {
        // Test the failure on stop case
        final TestBalancerService stopFailureService = new TestBalancerService("test", 4, zkClient, discoveryService, discoveryService, false, false);
        stopFailureService.startAndWait();
        // Should get all four partitions
        Tasks.waitFor(ImmutableSet.of(0, 1, 2, 3), new Callable<Set<Integer>>() {

            @Override
            public Set<Integer> call() throws Exception {
                return stopFailureService.getPartitions();
            }
        }, 10, TimeUnit.SECONDS, 100, TimeUnit.MILLISECONDS);
        // Register a new discoverable, this should trigger a partition change in the resource balancer service
        Cancellable cancellable = discoveryService.register(new Discoverable("test", new InetSocketAddress(InetAddress.getLoopbackAddress(), 1234)));
        try {
            // Should get reduced to two partitions
            Tasks.waitFor(2, new Callable<Integer>() {

                @Override
                public Integer call() throws Exception {
                    return stopFailureService.getPartitions().size();
                }
            }, 10, TimeUnit.SECONDS, 100, TimeUnit.MILLISECONDS);
        } finally {
            cancellable.cancel();
        }
    } finally {
        zkClient.stopAndWait();
    }
}

Example 8 with ZKClientService

use of org.apache.twill.zookeeper.ZKClientService in project cdap by caskdata.

the class ZKExtOperationsTest method testCreateOrSet.

@Test
public void testCreateOrSet() throws Exception {
    String path = "/parent/testCreateOrSet";
    ZKClientService zkClient = ZKClientService.Builder.of(zkServer.getConnectionStr()).build();
    zkClient.startAndWait();
    // Create with "1"
    Assert.assertEquals(1, ZKExtOperations.createOrSet(zkClient, path, Suppliers.ofInstance(1), INT_CODEC, 0).get().intValue());
    // Should get "1" back
    Assert.assertEquals(1, INT_CODEC.decode(zkClient.getData(path).get().getData()).intValue());
    // Set with "2"
    Assert.assertEquals(2, ZKExtOperations.createOrSet(zkClient, path, Suppliers.ofInstance(2), INT_CODEC, 0).get().intValue());
    // Should get "2" back
    Assert.assertEquals(2, INT_CODEC.decode(zkClient.getData(path).get().getData()).intValue());
    zkClient.stopAndWait();
}

Example 9 with ZKClientService

use of org.apache.twill.zookeeper.ZKClientService in project cdap by caskdata.

the class ResourceCoordinatorTest method testAssignment.

@Test
public void testAssignment() throws InterruptedException, ExecutionException {
    CConfiguration cConf = CConfiguration.create();
    cConf.set(Constants.Zookeeper.QUORUM, zkServer.getConnectionStr());
    String serviceName = "test-assignment";
    Injector injector = Guice.createInjector(new ConfigModule(cConf), new ZKClientModule(), new DiscoveryRuntimeModule().getDistributedModules());
    ZKClientService zkClient = injector.getInstance(ZKClientService.class);
    zkClient.startAndWait();
    DiscoveryService discoveryService = injector.getInstance(DiscoveryService.class);
    try {
        ResourceCoordinator coordinator = new ResourceCoordinator(zkClient, injector.getInstance(DiscoveryServiceClient.class), new BalancedAssignmentStrategy());
        coordinator.startAndWait();
        try {
            ResourceCoordinatorClient client = new ResourceCoordinatorClient(zkClient);
            client.startAndWait();
            try {
                // Create a requirement
                ResourceRequirement requirement = ResourceRequirement.builder(serviceName).addPartitions("p", 5, 1).build();
                client.submitRequirement(requirement).get();
                // Fetch the requirement, just to verify it's the same as the one get submitted.
                Assert.assertEquals(requirement, client.fetchRequirement(requirement.getName()).get());
                // Register a discovery endpoint
                final Discoverable discoverable1 = createDiscoverable(serviceName, 10000);
                Cancellable cancelDiscoverable1 = discoveryService.register(ResolvingDiscoverable.of(discoverable1));
                // Add a change handler for this discoverable.
                final BlockingQueue<Collection<PartitionReplica>> assignmentQueue = new SynchronousQueue<>();
                final Semaphore finishSemaphore = new Semaphore(0);
                Cancellable cancelSubscribe1 = subscribe(client, discoverable1, assignmentQueue, finishSemaphore);
                // Assert that it received the changes.
                Collection<PartitionReplica> assigned = assignmentQueue.poll(30, TimeUnit.SECONDS);
                Assert.assertNotNull(assigned);
                Assert.assertEquals(5, assigned.size());
                // Unregister from discovery, the handler should receive a change with empty collection
                cancelDiscoverable1.cancel();
                Assert.assertTrue(assignmentQueue.poll(30, TimeUnit.SECONDS).isEmpty());
                // Register to discovery again, would receive changes.
                cancelDiscoverable1 = discoveryService.register(ResolvingDiscoverable.of(discoverable1));
                assigned = assignmentQueue.poll(30, TimeUnit.SECONDS);
                Assert.assertNotNull(assigned);
                Assert.assertEquals(5, assigned.size());
                // Register another discoverable
                final Discoverable discoverable2 = createDiscoverable(serviceName, 10001);
                Cancellable cancelDiscoverable2 = discoveryService.register(ResolvingDiscoverable.of(discoverable2));
                // Changes should be received by the handler, with only 3 resources,
                // as 2 out of 5 should get moved to the new discoverable.
                assigned = assignmentQueue.poll(30, TimeUnit.SECONDS);
                Assert.assertNotNull(assigned);
                Assert.assertEquals(3, assigned.size());
                // Cancel the first discoverable again, should expect empty result.
                // This also make sure the latest assignment get cached in the ResourceCoordinatorClient.
                // It is the the next test step.
                cancelDiscoverable1.cancel();
                Assert.assertTrue(assignmentQueue.poll(30, TimeUnit.SECONDS).isEmpty());
                // Cancel the handler.
                cancelSubscribe1.cancel();
                Assert.assertTrue(finishSemaphore.tryAcquire(2, TimeUnit.SECONDS));
                // Subscribe to changes for the second discoverable,
                // it should see the latest assignment, even though no new fetch from ZK is triggered.
                Cancellable cancelSubscribe2 = subscribe(client, discoverable2, assignmentQueue, finishSemaphore);
                assigned = assignmentQueue.poll(30, TimeUnit.SECONDS);
                Assert.assertNotNull(assigned);
                Assert.assertEquals(5, assigned.size());
                // Update the requirement to be an empty requirement, the handler should receive an empty collection
                client.submitRequirement(ResourceRequirement.builder(serviceName).build());
                Assert.assertTrue(assignmentQueue.poll(30, TimeUnit.SECONDS).isEmpty());
                // Update the requirement to have one partition, the handler should receive one resource
                client.submitRequirement(ResourceRequirement.builder(serviceName).addPartitions("p", 1, 1).build());
                assigned = assignmentQueue.poll(30, TimeUnit.SECONDS);
                Assert.assertNotNull(assigned);
                Assert.assertEquals(1, assigned.size());
                // Delete the requirement, the handler should receive a empty collection
                client.deleteRequirement(requirement.getName());
                Assert.assertTrue(assignmentQueue.poll(30, TimeUnit.SECONDS).isEmpty());
                // Cancel the second handler.
                cancelSubscribe2.cancel();
                Assert.assertTrue(finishSemaphore.tryAcquire(2, TimeUnit.SECONDS));
                cancelDiscoverable2.cancel();
            } finally {
                client.stopAndWait();
            }
        } finally {
            coordinator.stopAndWait();
        }
    } finally {
        zkClient.stopAndWait();
    }
}

Example 10 with ZKClientService

use of org.apache.twill.zookeeper.ZKClientService in project cdap by caskdata.

the class LeaderElectionMessagingServiceTest method testTransition.

@Test
public void testTransition() throws Throwable {
    final TopicId topicId = NamespaceId.SYSTEM.topic("topic");
    Injector injector1 = createInjector(0);
    Injector injector2 = createInjector(1);
    // Start a messaging service, which would becomes leader
    ZKClientService zkClient1 = injector1.getInstance(ZKClientService.class);
    zkClient1.startAndWait();
    final MessagingService firstService = injector1.getInstance(MessagingService.class);
    if (firstService instanceof Service) {
        ((Service) firstService).startAndWait();
    }
    // Publish a message with the leader
    firstService.publish(StoreRequestBuilder.of(topicId).addPayloads("Testing1").build());
    // Start another messaging service, this one would be follower
    ZKClientService zkClient2 = injector2.getInstance(ZKClientService.class);
    zkClient2.startAndWait();
    final MessagingService secondService = injector2.getInstance(MessagingService.class);
    if (secondService instanceof Service) {
        ((Service) secondService).startAndWait();
    }
    // Try to call the follower, should get service unavailable.
    try {
        secondService.listTopics(NamespaceId.SYSTEM);
        Assert.fail("Expected service unavailable");
    } catch (ServiceUnavailableException e) {
    // Expected
    }
    // Make the ZK session timeout for the leader service. The second one should pickup.
    KillZKSession.kill(zkClient1.getZooKeeperSupplier().get(), zkClient1.getConnectString(), 10000);
    // Publish one more message and then fetch from the current leader
    List<String> messages = Retries.callWithRetries(new Retries.Callable<List<String>, Throwable>() {

        @Override
        public List<String> call() throws Throwable {
            secondService.publish(StoreRequestBuilder.of(topicId).addPayloads("Testing2").build());
            List<String> messages = new ArrayList<>();
            try (CloseableIterator<RawMessage> iterator = secondService.prepareFetch(topicId).fetch()) {
                while (iterator.hasNext()) {
                    messages.add(new String(iterator.next().getPayload(), "UTF-8"));
                }
            }
            return messages;
        }
    }, RetryStrategies.timeLimit(10, TimeUnit.SECONDS, RetryStrategies.fixDelay(1, TimeUnit.SECONDS)));
    Assert.assertEquals(Arrays.asList("Testing1", "Testing2"), messages);
    // Shutdown the current leader. The session timeout one should becomes leader again.
    if (secondService instanceof Service) {
        ((Service) secondService).stopAndWait();
    }
    // Try to fetch message from the current leader again.
    // Should see two messages (because the cache is cleared and fetch is from the backing store).
    messages = Retries.callWithRetries(new Retries.Callable<List<String>, Throwable>() {

        @Override
        public List<String> call() throws Throwable {
            List<String> messages = new ArrayList<>();
            try (CloseableIterator<RawMessage> iterator = firstService.prepareFetch(topicId).fetch()) {
                while (iterator.hasNext()) {
                    messages.add(new String(iterator.next().getPayload(), "UTF-8"));
                }
            }
            return messages;
        }
    }, RetryStrategies.timeLimit(10, TimeUnit.SECONDS, RetryStrategies.fixDelay(1, TimeUnit.SECONDS)));
    Assert.assertEquals(Arrays.asList("Testing1", "Testing2"), messages);
    zkClient1.stopAndWait();
    zkClient2.stopAndWait();
}