public static int computeNetworkBuffersForAnnouncing(
            final int numBuffersPerChannel,
            final int numFloatingBuffersPerGate,
            final Optional<Integer> maxRequiredBuffersPerGate,
            final int sortShuffleMinParallelism,
            final int sortShuffleMinBuffers,
            final Map<IntermediateDataSetID, Integer> inputChannelNums,
            final Map<IntermediateDataSetID, Integer> partitionReuseCount,
            final Map<IntermediateDataSetID, Integer> subpartitionNums,
            final Map<IntermediateDataSetID, ResultPartitionType> inputPartitionTypes,
            final Map<IntermediateDataSetID, ResultPartitionType> partitionTypes) {

        int requirementForInputs = 0;
        for (IntermediateDataSetID dataSetId : inputChannelNums.keySet()) {
            int numChannels = inputChannelNums.get(dataSetId);
            ResultPartitionType inputPartitionType = inputPartitionTypes.get(dataSetId);
            checkNotNull(inputPartitionType);

            int numSingleGateBuffers =
                    getNumBuffersToAnnounceForInputGate(
                            inputPartitionType,
                            numBuffersPerChannel,
                            numFloatingBuffersPerGate,
                            maxRequiredBuffersPerGate,
                            numChannels);
            checkState(partitionReuseCount.containsKey(dataSetId));
            requirementForInputs += numSingleGateBuffers * partitionReuseCount.get(dataSetId);
        }

        int requirementForOutputs = 0;
        for (IntermediateDataSetID dataSetId : subpartitionNums.keySet()) {
            int numSubs = subpartitionNums.get(dataSetId);
            ResultPartitionType partitionType = partitionTypes.get(dataSetId);
            checkNotNull(partitionType);

            requirementForOutputs +=
                    getNumBuffersToAnnounceForResultPartition(
                            partitionType,
                            numBuffersPerChannel,
                            numFloatingBuffersPerGate,
                            sortShuffleMinParallelism,
                            sortShuffleMinBuffers,
                            numSubs);
        }

        return requirementForInputs + requirementForOutputs;
    }
@Test
    void testComputeRequiredNetworkBuffers() throws Exception {
        int numBuffersPerChannel = 5;
        int numBuffersPerGate = 8;
        Optional<Integer> maxRequiredBuffersPerGate = Optional.of(Integer.MAX_VALUE);
        int sortShuffleMinParallelism = 8;
        int numSortShuffleMinBuffers = 12;

        IntermediateDataSetID ids1 = new IntermediateDataSetID();
        IntermediateDataSetID ids2 = new IntermediateDataSetID();

        int numChannels1 = 3;
        int numChannels2 = 4;

        IntermediateDataSetID ds1 = new IntermediateDataSetID();
        IntermediateDataSetID ds2 = new IntermediateDataSetID();
        IntermediateDataSetID ds3 = new IntermediateDataSetID();

        int numSubs1 = 5; // pipelined shuffle
        int numSubs2 = 6; // hash blocking shuffle
        int numSubs3 = 10; // sort blocking shuffle
        Map<IntermediateDataSetID, Integer> subpartitionNums =
                ImmutableMap.of(ds1, numSubs1, ds2, numSubs2, ds3, numSubs3);
        Map<IntermediateDataSetID, ResultPartitionType> partitionTypes =
                ImmutableMap.of(ds1, PIPELINED_BOUNDED, ds2, BLOCKING, ds3, BLOCKING);
        Map<IntermediateDataSetID, Integer> numInputChannels =
                ImmutableMap.of(ids1, numChannels1, ids2, numChannels2);
        Map<IntermediateDataSetID, Integer> partitionReuseCount = ImmutableMap.of(ids1, 1, ids2, 1);
        Map<IntermediateDataSetID, ResultPartitionType> inputPartitionTypes =
                ImmutableMap.of(ids1, PIPELINED_BOUNDED, ids2, BLOCKING);

        int numTotalBuffers =
                NettyShuffleUtils.computeNetworkBuffersForAnnouncing(
                        numBuffersPerChannel,
                        numBuffersPerGate,
                        maxRequiredBuffersPerGate,
                        sortShuffleMinParallelism,
                        numSortShuffleMinBuffers,
                        numInputChannels,
                        partitionReuseCount,
                        subpartitionNums,
                        inputPartitionTypes,
                        partitionTypes);

        NettyShuffleEnvironment sEnv =
                new NettyShuffleEnvironmentBuilder()
                        .setNumNetworkBuffers(numTotalBuffers)
                        .setNetworkBuffersPerChannel(numBuffersPerChannel)
                        .setSortShuffleMinBuffers(numSortShuffleMinBuffers)
                        .setSortShuffleMinParallelism(sortShuffleMinParallelism)
                        .build();

        SingleInputGate inputGate1 = createInputGate(sEnv, PIPELINED_BOUNDED, numChannels1);
        inputGate1.setup();

        SingleInputGate inputGate2 = createInputGate(sEnv, BLOCKING, numChannels2);
        inputGate2.setup();

        ResultPartition resultPartition1 = createResultPartition(sEnv, PIPELINED_BOUNDED, numSubs1);
        resultPartition1.setup();

        ResultPartition resultPartition2 = createResultPartition(sEnv, BLOCKING, numSubs2);
        resultPartition2.setup();

        ResultPartition resultPartition3 = createResultPartition(sEnv, BLOCKING, numSubs3);
        resultPartition3.setup();

        int expected =
                calculateBuffersConsumption(inputGate1)
                        + calculateBuffersConsumption(inputGate2)
                        + calculateBuffersConsumption(resultPartition1)
                        + calculateBuffersConsumption(resultPartition2)
                        + calculateBuffersConsumption(resultPartition3);
        assertThat(numTotalBuffers).isEqualTo(expected);

        inputGate1.close();
        inputGate2.close();
        resultPartition1.close();
        resultPartition2.close();
        resultPartition3.close();
    }