// Start first IPC service w/ counter
final ConcurrentMap<String, AtomicInteger> firstCounts =
new ConcurrentHashMap<String, AtomicInteger>();
final HelixIPCService firstIPC =
new NettyHelixIPCService(new NettyHelixIPCService.Config().setInstanceName(
firstNode.getInstanceName()).setPort(firstPort));
firstIPC.registerCallback(messageType, new HelixIPCCallback() {
@Override
public void onMessage(HelixMessageScope scope, UUID messageId, ByteBuf message) {
String key = scope.getPartition() + ":" + scope.getState();
firstCounts.putIfAbsent(key, new AtomicInteger());
firstCounts.get(key).incrementAndGet();
}
});
firstIPC.start();
// Start second IPC Service w/ counter
final ConcurrentMap<String, AtomicInteger> secondCounts =
new ConcurrentHashMap<String, AtomicInteger>();
final HelixIPCService secondIPC =
new NettyHelixIPCService(new NettyHelixIPCService.Config().setInstanceName(
secondNode.getInstanceName()).setPort(secondPort));
secondIPC.registerCallback(messageType, new HelixIPCCallback() {
@Override
public void onMessage(HelixMessageScope scope, UUID messageId, ByteBuf message) {
String key = scope.getPartition() + ":" + scope.getState();
secondCounts.putIfAbsent(key, new AtomicInteger());
secondCounts.get(key).incrementAndGet();
}
});
secondIPC.start();
// Allow resolver callbacks to fire
Thread.sleep(500);
// Find all partitions on second node...
String secondName = "localhost_" + secondPort;
Set<String> secondPartitions = new HashSet<String>();
IdealState idealState =
controller.getClusterManagmentTool().getResourceIdealState(CLUSTER_NAME, RESOURCE_NAME);
for (String partitionName : idealState.getPartitionSet()) {
for (Map.Entry<String, String> stateEntry : idealState.getInstanceStateMap(partitionName)
.entrySet()) {
if (stateEntry.getKey().equals(secondName)) {
secondPartitions.add(partitionName);
}
}
}
// And use first node to send messages to them
for (String partitionName : secondPartitions) {
for (int i = 0; i < numMessages; i++) {
HelixMessageScope scope =
new HelixMessageScope.Builder().cluster(firstNode.getClusterName())
.resource(RESOURCE_NAME).partition(partitionName).state("ONLINE").build();
Set<HelixAddress> destinations = firstResolver.getDestinations(scope);
for (HelixAddress destination : destinations) {
ByteBuf message = Unpooled.wrappedBuffer(("Hello" + i).getBytes());
firstIPC.send(destination, messageType, UUID.randomUUID(), message);
}
}
}
// Loopback
for (String partitionName : secondPartitions) {
for (int i = 0; i < numMessages; i++) {
HelixMessageScope scope =
new HelixMessageScope.Builder().cluster(secondNode.getClusterName())
.resource(RESOURCE_NAME).partition(partitionName).state("ONLINE").build();
Set<HelixAddress> destinations = secondResolver.getDestinations(scope);
for (HelixAddress destination : destinations) {
ByteBuf message = Unpooled.wrappedBuffer(("Hello" + i).getBytes());
secondIPC.send(destination, messageType, UUID.randomUUID(), message);
}
}
}
// Check
Thread.sleep(500); // just in case
for (String partitionName : secondPartitions) {
AtomicInteger count = secondCounts.get(partitionName + ":ONLINE");
Assert.assertNotNull(count);
Assert.assertEquals(count.get(), 2 * numMessages);
}
// Shutdown
firstIPC.shutdown();
secondIPC.shutdown();
}