ProxyClientManager implements TimerTask { private void handshake(SocketAddress address, ProxyClient sc, FutureEventListener<ServerInfo> listener, boolean logging, boolean getOwnerships) { if (clientConfig.getHandshakeWithClientInfo()) { ClientInfo clientInfo = new ClientInfo(); clientInfo.setGetOwnerships(getOwnerships); clientInfo.setStreamNameRegex(clientConfig.getStreamNameRegex()); if (logging) { logger.info("Handshaking with {} : {}", address, clientInfo); } sc.getService().handshakeWithClientInfo(clientInfo) .addEventListener(listener); } else { if (logging) { logger.info("Handshaking with {}", address); } sc.getService().handshake().addEventListener(listener); } } ProxyClientManager(ClientConfig clientConfig,
                              ProxyClient.Builder clientBuilder,
                              HashedWheelTimer timer,
                              HostProvider hostProvider,
                              ClientStats clientStats); @Override void run(Timeout timeout); void registerProxyListener(ProxyListener listener); ProxyClient getClient(final SocketAddress address); void removeClient(SocketAddress address); void removeClient(SocketAddress address, ProxyClient sc); ProxyClient createClient(final SocketAddress address); void handshake(); int getNumProxies(); Map<SocketAddress, ProxyClient> getAllClients(); void close();  }
@Test(timeout = 60000) public void testHandshake() throws Exception { final int numHosts = 3; final int numStreamsPerHost = 3; final int initialPort = 4000; MockProxyClientBuilder builder = new MockProxyClientBuilder(); Map<SocketAddress, ServerInfo> serverInfoMap = new HashMap<SocketAddress, ServerInfo>(); for (int i = 0; i < numHosts; i++) { SocketAddress address = createSocketAddress(initialPort + i); ServerInfo serverInfo = new ServerInfo(); for (int j = 0; j < numStreamsPerHost; j++) { serverInfo.putToOwnerships(runtime.getMethodName() + "_stream_" + j, address.toString()); } Pair<MockProxyClient, MockServerInfoService> mockProxyClient = createMockProxyClient(address, serverInfo); builder.provideProxyClient(address, mockProxyClient.getLeft()); serverInfoMap.put(address, serverInfo); } final Map<SocketAddress, ServerInfo> results = new HashMap<SocketAddress, ServerInfo>(); final CountDownLatch doneLatch = new CountDownLatch(2 * numHosts); ProxyListener listener = new ProxyListener() { @Override public void onHandshakeSuccess(SocketAddress address, ProxyClient client, ServerInfo serverInfo) { synchronized (results) { results.put(address, serverInfo); } doneLatch.countDown(); } @Override public void onHandshakeFailure(SocketAddress address, ProxyClient client, Throwable cause) { } }; TestHostProvider rs = new TestHostProvider(); ProxyClientManager clientManager = createProxyClientManager(builder, rs, 0L); clientManager.registerProxyListener(listener); assertEquals("There should be no clients in the manager", 0, clientManager.getNumProxies()); for (int i = 0; i < numHosts; i++) { rs.addHost(createSocketAddress(initialPort + i)); } clientManager.handshake(); doneLatch.await(); assertEquals("Handshake should return server info", numHosts, results.size()); assertTrue("Handshake should get all server infos", Maps.difference(serverInfoMap, results).areEqual()); }