OwnershipCache implements TimerTask { public void removeAllStreamsFromOwner(SocketAddress addr) { logger.info("Remove streams mapping for host {}", addr); Set<String> streamsForHost = address2Streams.get(addr); if (null != streamsForHost) { synchronized (streamsForHost) { for (String s : streamsForHost) { if (stream2Addresses.remove(s, addr)) { logger.info("Removing mapping for stream : {} from host : {}", s, addr); ownershipStatsLogger.onRemove(s); } } address2Streams.remove(addr, streamsForHost); } } } OwnershipCache(ClientConfig clientConfig,
                          HashedWheelTimer timer,
                          StatsReceiver statsReceiver,
                          StatsReceiver streamStatsReceiver); @Override void run(Timeout timeout); OwnershipStatsLogger getOwnershipStatsLogger(); boolean updateOwner(String stream, SocketAddress addr); SocketAddress getOwner(String stream); void removeOwnerFromStream(String stream, SocketAddress addr, String reason); void removeAllStreamsFromOwner(SocketAddress addr); int getNumCachedStreams(); Map<SocketAddress, Set<String>> getStreamOwnershipDistribution(); Map<String, SocketAddress> getStreamOwnerMapping();  }
@Test(timeout = 60000) public void testRemoveAllStreamsFromOwner() { OwnershipCache cache = createOwnershipCache(); int initialPort = 2000; int numProxies = 2; int numStreamsPerProxy = 2; for (int i = 0; i < numProxies; i++) { SocketAddress addr = createSocketAddress(initialPort + i); for (int j = 0; j < numStreamsPerProxy; j++) { String stream = runtime.getMethodName() + "_" + i + "_" + j; cache.updateOwner(stream, addr); } } Map<String, SocketAddress> ownershipMap = cache.getStreamOwnerMapping(); assertEquals("There should be " + (numProxies * numStreamsPerProxy) + " entries in cache", numProxies * numStreamsPerProxy, ownershipMap.size()); Map<SocketAddress, Set<String>> ownershipDistribution = cache.getStreamOwnershipDistribution(); assertEquals("There should be " + numProxies + " proxies cached", numProxies, ownershipDistribution.size()); SocketAddress owner = createSocketAddress(initialPort); SocketAddress nonExistentAddr = createSocketAddress(initialPort + 999); cache.removeAllStreamsFromOwner(nonExistentAddr); ownershipMap = cache.getStreamOwnerMapping(); assertEquals("There should still be " + (numProxies * numStreamsPerProxy) + " entries in cache", numProxies * numStreamsPerProxy, ownershipMap.size()); ownershipDistribution = cache.getStreamOwnershipDistribution(); assertEquals("There should still be " + numProxies + " proxies cached", numProxies, ownershipDistribution.size()); cache.removeAllStreamsFromOwner(owner); ownershipMap = cache.getStreamOwnerMapping(); assertEquals("There should be " + ((numProxies - 1) * numStreamsPerProxy) + " entries left in cache", (numProxies - 1) * numStreamsPerProxy, ownershipMap.size()); ownershipDistribution = cache.getStreamOwnershipDistribution(); assertEquals("There should be " + (numProxies - 1) + " proxies cached", numProxies - 1, ownershipDistribution.size()); assertFalse("Host " + owner + " should not be cached", ownershipDistribution.containsKey(owner)); }