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

Code before:
            // 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));
            List<TopicIdPartition> partitionsList = new ArrayList<>(partitions);
            int index = RANDOM.nextInt(partitionsList.size());
            int count = 0;
            while (count < requestMaxBytes) {
                partitionMaxBytes.put(partitionsList.get(index), 1);
                count += 1;
                index = (index + 1) % partitionsList.size();
            }
            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) {
            List<TopicIdPartition> partitionsList = new ArrayList<>(partitionMaxBytes.keySet());
            TopicIdPartition randomPartition = partitionsList.get(RANDOM.nextInt(partitionsList.size()));
            partitionMaxBytes.put(randomPartition, uniformPartitionBytes + remainingBytes);
        }
        return partitionMaxBytes;
    }

    // Visible for testing.
    static void checkValidArguments(int requestMaxBytes, Set<TopicIdPartition> partitions, int acquiredPartitionsSize) {
        if (partitions == null || partitions.isEmpty()) {
            throw new IllegalArgumentException("Partitions to generate max bytes is null or empty");
        }
        if (requestMaxBytes <= 0) {
            throw new IllegalArgumentException("Request max bytes must be greater than 0");
        }
        if (acquiredPartitionsSize <= 0) {
            throw new IllegalArgumentException("Acquired partitions size must be greater than 0");
        }
    }
}

Review comment (suggestion):
nit
```suggestion
        int partitionsSize
        ) {
```

Code after:
            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));
        }
        return partitionMaxBytes;
    }

    private static void allotBytesByPartitionIndex(
        int bytes,
        Set<TopicIdPartition> partitions,
        LinkedHashMap<TopicIdPartition, Integer> partitionMaxBytes,
        Set<Integer> partitionIndexesToAllot
    ) {
        int index = 0;
        int count = 0;
        for (TopicIdPartition partition : partitions) {
            if (partitionIndexesToAllot.contains(index)) {
                partitionMaxBytes.put(partition, bytes);
                count += 1;