use of org.apache.twill.common.Cancellable 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();
}
}
use of org.apache.twill.common.Cancellable 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();
}
}
use of org.apache.twill.common.Cancellable in project cdap by caskdata.
the class LoggingContextAccessor method setLoggingContext.
/**
* Sets the logging context.
* <p>
* NOTE: in work execution frameworks where threads are shared between workers (like Akka) we would have to init
* context very frequently (before every chunk of work is started). In that case we really want to re-use
* logging context object instance.
* </p>
* @param context context to set
* @return Cancellable that can be used to revert the logging context and MDC Map to its original value
*/
public static Cancellable setLoggingContext(LoggingContext context) {
final LoggingContext saveLoggingContext = loggingContext.get();
final Map saveContextMap = MDC.getCopyOfContextMap();
final Thread saveCurrentThread = Thread.currentThread();
loggingContext.set(context);
try {
// Also copy context tags to MDC
MDC.setContextMap(context.getSystemTagsAsString());
} catch (IllegalStateException e) {
// MDC will throw this if there is no valid binding. Normally this shouldn't happen, but in case it does,
// we'll just ignore it as it doesn't affect platform logic at all as we always use loggingContext.
}
return new Cancellable() {
private boolean cancelled;
@Override
public void cancel() {
if (Thread.currentThread() == saveCurrentThread && !cancelled) {
MDC.setContextMap(saveContextMap == null ? Collections.emptyMap() : saveContextMap);
loggingContext.set(saveLoggingContext);
cancelled = true;
}
}
};
}
use of org.apache.twill.common.Cancellable in project cdap by caskdata.
the class DistributedProgramRunner method logProgramStart.
/**
* Set up Logging Context so the Log is tagged correctly for the Program.
* Reset the context once done.
*/
private void logProgramStart(Program program, ProgramOptions options) {
LoggingContext loggingContext = LoggingContextHelper.getLoggingContext(program.getNamespaceId(), program.getApplicationId(), program.getName(), program.getType(), ProgramRunners.getRunId(options).getId(), options.getArguments().asMap());
Cancellable saveContextCancellable = LoggingContextAccessor.setLoggingContext(loggingContext);
String userArguments = Joiner.on(", ").withKeyValueSeparator("=").join(options.getUserArguments());
LOG.info("Starting {} Program '{}' with Arguments [{}]", program.getType(), program.getName(), userArguments);
saveContextCancellable.cancel();
}
use of org.apache.twill.common.Cancellable in project cdap by caskdata.
the class AbstractEndpointStrategy method pick.
@Override
public Discoverable pick(long timeout, TimeUnit timeoutUnit) {
Discoverable discoverable = pick();
if (discoverable != null) {
return discoverable;
}
final SettableFuture<Discoverable> future = SettableFuture.create();
Cancellable cancellable = serviceDiscovered.watchChanges(new ServiceDiscovered.ChangeListener() {
@Override
public void onChange(ServiceDiscovered serviceDiscovered) {
// The serviceDiscovered provided is the same as the one in the field, hence ok to just call pick().
Discoverable discoverable = pick();
if (discoverable != null) {
future.set(discoverable);
}
}
}, Threads.SAME_THREAD_EXECUTOR);
try {
return future.get(timeout, timeoutUnit);
} catch (Exception e) {
return null;
} finally {
cancellable.cancel();
}
}
Aggregations