Set<LocatableInputSplit> splits = new HashSet<LocatableInputSplit>();
for (int i = 0; i < NUM_SPLITS; i++) {
splits.add(new LocatableInputSplit(i, hosts[i%hosts.length]));
}
final LocatableInputSplitAssigner ia = new LocatableInputSplitAssigner(splits);
final AtomicInteger splitsRetrieved = new AtomicInteger(0);
final AtomicInteger sumOfIds = new AtomicInteger(0);
Runnable retriever = new Runnable() {
@Override
public void run() {
final String threadHost = hosts[(int) (Math.random() * hosts.length)];
LocatableInputSplit split;
while ((split = ia.getNextInputSplit(threadHost)) != null) {
splitsRetrieved.incrementAndGet();
sumOfIds.addAndGet(split.getSplitNumber());
}
}
};
// create the threads
Thread[] threads = new Thread[NUM_THREADS];
for (int i = 0; i < NUM_THREADS; i++) {
threads[i] = new Thread(retriever);
threads[i].setDaemon(true);
}
// launch concurrently
for (int i = 0; i < NUM_THREADS; i++) {
threads[i].start();
}
// sync
for (int i = 0; i < NUM_THREADS; i++) {
threads[i].join(5000);
}
// verify
for (int i = 0; i < NUM_THREADS; i++) {
if (threads[i].isAlive()) {
fail("The concurrency test case is erroneous, the thread did not respond in time.");
}
}
assertEquals(NUM_SPLITS, splitsRetrieved.get());
assertEquals(SUM_OF_IDS, sumOfIds.get());
// nothing left
assertNull(ia.getNextInputSplit("testhost"));
// at least one fraction of hosts needs be local, no matter how bad the thread races
assertTrue(ia.getNumberOfLocalAssignments() >= NUM_SPLITS / hosts.length);
}
catch (Exception e) {
e.printStackTrace();
fail(e.getMessage());
}