Source:
public static AbstractTopology shiftPartitionLeaders(AbstractTopology topology, Set<Integer> missingHosts) { if (missingHosts == null || missingHosts.isEmpty()) { return topology; } Map<Integer, MutableHost> mutableHostMap = new TreeMap<>(); Map<Integer, MutablePartition> mutablePartitionMap = new TreeMap<>(); for (Host host : topology.hostsById.values()) { final MutableHost mutableHost = new MutableHost(host.id, host.targetSiteCount, host.haGroup); if (missingHosts.contains(host.id)) { mutableHost.markHostMissing(true); } mutableHostMap.put(host.id, mutableHost); } for (Partition partition : topology.partitionsById.values()) { MutablePartition mp = new MutablePartition(partition.id, partition.k); mutablePartitionMap.put(mp.id, mp); for (Integer hId : partition.hostIds) { final MutableHost mutableHost = mutableHostMap.get(hId); mp.hosts.add(mutableHost); mutableHost.partitions.add(mp); } int leaderId = partition.leaderHostId; if (missingHosts.contains(leaderId)) { List<Host> perspectiveHosts = Lists.newArrayList(); for (Host host : topology.hostsById.values()) { if (!missingHosts.contains(host.id)) { List<Integer> list = topology.getPartitionIdList(host.id); if (list.contains(partition.id)) { perspectiveHosts.add(host); } } } assert(!perspectiveHosts.isEmpty()); perspectiveHosts.sort((Host a, Host b) -> { int result = 0; for (Partition p : topology.partitionsById.values()) { if ( a.id == p.leaderHostId) { result++; } else if ( b.id == p.leaderHostId) { result--; } } return result; }); leaderId = perspectiveHosts.get(0).id; partition.leaderHostId = leaderId; } mp.leader = mutableHostMap.get(leaderId); } return convertMutablesToTopology(topology.version, mutableHostMap, mutablePartitionMap); }

Target:
public static AbstractTopology shiftPartitionLeaders(AbstractTopology topology, Set<Integer> missingHosts) { if (missingHosts == null || missingHosts.isEmpty()) { return topology; } Map<Integer, MutableHost> mutableHostMap = new TreeMap<>(); Map<Integer, MutablePartition> mutablePartitionMap = new TreeMap<>(); for (Host host : topology.hostsById.values()) { final MutableHost mutableHost = new MutableHost(host.id, host.targetSiteCount, host.haGroup); if (missingHosts.contains(host.id)) { mutableHost.markHostMissing(true); } mutableHostMap.put(host.id, mutableHost); } for (Partition partition : topology.partitionsById.values()) { MutablePartition mp = new MutablePartition(partition.id, partition.k); mutablePartitionMap.put(mp.id, mp); for (Integer hId : partition.hostIds) { final MutableHost mutableHost = mutableHostMap.get(hId); mp.hosts.add(mutableHost); mutableHost.partitions.add(mp); } int leaderId = partition.leaderHostId; if (missingHosts.contains(leaderId)) { List<Host> perspectiveHosts = Lists.newArrayList(); for (Host host : topology.hostsById.values()) { if (!missingHosts.contains(host.id)) { List<Integer> partitionListOnNonMissingHost = topology.getPartitionIdList(host.id); if (partitionListOnNonMissingHost.contains(partition.id)) { perspectiveHosts.add(host); } } } assert(!perspectiveHosts.isEmpty()); if (perspectiveHosts.size() > 1) { perspectiveHosts.sort((Host a, Host b) -> { return (a.getleaderCount() - b.getleaderCount()); }); } leaderId = perspectiveHosts.get(0).id; partition.leaderHostId = leaderId; } mp.leader = mutableHostMap.get(leaderId); } return convertMutablesToTopology(topology.version, mutableHostMap, mutablePartitionMap); }