OwnershipCache implements TimerTask { public void removeOwnerFromStream(String stream, SocketAddress addr, String reason) { if (stream2Addresses.remove(stream, addr)) { logger.info("Removed stream to host mapping for (stream: {} -> host: {}) : reason = '{}'.", new Object[] { stream, addr, reason }); } Set<String> streamsForHost = address2Streams.get(addr); if (null != streamsForHost) { synchronized (streamsForHost) { if (streamsForHost.remove(stream)) { logger.info("Removed stream ({}) from host {} : reason = '{}'.", new Object[] { stream, addr, reason }); if (streamsForHost.isEmpty()) { address2Streams.remove(addr, streamsForHost); } ownershipStatsLogger.onRemove(stream); } } } } 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 testRemoveOwnerFromStream() { 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()); String stream = runtime.getMethodName() + "_0_0"; SocketAddress owner = createSocketAddress(initialPort); SocketAddress nonExistentAddr = createSocketAddress(initialPort + 999); cache.removeOwnerFromStream(stream, nonExistentAddr, "remove-non-existent-addr"); assertEquals("Owner " + owner + " should not be removed", owner, cache.getOwner(stream)); ownershipMap = cache.getStreamOwnerMapping(); assertEquals("There should be " + (numProxies * numStreamsPerProxy) + " entries in cache", numProxies * numStreamsPerProxy, ownershipMap.size()); cache.removeOwnerFromStream(stream, owner, "remove-owner"); assertNull("Owner " + owner + " should be removed", cache.getOwner(stream)); ownershipMap = cache.getStreamOwnerMapping(); assertEquals("There should be " + (numProxies * numStreamsPerProxy - 1) + " entries left in cache", numProxies * numStreamsPerProxy - 1, ownershipMap.size()); ownershipDistribution = cache.getStreamOwnershipDistribution(); assertEquals("There should still be " + numProxies + " proxies cached", numProxies, ownershipDistribution.size()); Set<String> ownedStreams = ownershipDistribution.get(owner); assertEquals("There should be only " + (numStreamsPerProxy - 1) + " streams owned for " + owner, numStreamsPerProxy - 1, ownedStreams.size()); assertFalse("Stream " + stream + " should not be owned by " + owner, ownedStreams.contains(stream)); }