zhijiangW commented on a change in pull request #11155: 
[FLINK-14818][benchmark] Fix receiving InputGate setup of 
StreamNetworkBenchmarkEnvironment.
URL: https://github.com/apache/flink/pull/11155#discussion_r383090260
 
 

 ##########
 File path: 
flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/benchmark/StreamNetworkBenchmarkEnvironment.java
 ##########
 @@ -272,11 +287,161 @@ private static ShuffleDescriptor 
createShuffleDescriptor(
                        ResultPartitionID resultPartitionID,
                        ResourceID location,
                        TaskManagerLocation senderLocation,
-                       int channel) {
+                       int connectionIndex) {
                final NettyShuffleDescriptorBuilder builder = 
NettyShuffleDescriptorBuilder.newBuilder()
                        .setId(resultPartitionID)
                        .setProducerInfoFromTaskManagerLocation(senderLocation)
-                       .setConnectionIndex(channel);
+                       .setConnectionIndex(connectionIndex);
                return localMode ? 
builder.setProducerLocation(location).buildLocal() : builder.buildRemote();
        }
+
+       /**
+        * A {@link SingleInputGateFactory} which replaces the default {@link 
RemoteInputChannel} and
+        * {@link LocalInputChannel} implementation with costume ones.
+        */
+       private static class TestSingleInputGateFactory extends 
SingleInputGateFactory {
+
+               private final ResourceID taskExecutorResourceId;
+               private final int partitionRequestInitialBackoff;
+               private final int partitionRequestMaxBackoff;
+               private final ConnectionManager connectionManager;
+               private final ResultPartitionManager partitionManager;
+               private final TaskEventPublisher taskEventPublisher;
+               private final NetworkBufferPool networkBufferPool;
+
+               public TestSingleInputGateFactory(
+                               @Nonnull ResourceID taskExecutorResourceId,
+                               @Nonnull NettyShuffleEnvironmentConfiguration 
networkConfig,
+                               @Nonnull ConnectionManager connectionManager,
+                               @Nonnull ResultPartitionManager 
partitionManager,
+                               @Nonnull TaskEventPublisher taskEventPublisher,
+                               @Nonnull NetworkBufferPool networkBufferPool) {
+                       super(
+                               taskExecutorResourceId,
+                               networkConfig,
+                               connectionManager,
+                               partitionManager,
+                               taskEventPublisher,
+                               networkBufferPool);
+                       this.networkBufferPool = networkBufferPool;
+                       this.taskEventPublisher = taskEventPublisher;
+                       this.partitionManager = partitionManager;
+                       this.connectionManager = connectionManager;
+                       this.partitionRequestMaxBackoff = 
networkConfig.partitionRequestMaxBackoff();
+                       this.partitionRequestInitialBackoff = 
networkConfig.partitionRequestInitialBackoff();
+                       this.taskExecutorResourceId = taskExecutorResourceId;
+               }
+
+               @Override
+               protected InputChannel createKnownInputChannel(
+                       SingleInputGate inputGate,
+                       int index,
+                       NettyShuffleDescriptor inputChannelDescriptor,
+                       ChannelStatistics channelStatistics,
+                       InputChannelMetrics metrics) {
+                       ResultPartitionID partitionId = 
inputChannelDescriptor.getResultPartitionID();
+                       if 
(inputChannelDescriptor.isLocalTo(taskExecutorResourceId)) {
+                               return new TestLocalInputChannel(
+                                       inputGate,
+                                       index,
+                                       partitionId,
+                                       partitionManager,
+                                       taskEventPublisher,
+                                       partitionRequestInitialBackoff,
+                                       partitionRequestMaxBackoff,
+                                       metrics);
+                       } else {
+                               return new TestRemoteInputChannel(
+                                       inputGate,
+                                       index,
+                                       partitionId,
+                                       
inputChannelDescriptor.getConnectionId(),
+                                       connectionManager,
+                                       partitionRequestInitialBackoff,
+                                       partitionRequestMaxBackoff,
+                                       metrics,
+                                       networkBufferPool);
+                       }
+               }
+       }
+
+       /**
+        * A {@link LocalInputChannel} which ignores the given subpartition 
index and uses channel index
+        * instead when requesting subpartition.
+        */
+       private static class TestLocalInputChannel extends LocalInputChannel {
+
+               private ResultPartitionID newPartitionID = new 
ResultPartitionID();
+
+               public TestLocalInputChannel(
+                               SingleInputGate inputGate,
+                               int channelIndex,
+                               ResultPartitionID partitionId,
+                               ResultPartitionManager partitionManager,
+                               TaskEventPublisher taskEventPublisher,
+                               int initialBackoff,
+                               int maxBackoff,
+                               InputChannelMetrics metrics) {
+                       super(
+                               inputGate,
+                               channelIndex,
+                               partitionId,
+                               partitionManager,
+                               taskEventPublisher,
+                               initialBackoff,
+                               maxBackoff,
+                               metrics);
+               }
+
+               @Override
+               public void requestSubpartition(int subpartitionIndex) throws 
IOException, InterruptedException {
+                       super.requestSubpartition(channelIndex);
+               }
+
+               @Override
+               public ResultPartitionID getPartitionId() {
 
 Review comment:
   I have some basic thoughts for the concerns of `ResultPartitionID` AND 
`subpartitionIndex`.
   
   The baseline is that we can not break the existing rules that 
`ResultPartitionID` is different for every input channel but 
`subpartitionIndex` is same for all. 
   
   - `ResultPartitionID`: We can pass a desired value directly while 
constructing the `InputChannel` instance instead of getting it from 
`NettyShuffleDescriptor` now. Although it is still not consistent between 
`NettyShuffleDescriptor` and `InputChannel`, it becomes consistent in all the 
places inside channel, and we also avoid hacking the override method of 
`getPartitionId` in `TestRemoteInputChannel` below.
   
   - `subpartitionIndex`: We can keep it same as now. Although we can also 
maintain this info in the `InputChannel` layer as did above, then we can also 
pass the desired subpartition index while constructing the `InputChannel` to 
further void hack `requestSubpartition` method. But considering it is same for 
all the channels in normal case, so it might be more reasonable to maintain it 
in gate layer as now. 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


With regards,
Apache Git Services

Reply via email to