ElasticSourcePMetaSupplier implements ProcessorMetaSupplier { static Map<Address, List<Shard>> assignShards(Collection<Shard> shards, Collection<Address> addresses) { Map<String, List<Shard>> nodeCandidates = shards.stream() .collect(groupingBy(Shard::getIp)); Map<Address, List<Shard>> nodeAssigned = new HashMap<>(); if (!addresses.stream().map(Address::getHost).collect(toSet()) .containsAll(nodeCandidates.keySet())) { throw new JetException("Shard locations are not equal to Jet nodes locations, " + "shards=" + nodeCandidates.keySet() + ", Jet nodes=" + addresses); } int uniqueShards = (int) shards.stream().map(Shard::indexShard).distinct().count(); Set<String> assignedShards = new HashSet<>(); int candidatesSize = nodeCandidates.size(); int iterations = (uniqueShards + candidatesSize - 1) / candidatesSize; for (int i = 0; i < iterations; i++) { for (Address address : addresses) { String host = address.getHost(); List<Shard> thisNodeCandidates = nodeCandidates.getOrDefault(host, emptyList()); if (thisNodeCandidates.isEmpty()) { continue; } Shard shard = thisNodeCandidates.remove(0); List<Shard> nodeShards = nodeAssigned.computeIfAbsent(address, (key) -> new ArrayList<>()); nodeShards.add(shard); nodeCandidates.values().forEach(candidates -> candidates.removeIf(next -> next.indexShard().equals(shard.indexShard()))); assignedShards.add(shard.indexShard()); } } if (assignedShards.size() != uniqueShards) { throw new JetException("Not all shards have been assigned"); } return nodeAssigned; } ElasticSourcePMetaSupplier(@Nonnull ElasticSourceConfiguration<T> configuration); @Override int preferredLocalParallelism(); @Override void init(@Nonnull Context context); @Nonnull @Override Function<? super Address, ? extends ProcessorSupplier> get(@Nonnull List<Address> addresses);  }
@Test public void given_singleNodeAddress_when_assignShards_then_shouldAssignAllShardsToSingleAddress() { List<Shard> shards = newArrayList( new Shard("elastic-index", 0, Prirep.p, 10, "STARTED", "10.0.0.1", "10.0.0.1:9200", "node1"), new Shard("elastic-index", 1, Prirep.p, 10, "STARTED", "10.0.0.1", "10.0.0.1:9200", "node2"), new Shard("elastic-index", 2, Prirep.p, 10, "STARTED", "10.0.0.1", "10.0.0.1:9200", "node3") ); List<Address> addresses = addresses("10.0.0.1"); Map<Address, List<Shard>> assignment = ElasticSourcePMetaSupplier.assignShards(shards, addresses); assertThat(assignment).contains( entry(address("10.0.0.1"), shards) ); }
@Test(expected = JetException.class) public void given_noMatchingNode_when_assignShards_thenThrowException() { List<Shard> shards = newArrayList( new Shard("elastic-index", 0, Prirep.p, 10, "STARTED", "10.0.0.1", "10.0.0.1:9200", "node1") ); List<Address> addresses = addresses("10.0.0.2"); ElasticSourcePMetaSupplier.assignShards(shards, addresses); }