List<Address> members1 = Arrays.asList(addresses[0], addresses[1], addresses[2], addresses[3]);
List<Address> members2 = Arrays.asList(addresses[0], addresses[1], addresses[2]);
// create CHes
DefaultConsistentHashFactory chf = new DefaultConsistentHashFactory();
DefaultConsistentHash ch1 = chf.create(new MurmurHash3(), 2, 40, members1, null);
final DefaultConsistentHash ch2 = chf.updateMembers(ch1, members2, null);
DefaultConsistentHash ch3 = chf.rebalance(ch2);
log.debug(ch1);
log.debug(ch2);
// create dependencies
Cache cache = mock(Cache.class);
when(cache.getName()).thenReturn("testCache");
ThreadFactory threadFactory = new ThreadFactory() {
@Override
public Thread newThread(Runnable r) {
String name = "PooledExecutorThread-" + StateConsumerTest.class.getSimpleName() + "-" + r.hashCode();
return new Thread(r, name);
}
};
pooledExecutorService = new ThreadPoolExecutor(10, 20, 0L, TimeUnit.MILLISECONDS,
new LinkedBlockingDeque<Runnable>(), threadFactory, new ThreadPoolExecutor.CallerRunsPolicy());
StateTransferManager stateTransferManager = mock(StateTransferManager.class);
CacheNotifier cacheNotifier = mock(CacheNotifier.class);
RpcManager rpcManager = mock(RpcManager.class);
Transport transport = mock(Transport.class);
CommandsFactory commandsFactory = mock(CommandsFactory.class);
PersistenceManager persistenceManager = mock(PersistenceManager.class);
DataContainer dataContainer = mock(DataContainer.class);
TransactionTable transactionTable = mock(TransactionTable.class);
StateTransferLock stateTransferLock = mock(StateTransferLock.class);
InterceptorChain interceptorChain = mock(InterceptorChain.class);
InvocationContextContainer icc = mock(InvocationContextContainer.class);
TotalOrderManager totalOrderManager = mock(TotalOrderManager.class);
BlockingTaskAwareExecutorService remoteCommandsExecutor = mock(BlockingTaskAwareExecutorService.class);
L1Manager l1Manager = mock(L1Manager.class);
when(commandsFactory.buildStateRequestCommand(any(StateRequestCommand.Type.class), any(Address.class), anyInt(), any(Set.class))).thenAnswer(new Answer<StateRequestCommand>() {
@Override
public StateRequestCommand answer(InvocationOnMock invocation) {
return new StateRequestCommand("cache1", (StateRequestCommand.Type) invocation.getArguments()[0], (Address) invocation.getArguments()[1], (Integer) invocation.getArguments()[2], (Set) invocation.getArguments()[3]);
}
});
when(transport.getViewId()).thenReturn(1);
when(rpcManager.getAddress()).thenReturn(addresses[0]);
when(rpcManager.getTransport()).thenReturn(transport);
final Map<Address, Set<Integer>> requestedSegments = CollectionFactory.makeConcurrentMap();
final Set<Integer> flatRequestedSegments = new ConcurrentSkipListSet<Integer>();
when(rpcManager.invokeRemotely(any(Collection.class), any(StateRequestCommand.class), any(RpcOptions.class)))
.thenAnswer(new Answer<Map<Address, Response>>() {
@Override
public Map<Address, Response> answer(InvocationOnMock invocation) {
Collection<Address> recipients = (Collection<Address>) invocation.getArguments()[0];
Address recipient = recipients.iterator().next();
StateRequestCommand cmd = (StateRequestCommand) invocation.getArguments()[1];
Map<Address, Response> results = new HashMap<Address, Response>(1);
if (cmd.getType().equals(StateRequestCommand.Type.GET_TRANSACTIONS)) {
results.put(recipient, SuccessfulResponse.create(new ArrayList<TransactionInfo>()));
Set<Integer> segments = (Set<Integer>) cmd.getParameters()[3];
requestedSegments.put(recipient, segments);
flatRequestedSegments.addAll(segments);
} else if (cmd.getType().equals(StateRequestCommand.Type.START_STATE_TRANSFER)
|| cmd.getType().equals(StateRequestCommand.Type.CANCEL_STATE_TRANSFER)) {
results.put(recipient, SuccessfulResponse.SUCCESSFUL_EMPTY_RESPONSE);
}
return results;
}
});
when(rpcManager.getRpcOptionsBuilder(any(ResponseMode.class))).thenAnswer(new Answer<RpcOptionsBuilder>() {
public RpcOptionsBuilder answer(InvocationOnMock invocation) {
Object[] args = invocation.getArguments();
return new RpcOptionsBuilder(10000, TimeUnit.MILLISECONDS, (ResponseMode) args[0], true);
}
});
// create state provider
final StateConsumerImpl stateConsumer = new StateConsumerImpl();
stateConsumer.init(cache, pooledExecutorService, stateTransferManager, interceptorChain, icc, configuration, rpcManager, null,
commandsFactory, persistenceManager, dataContainer, transactionTable, stateTransferLock, cacheNotifier,
totalOrderManager, remoteCommandsExecutor, l1Manager);
stateConsumer.start();
final List<InternalCacheEntry> cacheEntries = new ArrayList<InternalCacheEntry>();
Object key1 = new TestKey("key1", 0, ch1);
Object key2 = new TestKey("key2", 0, ch1);
cacheEntries.add(new ImmortalCacheEntry(key1, "value1"));
cacheEntries.add(new ImmortalCacheEntry(key2, "value2"));
when(dataContainer.iterator()).thenAnswer(new Answer<Iterator<InternalCacheEntry>>() {
@Override
public Iterator<InternalCacheEntry> answer(InvocationOnMock invocation) {
return cacheEntries.iterator();
}
});
when(transactionTable.getLocalTransactions()).thenReturn(Collections.<LocalTransaction>emptyList());
when(transactionTable.getRemoteTransactions()).thenReturn(Collections.<RemoteTransaction>emptyList());
assertFalse(stateConsumer.hasActiveTransfers());
// node 481 leaves
stateConsumer.onTopologyUpdate(new CacheTopology(1, ch2, null), false);
assertFalse(stateConsumer.hasActiveTransfers());
// start a rebalance
stateConsumer.onTopologyUpdate(new CacheTopology(2, ch2, ch3), true);
assertTrue(stateConsumer.hasActiveTransfers());
// check that all segments have been requested
Set<Integer> oldSegments = ch2.getSegmentsForOwner(addresses[0]);
final Set<Integer> newSegments = ch3.getSegmentsForOwner(addresses[0]);
newSegments.removeAll(oldSegments);
log.debugf("Rebalancing. Added segments=%s, old segments=%s", newSegments, oldSegments);
assertEquals(flatRequestedSegments, newSegments);
// simulate a cluster state recovery and return to ch2