Repository: apache/kafka
File: tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java (Java)

Code before:
        return results;
    }

    static Map<TopicPartitionReplica, String> getReplicaToLogDir(
        Admin adminClient,
        Map<TopicPartition, List<Integer>> topicPartitionToReplicas
    ) {
        var replicaLogDirs = topicPartitionToReplicas
                .entrySet()
                .stream()
                .flatMap(entry -> entry.getValue()
                    .stream()
                    .map(id -> new TopicPartitionReplica(entry.getKey().topic(), entry.getKey().partition(), id)))
                .collect(Collectors.toUnmodifiableSet());

        var futureMap = adminClient.describeReplicaLogDirs(replicaLogDirs).values();

        return futureMap.entrySet()
                .stream()
                .collect(Collectors.toMap(
                        Entry::getKey,
                        entry -> {
                            try {
                                var logDir = entry.getValue().get().getCurrentReplicaLogDir();
                                return logDir != null ? logDir : "any";
                            } catch (ExecutionException e) {
                                if (e.getCause() instanceof TimeoutException) {
                                    System.out.println("Failed to get log dir for " + entry.getKey() + ": " + e.getMessage());
                                    return "any";
                                } else {
                                    throw new AdminCommandFailedException("Failed to describe replica log dirs", e.getCause());
                                }
                            } catch (InterruptedException e) {
                                throw new AdminCommandFailedException("Failed to describe replica log dirs", e.getCause());
                            }
                        }
                ));
    }
}


Review comment (suggestion):
I'm not a fan of this solution because it leads to timeouts, which we want to avoid. Perhaps `getReplicaAssignmentForPartitions` could return `Node` instead of `Integer`? Then we can reuse the `Node#isEmpty` check to skip fetching directories from offline nodes

Code after:
        });
        return opts;
    }

    static Set<TopicPartitionReplica> alterReplicaLogDirs(Admin adminClient,
                                                          Map<TopicPartitionReplica, String> assignment) throws InterruptedException {
        Set<TopicPartitionReplica> results = new HashSet<>();
        Map<TopicPartitionReplica, KafkaFuture<Void>> values = adminClient.alterReplicaLogDirs(assignment).values();

        for (Entry<TopicPartitionReplica, KafkaFuture<Void>> e : values.entrySet()) {
            TopicPartitionReplica replica = e.getKey();
            KafkaFuture<Void> future = e.getValue();
            try {
                future.get();
                results.add(replica);
            } catch (ExecutionException t) {
                // Ignore ReplicaNotAvailableException.  It is OK if the replica is not
                // available at this moment.
                if (t.getCause() instanceof ReplicaNotAvailableException)
                    continue;
                throw new AdminCommandFailedException("Failed to alter dir for " + replica, t);
            }
        }
        return results;
    }

    /**
     * Get the log directory for each replica.
     *
     * @param adminClient The admin client object to use.
     * @param current The current partition assignment with Node information.
     * @return Map of TopicPartitionReplica to log directory path.
     */
    static Map<TopicPartitionReplica, String> getReplicaToLogDir(
        Admin adminClient,
        Map<TopicPartition, List<Node>> current
    ) throws ExecutionException, InterruptedException {
        List<TopicPartitionReplica> availableReplicas = available(current);

        if (availableReplicas.isEmpty()) {
            return Map.of();
        }

        return adminClient.describeReplicaLogDirs(availableReplicas).all().get()
                .entrySet()
                .stream()
                .filter(e -> e.getValue().getCurrentReplicaLogDir() != null)
                .collect(Collectors.toMap(
                    Entry::getKey,
                    e -> e.getValue().getCurrentReplicaLogDir())
                );