Source:
<code>public static ImmutableWorkerInfo selectWorker( final Task task, final Map<String, ImmutableWorkerInfo> allWorkers, final WorkerTaskRunnerConfig workerTaskRunnerConfig, @Nullable final AffinityConfig affinityConfig, final Function<ImmutableMap<String, ImmutableWorkerInfo>, ImmutableWorkerInfo> workerSelector ) { final Map<String, ImmutableWorkerInfo> runnableWorkers = allWorkers .values() .stream() .filter(worker -> worker.canRunTask(task) && worker.isValidVersion(workerTaskRunnerConfig.getMinWorkerVersion())) .collect(Collectors.toMap(w -> w.getWorker().getHost(), Function.identity())); if (affinityConfig == null) { return workerSelector.apply(ImmutableMap.copyOf(runnableWorkers)); } else { final Map<String, ImmutableWorkerInfo> nonAffinityWorkers = runnableWorkers .entrySet() .stream() .filter(entry -> !affinityConfig.getAffinityWorkers().contains(entry.getKey())) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); final List<String> dataSourceWorkers = affinityConfig.getAffinity().get(task.getDataSource()); if (dataSourceWorkers == null) { return workerSelector.apply(ImmutableMap.copyOf(nonAffinityWorkers)); } else { <START> final Set<String> dataSourceWorkerHosts = ImmutableSet.copyOf(dataSourceWorkers); <END> final Map<String, ImmutableWorkerInfo> dataSourceWorkerMap = runnableWorkers .entrySet() .stream() .filter(entry -> dataSourceWorkerHosts.contains(entry.getKey())) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); final ImmutableWorkerInfo selected = workerSelector.apply(ImmutableMap.copyOf(dataSourceWorkerMap)); if (selected != null) { return selected; } else if (affinityConfig.isStrong()) { return null; } else { return workerSelector.apply(ImmutableMap.copyOf(nonAffinityWorkers)); } } } }</code><technical_language>copy if AffinityConfig.affinity is a Map<String, Set<String>></technical_language>

Target:
public static ImmutableWorkerInfo selectWorker( final Task task, final Map<String, ImmutableWorkerInfo> allWorkers, final WorkerTaskRunnerConfig workerTaskRunnerConfig, @Nullable final AffinityConfig affinityConfig, final Function<ImmutableMap<String, ImmutableWorkerInfo>, ImmutableWorkerInfo> workerSelector ) { final Map<String, ImmutableWorkerInfo> runnableWorkers = allWorkers .values() .stream() .filter(worker -> worker.canRunTask(task) && worker.isValidVersion(workerTaskRunnerConfig.getMinWorkerVersion())) .collect(Collectors.toMap(w -> w.getWorker().getHost(), Function.identity())); if (affinityConfig == null) { return workerSelector.apply(ImmutableMap.copyOf(runnableWorkers)); } else { final Set<String> dataSourceWorkers = affinityConfig.getAffinity().get(task.getDataSource()); if (dataSourceWorkers == null) { return workerSelector.apply(getNonAffinityWorkers(affinityConfig, runnableWorkers)); } else { final ImmutableMap<String, ImmutableWorkerInfo> dataSourceWorkerMap = ImmutableMap.copyOf(Maps.filterKeys(runnableWorkers, dataSourceWorkers::contains)); final ImmutableWorkerInfo selected = workerSelector.apply(dataSourceWorkerMap); if (selected != null) { return selected; } else if (affinityConfig.isStrong()) { return null; } else { return workerSelector.apply(getNonAffinityWorkers(affinityConfig, runnableWorkers)); } } } }