Repository: apache/kafka
File: server/src/main/java/org/apache/kafka/server/share/fetch/PartitionMaxBytesStrategy.java (Java)

Code before:
     * @param partitions - The topic partitions in the order for which we compute the partition max bytes.
     * @param acquiredPartitionsSize - The total partitions that have been acquired.
     * @return the partition max bytes for the topic partitions
     */
    LinkedHashMap<TopicIdPartition, Integer> maxBytes(int requestMaxBytes, Set<TopicIdPartition> partitions, int acquiredPartitionsSize);

    static PartitionMaxBytesStrategy type(StrategyType type) {
        if (type == null)
            throw new IllegalArgumentException("Strategy type cannot be null");
        return switch (type) {
            case UNIFORM -> PartitionMaxBytesStrategy::uniformPartitionMaxBytes;
        };
    }


    private static LinkedHashMap<TopicIdPartition, Integer> uniformPartitionMaxBytes(int requestMaxBytes, Set<TopicIdPartition> partitions, int acquiredPartitionsSize) {
        checkValidArguments(requestMaxBytes, partitions, acquiredPartitionsSize);
        LinkedHashMap<TopicIdPartition, Integer> partitionMaxBytes = new LinkedHashMap<>();
        if (requestMaxBytes >= acquiredPartitionsSize) {
            // Case 1: requestMaxBytes can be evenly distributed within partitions. If there is extra bytes left post
            // dividing it uniformly, assign it randomly to any one of the partitions.
            partitions.forEach(partition -> partitionMaxBytes.put(partition, requestMaxBytes / acquiredPartitionsSize));
            if (requestMaxBytes % acquiredPartitionsSize != 0) {
                TopicIdPartition randomPartition = selectPartitionRandomly(partitionMaxBytes);
                partitionMaxBytes.put(randomPartition,
                    (requestMaxBytes / acquiredPartitionsSize) + (requestMaxBytes % acquiredPartitionsSize));
            }
        } else if (requestMaxBytes >= partitions.size()) {
            // Case 2: we will be distributing requestMaxBytes greedily in this scenario to prevent any starvation. If
            // there is extra bytes left post dividing it uniformly, assign it randomly to any one of the partitions.
            partitions.forEach(partition -> partitionMaxBytes.put(partition, requestMaxBytes / partitions.size()));
            if (requestMaxBytes % partitions.size() != 0) {
                TopicIdPartition randomPartition = selectPartitionRandomly(partitionMaxBytes);
                partitionMaxBytes.put(randomPartition,
                    (requestMaxBytes / partitions.size()) + (requestMaxBytes % partitions.size()));
            }
        } else {
            // Case 3: we will distribute requestMaxBytes to as many partitions possible randomly to avoid starvation.
            List<TopicIdPartition> partitionsList = new ArrayList<>(partitions);
            Collections.shuffle(partitionsList);
            Set<TopicIdPartition> nonEmptyPartitions = new HashSet<>(partitionsList.subList(0, requestMaxBytes));
            partitions.forEach(
                partition -> {
                    if (nonEmptyPartitions.contains(partition)) {
                        partitionMaxBytes.put(partition, 1);
                    } else {
                        partitionMaxBytes.put(partition, 0);
                    }
                }
            );
        }

Review comment (suggestion):
nit: Though the calculation is trivial but not required to be repeated i.e. `requestMaxBytes / acquiredPartitionsSize` rather fetch the existing and add.

If you don't like it then ignore the suggestion.

```suggestion
            int remainingBytes = requestMaxBytes % acquiredPartitionsSize;
            if (remainingBytes != 0) {
                TopicIdPartition randomPartition = selectPartitionRandomly(partitionMaxBytes);
                partitionMaxBytes.put(randomPartition,
                    partitionMaxBytes.get(randomPartition) + remainingBytes);
```

Code after:
    LinkedHashMap<TopicIdPartition, Integer> maxBytes(int requestMaxBytes, Set<TopicIdPartition> partitions, int acquiredPartitionsSize);

    static PartitionMaxBytesStrategy type(StrategyType type) {
        if (type == null)
            throw new IllegalArgumentException("Strategy type cannot be null");
        return switch (type) {
            case UNIFORM -> PartitionMaxBytesStrategy::uniformPartitionMaxBytes;
        };
    }


    private static LinkedHashMap<TopicIdPartition, Integer> uniformPartitionMaxBytes(int requestMaxBytes, Set<TopicIdPartition> partitions, int acquiredPartitionsSize) {
        checkValidArguments(requestMaxBytes, partitions, acquiredPartitionsSize);
        if (requestMaxBytes >= acquiredPartitionsSize) {
            // Case 1: requestMaxBytes can be evenly distributed within partitions. If there is extra bytes left post
            // dividing it uniformly, assign it randomly to any one of the partitions.
            return allotUniformBytesToPartitions(partitions, requestMaxBytes, acquiredPartitionsSize);
        } else if (requestMaxBytes >= partitions.size()) {
            // Case 2: requestMaxBytes can be distributed greedily in this scenario to prevent any starvation. If
            // there is extra bytes left post dividing it uniformly, assign it randomly to any one of the partitions.
            return allotUniformBytesToPartitions(partitions, requestMaxBytes, partitions.size());
        } else {
            // Case 3: we will distribute requestMaxBytes to as many partitions possible randomly to avoid starvation.
            LinkedHashMap<TopicIdPartition, Integer> partitionMaxBytes = new LinkedHashMap<>();
            partitions.forEach(partition -> partitionMaxBytes.put(partition, 0));
            int randomIndex = RANDOM.nextInt(partitions.size());
            int nonEmptyPartitionsCount = 0;
            Set<Integer> nonEmptyPartitionIndexes = new HashSet<>();
            while (nonEmptyPartitionsCount < requestMaxBytes) {
                nonEmptyPartitionIndexes.add(randomIndex);
                nonEmptyPartitionsCount += 1;
                randomIndex = (randomIndex + 1) % partitions.size();
            }
            allotBytesByPartitionIndex(1, partitions, partitionMaxBytes, nonEmptyPartitionIndexes);
            return partitionMaxBytes;
        }
    }

    private static LinkedHashMap<TopicIdPartition, Integer> allotUniformBytesToPartitions(
        Set<TopicIdPartition> partitions,
        int requestMaxBytes,
        int partitionsSize
    ) {
        LinkedHashMap<TopicIdPartition, Integer> partitionMaxBytes = new LinkedHashMap<>();
        int uniformPartitionBytes = requestMaxBytes / partitionsSize;
        int remainingBytes = requestMaxBytes % partitionsSize;
        partitions.forEach(partition -> partitionMaxBytes.put(partition, uniformPartitionBytes));
        if (remainingBytes != 0) {
            int randomPartitionIndex = RANDOM.nextInt(partitionMaxBytes.keySet().size());
            allotBytesByPartitionIndex(uniformPartitionBytes + remainingBytes, partitions,
                partitionMaxBytes, Set.of(randomPartitionIndex));