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

Code before:
            return super.toString().toLowerCase(Locale.ROOT);
        }
    }

    /**
     * Returns the partition max bytes for a given partition based on the strategy type.
     * The partitions passed for maxBytes calculation are a subset of total acquired partitions for the share fetch request.
     * Thus, partitions for which we want to compute the max bytes <= acquired partitions.
     *
     * @param requestMaxBytes - The total max bytes available for the share fetch request
     * @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));

Review comment (suggestion):
> If there is extra bytes left post dividing it uniformly, assign it randomly to any one of the partitions

Should we consider splitting the extra bytes like Case 3 (one byte to as many partitions as possible) instead of assigning to randomly to one partition.
For example: if `requestMaxBytes` = 10, `acquiredPartitionsSize` = 6. 

- Then in Case 1, we would currently give 1 byte each to 6 partitions and then randomly give 4 bytes to a select partition making the final split as {5,1,1,1,1,1}.

- If we can split the remaining 4 bytes equally as well(like Case 3), final split would be {2, 2, 2, 2, 1, 1}. We would prefer this right?

(Same for Case 2 as well, when we have extra bytes)

Code after:

    /**
     * Returns the partition max bytes for a given partition based on the strategy type.
     * The partitions passed for maxBytes calculation can be a subset of total acquired partitions for the share fetch request.
     *
     * @param requestMaxBytes - The total max bytes available for the share fetch request
     * @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);
        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,